From c50d6aa09c9028dca3365516d30f1242cfd43306 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Hoguin?= <essen@ninenines.eu>
Date: Sat, 5 Oct 2019 13:04:21 +0200
Subject: Don't discard data following a Websocket upgrade request

While the protocol does not allow sending data before
receiving a successful Websocket upgrade response, we
do not want to discard that data if it does come in.
---
 test/ws_SUITE.erl | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

(limited to 'test/ws_SUITE.erl')

diff --git a/test/ws_SUITE.erl b/test/ws_SUITE.erl
index 4e35ec7..c99830c 100644
--- a/test/ws_SUITE.erl
+++ b/test/ws_SUITE.erl
@@ -304,6 +304,18 @@ do_ws_deflate_opts_z(Path, Config) ->
 	{error, closed} = gen_tcp:recv(Socket, 0, 6000),
 	ok.
 
+ws_first_frame_with_handshake(Config) ->
+	doc("Client sends the first frame immediately with the handshake. "
+		"This is invalid according to the protocol but we still want "
+		"to accept it if the handshake is successful."),
+	Mask = 16#37fa213d,
+	MaskedHello = do_mask(<<"Hello">>, Mask, <<>>),
+	{ok, Socket, _} = do_handshake("/ws_echo", "",
+		<<1:1, 0:3, 1:4, 1:1, 5:7, Mask:32, MaskedHello/binary>>,
+		Config),
+	{ok, <<1:1, 0:3, 1:4, 0:1, 5:7, "Hello">>} = gen_tcp:recv(Socket, 0, 6000),
+	ok.
+
 ws_init_return_ok(Config) ->
 	doc("Handler does nothing."),
 	{ok, Socket, _} = do_handshake("/ws_init?ok", Config),
@@ -636,9 +648,12 @@ ws_webkit_deflate_single_bytes(Config) ->
 %% Internal.
 
 do_handshake(Path, Config) ->
-	do_handshake(Path, "", Config).
+	do_handshake(Path, "", "", Config).
 
 do_handshake(Path, ExtraHeaders, Config) ->
+	do_handshake(Path, ExtraHeaders, "", Config).
+
+do_handshake(Path, ExtraHeaders, ExtraData, Config) ->
 	{ok, Socket} = gen_tcp:connect("localhost", config(port, Config),
 		[binary, {active, false}]),
 	ok = gen_tcp:send(Socket, [
@@ -650,10 +665,16 @@ do_handshake(Path, ExtraHeaders, Config) ->
 		"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
 		"Upgrade: websocket\r\n",
 		ExtraHeaders,
-		"\r\n"]),
+		"\r\n",
+		ExtraData]),
 	{ok, Handshake} = gen_tcp:recv(Socket, 0, 6000),
 	{ok, {http_response, {1, 1}, 101, _}, Rest} = erlang:decode_packet(http, Handshake, []),
-	[Headers, <<>>] = do_decode_headers(erlang:decode_packet(httph, Rest, []), []),
+	[Headers, Data] = do_decode_headers(erlang:decode_packet(httph, Rest, []), []),
+	%% Queue extra data back, if any. We don't want to receive it yet.
+	case Data of
+		<<>> -> ok;
+		_ -> gen_tcp:unrecv(Socket, Data)
+	end,
 	{_, "Upgrade"} = lists:keyfind('Connection', 1, Headers),
 	{_, "websocket"} = lists:keyfind('Upgrade', 1, Headers),
 	{_, "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="} = lists:keyfind("sec-websocket-accept", 1, Headers),
-- 
cgit v1.2.3