Add more frame types available in websocket replies

We can now reply empty close, ping and pong frames, or close
frames with a payload.

This means that we can send a frame and then close the connection
in a single operation.

If a close packet is sent, the connection is closed immediately,
even if there was frames that remained to be sent. Cowboy will
silently drop any extra frames in the list given as a reply.

4 files changed, 140 insertions, 27 deletions

diff --git a/src/cowboy_websocket.erl b/src/cowboy_websocket.erl
index 22bb8cd..0e41279 100644
--- a/src/cowboy_websocket.erl
+++ b/src/cowboy_websocket.erl
@@ -21,6 +21,10 @@
 %% Internal.
 -export([handler_loop/4]).
 
+-type frame() :: close | ping | pong
+	| {text | binary | close | ping | pong, binary()}.
+-export_type([frame/0]).
+
 -type opcode() :: 0 | 1 | 2 | 8 | 9 | 10.
 -type mask_key() :: 0..16#ffffffff.
 
@@ -455,6 +459,9 @@ handler_call(State=#state{handler=Handler, opts=Opts}, Req, HandlerState,
 			case websocket_send(Payload, State) of
 				ok ->
 					NextState(State, Req2, HandlerState2, RemainingData);
+				shutdown ->
+					handler_terminate(State, Req2, HandlerState,
+						{normal, shutdown});
 				{error, _} = Error ->
 					handler_terminate(State, Req2, HandlerState2, Error)
 			end;
@@ -464,6 +471,9 @@ handler_call(State=#state{handler=Handler, opts=Opts}, Req, HandlerState,
 				ok ->
 					NextState(State#state{hibernate=true},
 						Req2, HandlerState2, RemainingData);
+				shutdown ->
+					handler_terminate(State, Req2, HandlerState,
+						{normal, shutdown});
 				{error, _} = Error ->
 					handler_terminate(State, Req2, HandlerState2, Error)
 			end;
@@ -472,6 +482,9 @@ handler_call(State=#state{handler=Handler, opts=Opts}, Req, HandlerState,
 			case websocket_send_many(Payload, State) of
 				ok ->
 					NextState(State, Req2, HandlerState2, RemainingData);
+				shutdown ->
+					handler_terminate(State, Req2, HandlerState,
+						{normal, shutdown});
 				{error, _} = Error ->
 					handler_terminate(State, Req2, HandlerState2, Error)
 			end;
@@ -481,6 +494,9 @@ handler_call(State=#state{handler=Handler, opts=Opts}, Req, HandlerState,
 				ok ->
 					NextState(State#state{hibernate=true},
 						Req2, HandlerState2, RemainingData);
+				shutdown ->
+					handler_terminate(State, Req2, HandlerState,
+						{normal, shutdown});
 				{error, _} = Error ->
 					handler_terminate(State, Req2, HandlerState2, Error)
 			end;
@@ -498,8 +514,14 @@ handler_call(State=#state{handler=Handler, opts=Opts}, Req, HandlerState,
 		websocket_close(State, Req, HandlerState, {error, handler})
 	end.
 
--spec websocket_send({text | binary | ping | pong, binary()}, #state{})
-	-> ok | {error, atom()}.
+websocket_opcode(text) -> 1;
+websocket_opcode(binary) -> 2;
+websocket_opcode(close) -> 8;
+websocket_opcode(ping) -> 9;
+websocket_opcode(pong) -> 10.
+
+-spec websocket_send(frame(), #state{})
+	-> ok | shutdown | {error, atom()}.
 %% hixie-76 text frame.
 websocket_send({text, Payload}, #state{
 		socket=Socket, transport=Transport, version=0}) ->
@@ -507,24 +529,42 @@ websocket_send({text, Payload}, #state{
 %% Ignore all unknown frame types for compatibility with hixie 76.
 websocket_send(_Any, #state{version=0}) ->
 	ok;
+websocket_send(Type, #state{socket=Socket, transport=Transport})
+		when Type =:= close ->
+	Opcode = websocket_opcode(Type),
+	case Transport:send(Socket, << 1:1, 0:3, Opcode:4, 0:8 >>) of
+		ok -> shutdown;
+		Error -> Error
+	end;
+websocket_send(Type, #state{socket=Socket, transport=Transport})
+		when Type =:= ping; Type =:= pong ->
+	Opcode = websocket_opcode(Type),
+	Transport:send(Socket, << 1:1, 0:3, Opcode:4, 0:8 >>);
 websocket_send({Type, Payload}, #state{socket=Socket, transport=Transport}) ->
-	Opcode = case Type of
-		text -> 1;
-		binary -> 2;
-		ping -> 9;
-		pong -> 10
+	Opcode = websocket_opcode(Type),
+	Len = iolist_size(Payload),
+	%% Control packets must not be > 125 in length.
+	true = if Type =:= close; Type =:= ping; Type =:= pong ->
+			Len =< 125;
+		true ->
+			true
 	end,
-	Len = hybi_payload_length(iolist_size(Payload)),
-	Transport:send(Socket, [<< 1:1, 0:3, Opcode:4, 0:1, Len/bits >>,
-		Payload]).
+	BinLen = hybi_payload_length(Len),
+	Ret = Transport:send(Socket,
+		[<< 1:1, 0:3, Opcode:4, 0:1, BinLen/bits >>, Payload]),
+	case Type of
+		close -> shutdown;
+		_ -> Ret
+	end.
 
--spec websocket_send_many([{text | binary | ping | pong, binary()}], #state{})
-	-> ok | {error, atom()}.
+-spec websocket_send_many([frame()], #state{})
+	-> ok | shutdown | {error, atom()}.
 websocket_send_many([], _) ->
 	ok;
 websocket_send_many([Frame|Tail], State) ->
 	case websocket_send(Frame, State) of
 		ok -> websocket_send_many(Tail, State);
+		shutdown -> shutdown;
 		Error -> Error
 	end.
 
@@ -534,7 +574,6 @@ websocket_close(State=#state{socket=Socket, transport=Transport, version=0},
 		Req, HandlerState, Reason) ->
 	Transport:send(Socket, << 255, 0 >>),
 	handler_terminate(State, Req, HandlerState, Reason);
-%% @todo Send a Payload? Using Reason is usually good but we're quite careless.
 websocket_close(State=#state{socket=Socket, transport=Transport},
 		Req, HandlerState, Reason) ->
 	Transport:send(Socket, << 1:1, 0:3, 8:4, 0:8 >>),
diff --git a/src/cowboy_websocket_handler.erl b/src/cowboy_websocket_handler.erl
index 1b942fd..6d7f9de 100644
--- a/src/cowboy_websocket_handler.erl
+++ b/src/cowboy_websocket_handler.erl
@@ -50,7 +50,6 @@
 
 -type opts() :: any().
 -type state() :: any().
--type payload() :: {text | binary | ping | pong, binary()}.
 -type terminate_reason() :: {normal, closed}
 	| {normal, timeout}
 	| {error, closed}
@@ -67,15 +66,15 @@
 -callback websocket_handle({text | binary | ping | pong, binary()}, Req, State)
 	-> {ok, Req, State}
 	| {ok, Req, State, hibernate}
-	| {reply, payload() | [payload()], Req, State}
-	| {reply, payload() | [payload()], Req, State, hibernate}
+	| {reply, cowboy_websocket:frame() | [cowboy_websocket:frame()], Req, State}
+	| {reply, cowboy_websocket:frame() | [cowboy_websocket:frame()], Req, State, hibernate}
 	| {shutdown, Req, State}
 	when Req::cowboy_req:req(), State::state().
 -callback websocket_info(any(), Req, State)
 	-> {ok, Req, State}
 	| {ok, Req, State, hibernate}
-	| {reply, payload() | [payload()], Req, State}
-	| {reply, payload() | [payload()], Req, State, hibernate}
+	| {reply, cowboy_websocket:frame() | [cowboy_websocket:frame()], Req, State}
+	| {reply, cowboy_websocket:frame() | [cowboy_websocket:frame()], Req, State, hibernate}
 	| {shutdown, Req, State}
 	when Req::cowboy_req:req(), State::state().
 -callback websocket_terminate(terminate_reason(), cowboy_req:req(), state())
diff --git a/test/ws_SUITE.erl b/test/ws_SUITE.erl
index 7f1a18f..b63f41c 100644
--- a/test/ws_SUITE.erl
+++ b/test/ws_SUITE.erl
@@ -30,6 +30,8 @@
 -export([ws8_init_shutdown/1]).
 -export([ws8_single_bytes/1]).
 -export([ws13/1]).
+-export([ws_send_close/1]).
+-export([ws_send_close_payload/1]).
 -export([ws_send_many/1]).
 -export([ws_text_fragments/1]).
 -export([ws_timeout_hibernate/1]).
@@ -46,6 +48,8 @@ groups() ->
 		ws8_init_shutdown,
 		ws8_single_bytes,
 		ws13,
+		ws_send_close,
+		ws_send_close_payload,
 		ws_send_many,
 		ws_text_fragments,
 		ws_timeout_hibernate
@@ -85,7 +89,24 @@ init_dispatch() ->
 			{[<<"websocket">>], websocket_handler, []},
 			{[<<"ws_echo_handler">>], websocket_echo_handler, []},
 			{[<<"ws_init_shutdown">>], websocket_handler_init_shutdown, []},
-			{[<<"ws_send_many">>], ws_send_many_handler, []},
+			{[<<"ws_send_many">>], ws_send_many_handler, [
+				{sequence, [
+					{text, <<"one">>},
+					{text, <<"two">>},
+					{text, <<"seven!">>}]}
+			]},
+			{[<<"ws_send_close">>], ws_send_many_handler, [
+				{sequence, [
+					{text, <<"send">>},
+					close,
+					{text, <<"won't be received">>}]}
+			]},
+			{[<<"ws_send_close_payload">>], ws_send_many_handler, [
+				{sequence, [
+					{text, <<"send">>},
+					{close, <<"some text!">>},
+					{text, <<"won't be received">>}]}
+			]},
 			{[<<"ws_timeout_hibernate">>], ws_timeout_hibernate_handler, []}
 		]}
 	].
@@ -310,6 +331,64 @@ ws13(Config) ->
 	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
 	ok.
 
+ws_send_close(Config) ->
+	{port, Port} = lists:keyfind(port, 1, Config),
+	{ok, Socket} = gen_tcp:connect("localhost", Port,
+		[binary, {active, false}, {packet, raw}]),
+	ok = gen_tcp:send(Socket, [
+		"GET /ws_send_close HTTP/1.1\r\n"
+		"Host: localhost\r\n"
+		"Connection: Upgrade\r\n"
+		"Upgrade: websocket\r\n"
+		"Sec-WebSocket-Origin: http://localhost\r\n"
+		"Sec-WebSocket-Version: 8\r\n"
+		"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
+		"\r\n"]),
+	{ok, Handshake} = gen_tcp:recv(Socket, 0, 6000),
+	{ok, {http_response, {1, 1}, 101, "Switching Protocols"}, Rest}
+		= erlang:decode_packet(http, Handshake, []),
+	[Headers, <<>>] = websocket_headers(
+		erlang:decode_packet(httph, Rest, []), []),
+	{'Connection', "Upgrade"} = lists:keyfind('Connection', 1, Headers),
+	{'Upgrade', "websocket"} = lists:keyfind('Upgrade', 1, Headers),
+	{"sec-websocket-accept", "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="}
+		= lists:keyfind("sec-websocket-accept", 1, Headers),
+	%% We catch all frames at once and check them directly.
+	{ok, Many} = gen_tcp:recv(Socket, 8, 6000),
+	<< 1:1, 0:3, 1:4, 0:1, 4:7, "send",
+		1:1, 0:3, 8:4, 0:8 >> = Many,
+	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
+ws_send_close_payload(Config) ->
+	{port, Port} = lists:keyfind(port, 1, Config),
+	{ok, Socket} = gen_tcp:connect("localhost", Port,
+		[binary, {active, false}, {packet, raw}]),
+	ok = gen_tcp:send(Socket, [
+		"GET /ws_send_close_payload HTTP/1.1\r\n"
+		"Host: localhost\r\n"
+		"Connection: Upgrade\r\n"
+		"Upgrade: websocket\r\n"
+		"Sec-WebSocket-Origin: http://localhost\r\n"
+		"Sec-WebSocket-Version: 8\r\n"
+		"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
+		"\r\n"]),
+	{ok, Handshake} = gen_tcp:recv(Socket, 0, 6000),
+	{ok, {http_response, {1, 1}, 101, "Switching Protocols"}, Rest}
+		= erlang:decode_packet(http, Handshake, []),
+	[Headers, <<>>] = websocket_headers(
+		erlang:decode_packet(httph, Rest, []), []),
+	{'Connection', "Upgrade"} = lists:keyfind('Connection', 1, Headers),
+	{'Upgrade', "websocket"} = lists:keyfind('Upgrade', 1, Headers),
+	{"sec-websocket-accept", "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="}
+		= lists:keyfind("sec-websocket-accept", 1, Headers),
+	%% We catch all frames at once and check them directly.
+	{ok, Many} = gen_tcp:recv(Socket, 18, 6000),
+	<< 1:1, 0:3, 1:4, 0:1, 4:7, "send",
+		1:1, 0:3, 8:4, 0:1, 10:7, "some text!" >> = Many,
+	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
 ws_send_many(Config) ->
 	{port, Port} = lists:keyfind(port, 1, Config),
 	{ok, Socket} = gen_tcp:connect("localhost", Port,
diff --git a/test/ws_send_many_handler.erl b/test/ws_send_many_handler.erl
index ee386ba..bd67814 100644
--- a/test/ws_send_many_handler.erl
+++ b/test/ws_send_many_handler.erl
@@ -12,20 +12,16 @@
 init(_Any, _Req, _Opts) ->
 	{upgrade, protocol, cowboy_websocket}.
 
-websocket_init(_TransportName, Req, _Opts) ->
+websocket_init(_TransportName, Req, Sequence) ->
 	Req2 = cowboy_req:compact(Req),
 	erlang:send_after(10, self(), send_many),
-	{ok, Req2, undefined}.
+	{ok, Req2, Sequence}.
 
 websocket_handle(_Frame, Req, State) ->
 	{ok, Req, State}.
 
-websocket_info(send_many, Req, State) ->
-	{reply, [
-		{text, <<"one">>},
-		{text, <<"two">>},
-		{text, <<"seven!">>}
-	], Req, State}.
+websocket_info(send_many, Req, State = [{sequence, Sequence}]) ->
+	{reply, Sequence, Req, State}.
 
 websocket_terminate(_Reason, _Req, _State) ->
 	ok.