Make sure the hixie-76 websocket code works properly with proxies

2 files changed, 28 insertions, 12 deletions

diff --git a/src/cowboy_http_websocket.erl b/src/cowboy_http_websocket.erl
index d3c9c2d..59be2e5 100644
--- a/src/cowboy_http_websocket.erl
+++ b/src/cowboy_http_websocket.erl
@@ -50,7 +50,7 @@
 	version :: 0 | 7 | 8,
 	handler :: module(),
 	opts :: any(),
-	challenge = undefined :: undefined | binary(),
+	challenge = undefined :: undefined | binary() | {binary(), binary()},
 	timeout = infinity :: timeout(),
 	timeout_ref = undefined :: undefined | reference(),
 	messages = undefined :: undefined | {atom(), atom(), atom()},
@@ -86,6 +86,11 @@ websocket_upgrade(State, Req) ->
 -spec websocket_upgrade(undefined | <<_:8>>, #state{}, #http_req{})
 	-> {ok, #state{}, #http_req{}}.
 %% No version given. Assuming hixie-76 draft.
+%%
+%% We need to wait to send a reply back before trying to read the
+%% third part of the challenge key, because proxies will wait for
+%% a reply before sending it. Therefore we calculate the challenge
+%% key only in websocket_handshake/3.
 %% @todo Check Origin?
 websocket_upgrade(undefined, State, Req) ->
 	{<<"WebSocket">>, Req2} = cowboy_http_req:header('Upgrade', Req),
@@ -93,11 +98,9 @@ websocket_upgrade(undefined, State, Req) ->
 	{Key1, Req4} = cowboy_http_req:header(<<"Sec-Websocket-Key1">>, Req3),
 	{Key2, Req5} = cowboy_http_req:header(<<"Sec-Websocket-Key2">>, Req4),
 	false = lists:member(undefined, [Origin, Key1, Key2]),
-	{ok, Key3, Req6} = cowboy_http_req:body(8, Req5),
-	Challenge = hixie76_challenge(Key1, Key2, Key3),
 	EOP = binary:compile_pattern(<< 255 >>),
-	{ok, State#state{version=0, origin=Origin, challenge=Challenge,
-		eop=EOP}, Req6};
+	{ok, State#state{version=0, origin=Origin, challenge={Key1, Key2},
+		eop=EOP}, Req5};
 %% Versions 7 and 8. Implementation follows the hybi 7 through 10 drafts.
 %% @todo We don't need Origin?
 websocket_upgrade(<< Version >>, State, Req)
@@ -162,8 +165,9 @@ upgrade_terminate(#http_req{socket=Socket, transport=Transport}) ->
 
 -spec websocket_handshake(#state{}, #http_req{}, any()) -> ok | none().
 websocket_handshake(State=#state{version=0, origin=Origin,
-		challenge=Challenge}, Req=#http_req{transport=Transport,
-		raw_host=Host, port=Port, raw_path=Path, raw_qs=QS}, HandlerState) ->
+		challenge={Key1, Key2}}, Req=#http_req{socket=Socket,
+		transport=Transport, raw_host=Host, port=Port,
+		raw_path=Path, raw_qs=QS}, HandlerState) ->
 	Location = hixie76_location(Transport:name(), Host, Port, Path, QS),
 	{ok, Req2} = cowboy_http_req:reply(
 		<<"101 WebSocket Protocol Handshake">>,
@@ -171,9 +175,15 @@ websocket_handshake(State=#state{version=0, origin=Origin,
 		 {<<"Upgrade">>, <<"WebSocket">>},
 		 {<<"Sec-Websocket-Location">>, Location},
 		 {<<"Sec-Websocket-Origin">>, Origin}],
-		Challenge, Req#http_req{resp_state=waiting}),
+		[], Req#http_req{resp_state=waiting}),
+	%% We replied with a proper response. Proxies should be happy enough,
+	%% we can now read the 8 last bytes of the challenge keys and send
+	%% the challenge response directly to the socket.
+	{ok, Key3, Req3} = cowboy_http_req:body(8, Req2),
+	Challenge = hixie76_challenge(Key1, Key2, Key3),
+	Transport:send(Socket, Challenge),
 	handler_before_loop(State#state{messages=Transport:messages()},
-		Req2, HandlerState, <<>>);
+		Req3, HandlerState, <<>>);
 websocket_handshake(State=#state{challenge=Challenge},
 		Req=#http_req{transport=Transport}, HandlerState) ->
 	{ok, Req2} = cowboy_http_req:reply(
diff --git a/test/http_SUITE.erl b/test/http_SUITE.erl
index b5ee98b..de1045a 100644
--- a/test/http_SUITE.erl
+++ b/test/http_SUITE.erl
@@ -243,11 +243,15 @@ raw(Config) ->
 	[{Packet, StatusCode} = raw_req(Packet, Config)
 		|| {Packet, StatusCode} <- Tests].
 
+%% This test makes sure the code works even if we wait for a reply
+%% before sending the third challenge key in the GET body.
+%%
+%% This ensures that Cowboy will work fine with proxies on hixie.
 ws0(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, [
+	ok = gen_tcp:send(Socket,
 		"GET /websocket HTTP/1.1\r\n"
 		"Host: localhost\r\n"
 		"Connection: Upgrade\r\n"
@@ -255,11 +259,11 @@ ws0(Config) ->
 		"Origin: http://localhost\r\n"
 		"Sec-Websocket-Key1: Y\" 4 1Lj!957b8@0H756!i\r\n"
 		"Sec-Websocket-Key2: 1711 M;4\\74  80<6\r\n"
-		"\r\n", <<15,245,8,18,2,204,133,33>>]),
+		"\r\n"),
 	{ok, Handshake} = gen_tcp:recv(Socket, 0, 6000),
 	{ok, {http_response, {1, 1}, 101, "WebSocket Protocol Handshake"}, Rest}
 		= erlang:decode_packet(http, Handshake, []),
-	[Headers, Body] = websocket_headers(
+	[Headers, <<>>] = websocket_headers(
 		erlang:decode_packet(httph, Rest, []), []),
 	{'Connection', "Upgrade"} = lists:keyfind('Connection', 1, Headers),
 	{'Upgrade', "WebSocket"} = lists:keyfind('Upgrade', 1, Headers),
@@ -267,6 +271,8 @@ ws0(Config) ->
 		= lists:keyfind("sec-websocket-location", 1, Headers),
 	{"sec-websocket-origin", "http://localhost"}
 		= lists:keyfind("sec-websocket-origin", 1, Headers),
+	ok = gen_tcp:send(Socket, <<15,245,8,18,2,204,133,33>>),
+	{ok, Body} = gen_tcp:recv(Socket, 0, 6000),
 	<<169,244,191,103,146,33,149,59,74,104,67,5,99,118,171,236>> = Body,
 	ok = gen_tcp:send(Socket, << 0, "client_msg", 255 >>),
 	{ok, << 0, "client_msg", 255 >>} = gen_tcp:recv(Socket, 0, 6000),