Do not ignore data received immediately after switching to raw

LH: Minor tweaks.

4 files changed, 17 insertions, 23 deletions

diff --git a/src/gun.erl b/src/gun.erl
index ac02230..a7c0239 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -1235,7 +1235,7 @@ tls_handshake(internal, {tls_handshake, HandshakeEvent, Protocols, ReplyTo},
 			reply(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)
@@ -1272,7 +1272,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#{
@@ -1809,8 +1809,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, Buffer}], 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
@@ -1833,9 +1833,13 @@ commands([{switch_protocol, NewProtocol, ReplyTo}], State0=#state{
 	case active(State1) of
 		{ok, State2} ->
 			State = keepalive_cancel(State2),
+			Actions = case Buffer of
+				<<>> -> [];
+				_ -> [{next_event, info, {OK, Socket, Buffer}}]
+			end,
 			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
diff --git a/src/gun_http.erl b/src/gun_http.erl
index 2dc2d67..d4304b0 100644
--- a/src/gun_http.erl
+++ b/src/gun_http.erl
@@ -358,7 +358,7 @@ handle_connect(Rest, State=#http_state{
 			gun:reply(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),
 				gun:reply(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),
 	gun:reply(ReplyTo, {gun_upgrade, self(), RealStreamRef, [<<"websocket">>], Headers}),
@@ -1055,4 +1045,4 @@ ws_handshake_end(Buffer,
 		flow => InitialFlow,
 		handler => Handler,
 		opts => Opts
-	}}, ReplyTo}.
+	}}, ReplyTo, Buffer}.
diff --git a/src/gun_socks.erl b/src/gun_socks.erl
index 90b7042..d29f906 100644
--- a/src/gun_socks.erl
+++ b/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),
 			gun:reply(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
diff --git a/src/gun_tunnel.erl b/src/gun_tunnel.erl
index 559e8f7..4d7a904 100644
--- a/src/gun_tunnel.erl
+++ b/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, <<>>}|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, <<>>}|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}},