Rework state transitions resulting from from protocol changes

5 files changed, 31 insertions, 45 deletions

diff --git a/src/gun.erl b/src/gun.erl
index 5bf37fb..ca6c730 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -98,7 +98,7 @@
 -export([connecting/3]).
 -export([initial_tls_handshake/3]).
 -export([tls_handshake/3]).
--export([not_fully_connected/3]).
+-export([connected_no_input/3]).
 -export([connected/3]).
 -export([closing/3]).
 -export([terminate/3]).
@@ -951,11 +951,10 @@ tls_handshake(internal, {tls_handshake, HandshakeEvent, Protocols},
 			{keep_state, State} = commands([{switch_transport, gun_tls, TLSSocket}], State1),
 			{next_state, connected, State};
 		{ok, TLSSocket, NewProtocol, State1} ->
-			{keep_state, State} = commands([
+			commands([
 				{switch_transport, gun_tls, TLSSocket},
 				{switch_protocol, NewProtocol}
-			], State1),
-			{next_state, connected, State};
+			], State1);
 		{error, Reason, State} ->
 			commands({error, Reason}, State)
 	end;
@@ -985,17 +984,16 @@ tls_handshake(info, {gun_tls_proxy, Socket, {ok, Negotiated}, {HandshakeEvent, P
 		socket => Socket,
 		protocol => NewProtocol
 	}, EvHandlerState0),
-	State1 = State0#state{event_handler_state=EvHandlerState},
-	{keep_state, State} = case NewProtocol of
+	State = State0#state{event_handler_state=EvHandlerState},
+	case NewProtocol of
 		CurrentProtocol ->
 			%% We only need to switch the transport when the protocol remains the same.
 			%% The transport is given in Proto:init/4 in the other case.
 			ProtoState = CurrentProtocol:switch_transport(Transport, Socket, ProtoState0),
-			{keep_state, State1#state{protocol_state=ProtoState}};
+			{keep_state, State#state{protocol_state=ProtoState}};
 		_ ->
-			commands([{switch_protocol, NewProtocol}], State1)
-	end,
-	{next_state, connected, State};
+			commands([{switch_protocol, NewProtocol}], State)
+	end;
 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#{
@@ -1033,27 +1031,23 @@ protocol_negotiated({ok, <<"http/1.1">>}, _) -> http;
 protocol_negotiated({error, protocol_not_negotiated}, [{socks, _}]) -> socks;
 protocol_negotiated({error, protocol_not_negotiated}, _) -> http.
 
-not_fully_connected(Type, Event, State) ->
+connected_no_input(Type, Event, State) ->
 	handle_common_connected(Type, Event, ?FUNCTION_NAME, State).
 
-connected(internal, {connected, Socket, socks},
-		State=#state{owner=Owner, opts=Opts, transport=Transport}) ->
-	Protocol = gun_socks,
-	[{socks, ProtoOpts}] = maps:get(protocols, Opts),
-	ProtoState = Protocol:init(Owner, Socket, Transport, ProtoOpts),
-	Owner ! {gun_up, self(), Protocol:name()},
-	{next_state, not_fully_connected, active(State#state{socket=Socket,
-		protocol=Protocol, protocol_state=ProtoState})};
 connected(internal, {connected, Socket, Protocol0},
 		State0=#state{owner=Owner, opts=Opts, transport=Transport}) ->
 	Protocol = protocol_handler(Protocol0),
-	ProtoOpts = maps:get(Protocol:opts_name(), Opts, #{}),
-	ProtoState = Protocol:init(Owner, Socket, Transport, ProtoOpts),
+	%% Protocol options may have been given along the protocol name.
+	ProtoOpts = case lists:keyfind(Protocol0, 1, maps:get(protocols, Opts, [http])) of
+		{_, PO} -> PO;
+		false -> maps:get(Protocol:opts_name(), Opts, #{})
+	end,
+	{StateName, ProtoState} = Protocol:init(Owner, Socket, Transport, ProtoOpts),
 	Owner ! {gun_up, self(), Protocol:name()},
 	State = active(State0#state{socket=Socket, protocol=Protocol, protocol_state=ProtoState}),
 	case Protocol:has_keepalive() of
-		true -> {keep_state, keepalive_timeout(State)};
-		false -> {keep_state, State}
+		true -> {next_state, StateName, keepalive_timeout(State)};
+		false -> {next_state, StateName, State}
 	end;
 %% Public HTTP interface.
 connected(cast, {headers, ReplyTo, StreamRef, Method, Path, Headers, InitialFlow},
@@ -1299,7 +1293,7 @@ commands([{switch_transport, Transport, Socket}|Tail], State=#state{
 	commands(Tail, active(State#state{socket=Socket, transport=Transport,
 		messages=Transport:messages(), protocol_state=ProtoState,
 		event_handler_state=EvHandlerState}));
-commands([{switch_protocol, Protocol0}|Tail], State0=#state{
+commands([{switch_protocol, Protocol0}], State0=#state{
 		owner=Owner, opts=Opts, socket=Socket, transport=Transport, protocol=CurrentProtocol,
 		event_handler=EvHandler, event_handler_state=EvHandlerState0}) ->
 	{Protocol, ProtoOpts} = case Protocol0 of
@@ -1313,21 +1307,17 @@ commands([{switch_protocol, Protocol0}|Tail], State0=#state{
 		gun_socks -> Owner ! {gun_socks_connected, self(), Protocol:name()};
 		_ -> ok
 	end,
-	ProtoState = Protocol:init(Owner, Socket, Transport, ProtoOpts),
+	{StateName, ProtoState} = Protocol:init(Owner, Socket, Transport, ProtoOpts),
 	EvHandlerState = EvHandler:protocol_changed(#{protocol => Protocol:name()}, EvHandlerState0),
 	State = State0#state{protocol=Protocol, protocol_state=ProtoState, event_handler_state=EvHandlerState},
 	case Protocol:has_keepalive() of
-		true -> commands(Tail, keepalive_timeout(State));
-		false -> commands(Tail, keepalive_cancel(State))
+		true -> {next_state, StateName, keepalive_timeout(active(State))};
+		false -> {next_state, StateName, keepalive_cancel(active(State))}
 	end;
 %% Perform a TLS handshake.
 commands([TLSHandshake={tls_handshake, _, _}], State) ->
 	{next_state, tls_handshake, State,
-		{next_event, internal, TLSHandshake}};
-%% Switch from not_fully_connected to connected.
-%% @todo Do this in switch_protocol.
-commands([{mode, http}], State) ->
-	{next_state, connected, active(State)}.
+		{next_event, internal, TLSHandshake}}.
 
 disconnect(State0=#state{owner=Owner, status=Status, opts=Opts,
 		socket=Socket, transport=Transport,
@@ -1389,7 +1379,7 @@ status(State, NewStatus) ->
 	State#state{status=NewStatus}.
 
 keepalive_timeout(State=#state{opts=Opts, protocol=Protocol}) ->
-	ProtoOpts = maps:get(Protocols:opts_name(), Opts, #{}),
+	ProtoOpts = maps:get(Protocol:opts_name(), Opts, #{}),
 	Keepalive = maps:get(keepalive, ProtoOpts, Protocol:default_keepalive()),
 	KeepaliveRef = case Keepalive of
 		infinity -> undefined;
diff --git a/src/gun_http.erl b/src/gun_http.erl
index 2c26d01..85fd2c4 100644
--- a/src/gun_http.erl
+++ b/src/gun_http.erl
@@ -108,8 +108,8 @@ init(Owner, Socket, Transport, Opts) ->
 	Version = maps:get(version, Opts, 'HTTP/1.1'),
 	Handlers = maps:get(content_handlers, Opts, [gun_data_h]),
 	TransformHeaderName = maps:get(transform_header_name, Opts, fun (N) -> N end),
-	#http_state{owner=Owner, socket=Socket, transport=Transport, opts=Opts, version=Version,
-		content_handlers=Handlers, transform_header_name=TransformHeaderName}.
+	{connected, #http_state{owner=Owner, socket=Socket, transport=Transport, opts=Opts,
+		version=Version, content_handlers=Handlers, transform_header_name=TransformHeaderName}}.
 
 switch_transport(Transport, Socket, State) ->
 	State#http_state{socket=Socket, transport=Transport}.
diff --git a/src/gun_http2.erl b/src/gun_http2.erl
index 66e7732..9ddf9bb 100644
--- a/src/gun_http2.erl
+++ b/src/gun_http2.erl
@@ -118,7 +118,7 @@ init(Owner, Socket, Transport, Opts0) ->
 		transport=Transport, opts=Opts, content_handlers=Handlers,
 		http2_machine=HTTP2Machine},
 	Transport:send(Socket, Preface),
-	State.
+	{connected, State}.
 
 switch_transport(Transport, Socket, State) ->
 	State#http2_state{socket=Socket, transport=Transport}.
diff --git a/src/gun_socks.erl b/src/gun_socks.erl
index 6123953..0427cb1 100644
--- a/src/gun_socks.erl
+++ b/src/gun_socks.erl
@@ -91,8 +91,8 @@ init(Owner, Socket, Transport, Opts) ->
 		none -> <<0>>
 	end || A <- Auth>>,
 	Transport:send(Socket, [<<5, (length(Auth))>>, Methods]),
-	#socks_state{owner=Owner, socket=Socket, transport=Transport, opts=Opts,
-		version=Version, status=auth_method_select}.
+	{connected_no_input, #socks_state{owner=Owner, socket=Socket, transport=Transport,
+		opts=Opts, version=Version, status=auth_method_select}}.
 
 switch_transport(Transport, Socket, State) ->
 	State#socks_state{socket=Socket, transport=Transport}.
@@ -142,14 +142,10 @@ handle(<<5, 0, 0, Rest0/bits>>, #socks_state{opts=Opts, version=5, status=connec
 			},
 			[{origin, <<"https">>, NewHost, NewPort, socks5},
 				{tls_handshake, HandshakeEvent, maps:get(protocols, Opts, [http])}];
-		#{protocols := [Protocol={socks, _}]} ->
-			[{origin, <<"http">>, NewHost, NewPort, socks5},
-				{switch_protocol, Protocol}];
 		_ ->
 			[Protocol] = maps:get(protocols, Opts, [http]),
 			[{origin, <<"http">>, NewHost, NewPort, socks5},
-				{switch_protocol, Protocol},
-				{mode, http}]
+				{switch_protocol, Protocol}]
 	end;
 handle(<<5, Error, _/bits>>, #socks_state{version=5, status=connect}) ->
 	Reason = case Error of
diff --git a/src/gun_ws.erl b/src/gun_ws.erl
index 7d65be0..60926f2 100644
--- a/src/gun_ws.erl
+++ b/src/gun_ws.erl
@@ -85,9 +85,9 @@ has_keepalive() -> false.
 init(Owner, Socket, Transport, #{stream_ref := StreamRef, headers := Headers,
 		extensions := Extensions, flow := InitialFlow, handler := Handler, opts := Opts}) ->
 	{ok, HandlerState} = Handler:init(Owner, StreamRef, Headers, Opts),
-	#ws_state{owner=Owner, stream_ref=StreamRef,
+	{connected, #ws_state{owner=Owner, stream_ref=StreamRef,
 		socket=Socket, transport=Transport, opts=Opts, extensions=Extensions,
-		flow=InitialFlow, handler=Handler, handler_state=HandlerState}.
+		flow=InitialFlow, handler=Handler, handler_state=HandlerState}}.
 
 %% Do not handle anything if we received a close frame.
 %% Initiate or terminate the closing state depending on whether we sent a close yet.