do not ignore data received immediately after switching to raw

src/gun.erl

@@ -1222,7 +1222,7 @@ tls_handshake(internal, {tls_handshake, HandshakeEvent, Protocols, ReplyTo},
 			ReplyTo ! {gun_tunnel_up, self(), StreamRef, Protocol:name()},
 			commands([
 				{switch_transport, gun_tls, TLSSocket},
-				{switch_protocol, NewProtocol, ReplyTo}
+				{switch_protocol, NewProtocol, ReplyTo, <<>>}
 			], State);
 		{error, Reason, State} ->
 			commands({error, Reason}, State)
@@ -1259,7 +1259,7 @@ tls_handshake(info, {gun_tls_proxy, Socket, {ok, Negotiated}, {HandshakeEvent, P
 		socket => Socket,
 		protocol => Protocol:name()
 	}, EvHandlerState0),
-	commands([{switch_protocol, NewProtocol, ReplyTo}], State0#state{event_handler_state=EvHandlerState});
+	commands([{switch_protocol, NewProtocol, ReplyTo, <<>>}], State0#state{event_handler_state=EvHandlerState});
 tls_handshake(info, {gun_tls_proxy, Socket, Error = {error, Reason}, {HandshakeEvent, _, _}},
 		State=#state{socket=Socket, event_handler=EvHandler, event_handler_state=EvHandlerState0}) ->
 	EvHandlerState = EvHandler:tls_handshake_end(HandshakeEvent#{
@@ -1796,8 +1796,8 @@ commands([{switch_transport, Transport, Socket}|Tail], State0=#state{
 		Disconnect ->
 			Disconnect
 	end;
-commands([{switch_protocol, NewProtocol, ReplyTo}], State0=#state{
-		opts=Opts, socket=Socket, transport=Transport,
+commands([{switch_protocol, NewProtocol, ReplyTo, Rest}], State0=#state{
+		opts=Opts, socket=Socket, transport=Transport, messages={OK, _, _},
 		event_handler=EvHandler, event_handler_state=EvHandlerState0}) ->
 	{Protocol, ProtoOpts0} = gun_protocols:handler_and_opts(NewProtocol, Opts),
 	ProtoOpts = case ProtoOpts0 of
@@ -1820,9 +1820,10 @@ commands([{switch_protocol, NewProtocol, ReplyTo}], State0=#state{
 	case active(State1) of
 		{ok, State2} ->
 			State = keepalive_cancel(State2),
+			Actions = [{next_event, info, {OK, Socket, Rest}} || Rest /= <<>>],
 			case Protocol:has_keepalive() of
-				true -> {next_state, StateName, keepalive_timeout(State)};
-				false -> {next_state, StateName, State}
+				true -> {next_state, StateName, keepalive_timeout(State), Actions};
+				false -> {next_state, StateName, State, Actions}
 			end;
 		Disconnect ->
 			Disconnect

src/gun_http.erl

@@ -358,7 +358,7 @@ handle_connect(Rest, State=#http_state{
 			ReplyTo ! {gun_tunnel_up, self(), RealStreamRef, Protocol:name()},
 			{[
 				{origin, <<"http">>, NewHost, NewPort, connect},
-				{switch_protocol, NewProtocol, ReplyTo}
+				{switch_protocol, NewProtocol, ReplyTo, <<>>}
 			], CookieStore, EvHandlerState}
 	end.
 
@@ -379,12 +379,11 @@ handle_inform(Rest, State=#http_state{
 		%% @todo We should check that we asked for an upgrade before accepting it.
 		{'HTTP/1.1', 101, _} when is_reference(StreamRef) ->
 			try
-				%% @todo We shouldn't ignore Rest.
 				{_, Upgrade0} = lists:keyfind(<<"upgrade">>, 1, Headers),
 				Upgrade = cow_http_hd:parse_upgrade(Upgrade0),
 				ReplyTo ! {gun_upgrade, self(), stream_ref(State, StreamRef), Upgrade, Headers},
 				%% @todo We probably need to add_stream_ref?
-				{{switch_protocol, raw, ReplyTo}, CookieStore, EvHandlerState0}
+				{{switch_protocol, raw, ReplyTo, Rest}, CookieStore, EvHandlerState0}
 			catch _:_ ->
 				%% When the Upgrade header is missing or invalid we treat
 				%% the response as any other informational response.
@@ -1034,17 +1033,8 @@ ws_handshake_extensions_and_protocol(Buffer, State,
 	end.
 
 %% We know that the most recent stream is the Websocket one.
-ws_handshake_end(Buffer,
-		State=#http_state{socket=Socket, transport=Transport, streams=[#stream{flow=InitialFlow}|_]},
+ws_handshake_end(Buffer, State=#http_state{streams=[#stream{flow=InitialFlow}|_]},
 		#websocket{ref=StreamRef, reply_to=ReplyTo, opts=Opts}, Headers, Extensions, Handler) ->
-	%% Send ourselves the remaining buffer, if any.
-	_ = case Buffer of
-		<<>> ->
-			ok;
-		_ ->
-			{OK, _, _} = Transport:messages(),
-			self() ! {OK, Socket, Buffer}
-	end,
 	%% Inform the user that the upgrade was successful and switch the protocol.
 	RealStreamRef = stream_ref(State, StreamRef),
 	ReplyTo ! {gun_upgrade, self(), RealStreamRef, [<<"websocket">>], Headers},
@@ -1055,4 +1045,4 @@ ws_handshake_end(Buffer,
 		flow => InitialFlow,
 		handler => Handler,
 		opts => Opts
-	}}, ReplyTo}.
+	}}, ReplyTo, Buffer}.

src/gun_socks.erl

@@ -169,7 +169,7 @@ handle(<<5, 0, 0, Rest0/bits>>, #socks_state{ref=StreamRef, reply_to=ReplyTo, op
 			Protocol = gun_protocols:handler(NewProtocol),
 			ReplyTo ! {gun_tunnel_up, self(), StreamRef, Protocol:name()},
 			[{origin, <<"http">>, NewHost, NewPort, socks5},
-				{switch_protocol, NewProtocol, ReplyTo}]
+				{switch_protocol, NewProtocol, ReplyTo, <<>>}]
 	end;
 handle(<<5, Error, _/bits>>, #socks_state{version=5, status=connect}) ->
 	Reason = case Error of

src/gun_tunnel.erl

@@ -492,7 +492,7 @@ commands([Origin={origin, Scheme, Host, Port, Type}|Tail],
 		origin_port => Port
 	}, EvHandlerState0),
 	commands(Tail, State#tunnel_state{protocol_origin=Origin}, EvHandler, EvHandlerState);
-commands([{switch_protocol, NewProtocol, ReplyTo}|Tail],
+commands([{switch_protocol, NewProtocol, ReplyTo, <<>> = _Rest}|Tail],
 		State=#tunnel_state{socket=Socket, transport=Transport, opts=Opts,
 		protocol_origin=undefined},
 		EvHandler, EvHandlerState0) ->
@@ -510,7 +510,7 @@ commands([{switch_protocol, NewProtocol, ReplyTo}|Tail],
 		Error={error, _} ->
 			{Error, EvHandlerState0}
 	end;
-commands([{switch_protocol, NewProtocol, ReplyTo}|Tail],
+commands([{switch_protocol, NewProtocol, ReplyTo, <<>> = _Rest}|Tail],
 		State=#tunnel_state{transport=Transport, stream_ref=TunnelStreamRef,
 		info=#{origin_host := Host, origin_port := Port}, opts=Opts, protocol=CurrentProto,
 		protocol_origin={origin, _Scheme, OriginHost, OriginPort, Type}},

test/raw_SUITE.erl

@@ -242,13 +242,23 @@ connect_raw_reply_to(_) ->
 
 http11_upgrade_raw_tcp(_) ->
 	doc("Use the HTTP Upgrade mechanism to switch to the raw protocol over TCP."),
-	do_http11_upgrade_raw(tcp).
+	do_http11_upgrade_raw(tcp, "").
 
 http11_upgrade_raw_tls(_) ->
 	doc("Use the HTTP Upgrade mechanism to switch to the raw protocol over TLS."),
-	do_http11_upgrade_raw(tls).
+	do_http11_upgrade_raw(tls, "").
 
-do_http11_upgrade_raw(OriginTransport) ->
+http11_upgrade_raw_tcp_with_init_msg(_) ->
+	doc("Use the HTTP Upgrade mechanism to switch to the raw protocol over TCP,"
+		" with HTTP header immediately followed by the body."),
+	do_http11_upgrade_raw(tcp, "Initial message.").
+
+http11_upgrade_raw_tls_with_init_msg(_) ->
+	doc("Use the HTTP Upgrade mechanism to switch to the raw protocol over TLS,"
+		" with HTTP header immediately followed by the body."),
+	do_http11_upgrade_raw(tls, "Initial message.").
+
+do_http11_upgrade_raw(OriginTransport, InitMsg) ->
 	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw,
 		fun (Parent, ListenSocket, ClientSocket, ClientTransport) ->
 			%% We skip the request and send a 101 response unconditionally.
@@ -257,7 +267,7 @@ do_http11_upgrade_raw(OriginTransport) ->
 				"HTTP/1.1 101 Switching Protocols\r\n"
 				"Connection: upgrade\r\n"
 				"Upgrade: custom/1.0\r\n"
-				"\r\n"),
+				"\r\n" ++ InitMsg),
 			do_echo(Parent, ListenSocket, ClientSocket, ClientTransport)
 		end),
 	{ok, ConnPid} = gun:open("localhost", OriginPort, #{
@@ -271,6 +281,14 @@ do_http11_upgrade_raw(OriginTransport) ->
 		<<"upgrade">> => <<"custom/1.0">>
 	}),
 	{upgrade, [<<"custom/1.0">>], _} = gun:await(ConnPid, StreamRef),
+	case InitMsg of
+		[] ->
+			ok;
+		_  ->
+			InitMsgBin = list_to_binary(InitMsg),
+			{data, nofin, InitMsgBin} = gun:await(ConnPid, undefined),
+			ok
+	end,
 	%% When we take over the entire connection there is no stream reference.
 	gun:data(ConnPid, undefined, nofin, <<"Hello world!">>),
 	{data, nofin, <<"Hello world!">>} = gun:await(ConnPid, undefined),