From 6d4e15705fc1ac4a8cbb12cb9ea5ce62cf4df8b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Hoguin?= <essen@ninenines.eu>
Date: Sat, 8 Dec 2012 19:11:56 +0100
Subject: Add {close, StatusCode, Payload} and fix {close, Payload}

---
 src/cowboy_websocket.erl | 25 +++++++++++++++++--------
 test/ws_SUITE.erl        |  6 +++---
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/cowboy_websocket.erl b/src/cowboy_websocket.erl
index c53fe89..02431b8 100644
--- a/src/cowboy_websocket.erl
+++ b/src/cowboy_websocket.erl
@@ -22,7 +22,8 @@
 -export([handler_loop/4]).
 
 -type frame() :: close | ping | pong
-	| {text | binary | close | ping | pong, binary()}.
+	| {text | binary | close | ping | pong, binary()}
+	| {close, 1000..4999, binary()}.
 -export_type([frame/0]).
 
 -type opcode() :: 0 | 1 | 2 | 8 | 9 | 10.
@@ -541,22 +542,30 @@ 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({close, Payload}, State) ->
+	websocket_send({close, 1000, Payload}, State);
+websocket_send({Type = close, StatusCode, Payload}, #state{
+		socket=Socket, transport=Transport}) ->
+	Opcode = websocket_opcode(Type),
+	Len = 2 + iolist_size(Payload),
+	%% Control packets must not be > 125 in length.
+	true = Len =< 125,
+	BinLen = hybi_payload_length(Len),
+	Transport:send(Socket,
+		[<< 1:1, 0:3, Opcode:4, 0:1, BinLen/bits, StatusCode:16 >>, Payload]),
+	shutdown;
 websocket_send({Type, Payload}, #state{socket=Socket, transport=Transport}) ->
 	Opcode = websocket_opcode(Type),
 	Len = iolist_size(Payload),
 	%% Control packets must not be > 125 in length.
-	true = if Type =:= close; Type =:= ping; Type =:= pong ->
+	true = if Type =:= ping; Type =:= pong ->
 			Len =< 125;
 		true ->
 			true
 	end,
 	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.
+	Transport:send(Socket,
+		[<< 1:1, 0:3, Opcode:4, 0:1, BinLen/bits >>, Payload]).
 
 -spec websocket_send_many([frame()], #state{})
 	-> ok | shutdown | {error, atom()}.
diff --git a/test/ws_SUITE.erl b/test/ws_SUITE.erl
index 5047d97..1a0583d 100644
--- a/test/ws_SUITE.erl
+++ b/test/ws_SUITE.erl
@@ -106,7 +106,7 @@ init_dispatch() ->
 			{[<<"ws_send_close_payload">>], ws_send_many_handler, [
 				{sequence, [
 					{text, <<"send">>},
-					{close, <<"some text!">>},
+					{close, 1001, <<"some text!">>},
 					{text, <<"won't be received">>}]}
 			]},
 			{[<<"ws_timeout_hibernate">>], ws_timeout_hibernate_handler, []},
@@ -387,9 +387,9 @@ ws_send_close_payload(Config) ->
 	{"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),
+	{ok, Many} = gen_tcp:recv(Socket, 20, 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,
+		1:1, 0:3, 8:4, 0:1, 12:7, 1001:16, "some text!" >> = Many,
 	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
 	ok.
 
-- 
cgit v1.2.3