Supports going through multiple Socks proxies

This commit also reworks the switch_protocol command.
The `P | {P, Opts}` type is used here as well. This
allows us to remove the code specific to Websocket.

In addition a few new protocol functions allow us
to declare what's the name of the options key for
the protocol and what the capabilities are with
regard to keepalive.

5 files changed, 79 insertions, 52 deletions

diff --git a/src/gun.erl b/src/gun.erl
index 3802ee0..5a7750b 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -893,13 +893,12 @@ connecting(_, {retries, Retries, LookupInfo}, State=#state{opts=Opts,
 	case gun_tcp:connect(LookupInfo, ConnectTimeout) of
 		{ok, Socket} when Transport =:= gun_tcp ->
 			Protocol = case maps:get(protocols, Opts, [http]) of
-				[http] -> gun_http;
-				[http2] -> gun_http2;
-				[{socks, _}] -> gun_socks
+				[{P, _}] -> P;
+				[P] -> P
 			end,
 			EvHandlerState = EvHandler:connect_end(ConnectEvent#{
 				socket => Socket,
-				protocol => Protocol:name()
+				protocol => Protocol
 			}, EvHandlerState1),
 			{next_state, connected, State#state{event_handler_state=EvHandlerState},
 				{next_event, internal, {connected, Socket, Protocol}}};
@@ -951,10 +950,10 @@ tls_handshake(internal, {tls_handshake, HandshakeEvent, Protocols},
 			%% The transport is given in Proto:init/4 in the other case.
 			{keep_state, State} = commands([{switch_transport, gun_tls, TLSSocket}], State1),
 			{next_state, connected, State};
-		{ok, TLSSocket, NewProtocol, State1=#state{protocol_state=ProtoState}} ->
+		{ok, TLSSocket, NewProtocol, State1} ->
 			{keep_state, State} = commands([
 				{switch_transport, gun_tls, TLSSocket},
-				{switch_protocol, NewProtocol, ProtoState}
+				{switch_protocol, NewProtocol}
 			], State1),
 			{next_state, connected, State};
 		{error, Reason, State} ->
@@ -984,7 +983,7 @@ tls_handshake(info, {gun_tls_proxy, Socket, {ok, Negotiated}, {HandshakeEvent, P
 	NewProtocol = protocol_negotiated(Negotiated, Protocols),
 	EvHandlerState = EvHandler:tls_handshake_end(HandshakeEvent#{
 		socket => Socket,
-		protocol => NewProtocol:name()
+		protocol => NewProtocol
 	}, EvHandlerState0),
 	State1 = State0#state{event_handler_state=EvHandlerState},
 	{keep_state, State} = case NewProtocol of
@@ -994,7 +993,7 @@ tls_handshake(info, {gun_tls_proxy, Socket, {ok, Negotiated}, {HandshakeEvent, P
 			ProtoState = CurrentProtocol:switch_transport(Transport, Socket, ProtoState0),
 			{keep_state, State1#state{protocol_state=ProtoState}};
 		_ ->
-			commands([{switch_protocol, NewProtocol, ProtoState0}], State1)
+			commands([{switch_protocol, NewProtocol}], State1)
 	end,
 	{next_state, connected, State};
 tls_handshake(info, {gun_tls_proxy, Socket, Error = {error, Reason}, {HandshakeEvent, _}},
@@ -1019,7 +1018,7 @@ normal_tls_handshake(Socket, State=#state{event_handler=EvHandler, event_handler
 			Protocol = protocol_negotiated(ssl:negotiated_protocol(TLSSocket), Protocols),
 			EvHandlerState = EvHandler:tls_handshake_end(HandshakeEvent#{
 				socket => TLSSocket,
-				protocol => Protocol:name()
+				protocol => Protocol
 			}, EvHandlerState1),
 			{ok, TLSSocket, Protocol, State#state{event_handler_state=EvHandlerState}};
 		{error, Reason} ->
@@ -1029,32 +1028,33 @@ normal_tls_handshake(Socket, State=#state{event_handler=EvHandler, event_handler
 			{error, Reason, State#state{event_handler_state=EvHandlerState}}
 	end.
 
-protocol_negotiated({ok, <<"h2">>}, _) -> gun_http2;
-protocol_negotiated({ok, <<"http/1.1">>}, _) -> gun_http;
-protocol_negotiated({error, protocol_not_negotiated}, [{socks, _}]) -> gun_socks;
-protocol_negotiated({error, protocol_not_negotiated}, _) -> gun_http.
+protocol_negotiated({ok, <<"h2">>}, _) -> http2;
+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) ->
 	handle_common_connected(Type, Event, ?FUNCTION_NAME, State).
 
-connected(internal, {connected, Socket, Protocol=gun_socks},
+connected(internal, {connected, Socket, socks},
 		State=#state{owner=Owner, opts=Opts, transport=Transport}) ->
-	[{socks, ProtoOpts}] = [Proto || Proto = {socks, _} <- maps:get(protocols, Opts)],
+	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, Protocol},
-		State=#state{owner=Owner, opts=Opts, transport=Transport}) ->
-	ProtoOptsKey = case Protocol of
-		gun_http -> http_opts;
-		gun_http2 -> http2_opts
-	end,
-	ProtoOpts = maps:get(ProtoOptsKey, Opts, #{}),
+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),
 	Owner ! {gun_up, self(), Protocol:name()},
-	{keep_state, keepalive_timeout(active(State#state{socket=Socket,
-		protocol=Protocol, protocol_state=ProtoState}))};
+	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}
+	end;
 %% Public HTTP interface.
 connected(cast, {headers, ReplyTo, StreamRef, Method, Path, Headers, InitialFlow},
 		State=#state{origin_host=Host, origin_port=Port,
@@ -1299,32 +1299,33 @@ 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}));
-%% @todo The two loops should be reunified and this clause generalized.
-commands([{switch_protocol, Protocol=gun_ws, ProtoState}], State=#state{
-		event_handler=EvHandler, event_handler_state=EvHandlerState0}) ->
-	EvHandlerState = EvHandler:protocol_changed(#{protocol => Protocol:name()}, EvHandlerState0),
-	{keep_state, keepalive_cancel(State#state{protocol=Protocol, protocol_state=ProtoState,
-		event_handler_state=EvHandlerState})};
-%% @todo And this state should probably not be ignored.
-%% @todo Socks can be switching to *http* and we don't seem to support it properly yet.
-commands([{switch_protocol, Protocol, _ProtoState0}|Tail], State=#state{
+commands([{switch_protocol, Protocol0}|Tail], State0=#state{
 		owner=Owner, opts=Opts, socket=Socket, transport=Transport, protocol=CurrentProtocol,
 		event_handler=EvHandler, event_handler_state=EvHandlerState0}) ->
+	{Protocol, ProtoOpts} = case Protocol0 of
+		{P, PO} -> {protocol_handler(P), PO};
+		P ->
+			Protocol1 = protocol_handler(P),
+			{Protocol1, maps:get(Protocol1:opts_name(), Opts, #{})}
+	end,
 	%% When we switch_protocol from socks we must send a gun_socks_connected message.
 	_ = case CurrentProtocol of
 		gun_socks -> Owner ! {gun_socks_connected, self(), Protocol:name()};
 		_ -> ok
 	end,
-	ProtoOpts = maps:get(http2_opts, Opts, #{}),
 	ProtoState = Protocol:init(Owner, Socket, Transport, ProtoOpts),
 	EvHandlerState = EvHandler:protocol_changed(#{protocol => Protocol:name()}, EvHandlerState0),
-	commands(Tail, keepalive_timeout(State#state{protocol=Protocol, protocol_state=ProtoState,
-		event_handler_state=EvHandlerState}));
+	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))
+	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)}.
 
@@ -1370,6 +1371,11 @@ disconnect_flush(State=#state{socket=Socket, messages={OK, Closed, Error}}) ->
 		ok
 	end.
 
+protocol_handler(http) -> gun_http;
+protocol_handler(http2) -> gun_http2;
+protocol_handler(ws) -> gun_ws;
+protocol_handler(socks) -> gun_socks.
+
 active(State=#state{active=false}) ->
 	State;
 active(State=#state{socket=Socket, transport=Transport}) ->
diff --git a/src/gun_http.erl b/src/gun_http.erl
index 576c013..ba5c01d 100644
--- a/src/gun_http.erl
+++ b/src/gun_http.erl
@@ -16,6 +16,8 @@
 
 -export([check_options/1]).
 -export([name/0]).
+-export([opts_name/0]).
+-export([has_keepalive/0]).
 -export([init/4]).
 -export([switch_transport/3]).
 -export([handle/4]).
@@ -96,6 +98,8 @@ do_check_options([Opt|_]) ->
 	{error, {options, {http, Opt}}}.
 
 name() -> http.
+opts_name() -> http_opts.
+has_keepalive() -> true.
 
 init(Owner, Socket, Transport, Opts) ->
 	%% @todo If we keep the opts we don't need to add these to the state.
@@ -313,7 +317,7 @@ handle_head(Data, State=#http_state{version=ClientVersion, content_handlers=Hand
 								{origin, <<"http">>, NewHost, NewPort, connect}], EvHandlerState1};
 						[http2] ->
 							{[{origin, <<"http">>, NewHost, NewPort, connect},
-								{switch_protocol, gun_http2, State2}], EvHandlerState1}
+								{switch_protocol, http2}], EvHandlerState1}
 					end
 			end;
 		{_, _} when Status >= 100, Status =< 199 ->
@@ -885,4 +889,13 @@ ws_handshake_end(Buffer, #http_state{owner=Owner, socket=Socket, transport=Trans
 			{OK, _, _} = Transport:messages(),
 			self() ! {OK, Socket, Buffer}
 	end,
-	gun_ws:init(Owner, Socket, Transport, StreamRef, Headers, Extensions, InitialFlow, Handler, Opts).
+	%% Inform the user that the upgrade was successful and switch the protocol.
+	Owner ! {gun_upgrade, self(), StreamRef, [<<"websocket">>], Headers},
+	{switch_protocol, {ws, #{
+		stream_ref => StreamRef,
+		headers => Headers,
+		extensions => Extensions,
+		flow => InitialFlow,
+		handler => Handler,
+		opts => Opts
+	}}}.
diff --git a/src/gun_http2.erl b/src/gun_http2.erl
index 7dd369d..efa7afa 100644
--- a/src/gun_http2.erl
+++ b/src/gun_http2.erl
@@ -16,6 +16,8 @@
 
 -export([check_options/1]).
 -export([name/0]).
+-export([opts_name/0]).
+-export([has_keepalive/0]).
 -export([init/4]).
 -export([switch_transport/3]).
 -export([handle/4]).
@@ -97,6 +99,8 @@ do_check_options([Opt|_]) ->
 	{error, {options, {http2, Opt}}}.
 
 name() -> http2.
+opts_name() -> http2_opts.
+has_keepalive() -> true.
 
 init(Owner, Socket, Transport, Opts0) ->
 	%% We have different defaults than the protocol in order
diff --git a/src/gun_socks.erl b/src/gun_socks.erl
index 16684f8..6123953 100644
--- a/src/gun_socks.erl
+++ b/src/gun_socks.erl
@@ -16,6 +16,8 @@
 
 -export([check_options/1]).
 -export([name/0]).
+-export([opts_name/0]).
+-export([has_keepalive/0]).
 -export([init/4]).
 -export([switch_transport/3]).
 -export([handle/4]).
@@ -78,6 +80,8 @@ check_auth_opt(Methods) ->
 	end.
 
 name() -> socks.
+opts_name() -> socks_opts.
+has_keepalive() -> false.
 
 init(Owner, Socket, Transport, Opts) ->
 	5 = Version = maps:get(version, Opts, 5),
@@ -116,8 +120,7 @@ handle(<<1, 0>>, State=#socks_state{version=5, status=auth_username_password}) -
 handle(<<1, _>>, #socks_state{version=5, status=auth_username_password}) ->
 	{error, {socks5, username_password_auth_failure}};
 %% Connect reply.
-handle(<<5, 0, 0, Rest0/bits>>, State=#socks_state{owner=Owner, socket=Socket, transport=Transport, opts=Opts,
-		version=5, status=connect}) ->
+handle(<<5, 0, 0, Rest0/bits>>, #socks_state{opts=Opts, version=5, status=connect}) ->
 	%% @todo What to do with BoundAddr and BoundPort? Add as metadata to origin info?
 	{_BoundAddr, _BoundPort} = case Rest0 of
 		%% @todo Seen a server with <<1, 0:48>>.
@@ -139,16 +142,13 @@ handle(<<5, 0, 0, Rest0/bits>>, State=#socks_state{owner=Owner, socket=Socket, t
 			},
 			[{origin, <<"https">>, NewHost, NewPort, socks5},
 				{tls_handshake, HandshakeEvent, maps:get(protocols, Opts, [http])}];
-		#{protocols := [{socks, SockOpts}]} ->
+		#{protocols := [Protocol={socks, _}]} ->
 			[{origin, <<"http">>, NewHost, NewPort, socks5},
-				{switch_protocol, ?MODULE, init(Owner, Socket, Transport, SockOpts)}];
-		#{protocols := [http2]} ->
-			[{origin, <<"http">>, NewHost, NewPort, socks5},
-				{switch_protocol, gun_http2, State},
-				{mode, http}];
+				{switch_protocol, Protocol}];
 		_ ->
+			[Protocol] = maps:get(protocols, Opts, [http]),
 			[{origin, <<"http">>, NewHost, NewPort, socks5},
-				{switch_protocol, gun_http, State},
+				{switch_protocol, Protocol},
 				{mode, http}]
 	end;
 handle(<<5, Error, _/bits>>, #socks_state{version=5, status=connect}) ->
diff --git a/src/gun_ws.erl b/src/gun_ws.erl
index 49911dc..7d65be0 100644
--- a/src/gun_ws.erl
+++ b/src/gun_ws.erl
@@ -16,7 +16,9 @@
 
 -export([check_options/1]).
 -export([name/0]).
--export([init/9]).
+-export([opts_name/0]).
+-export([has_keepalive/0]).
+-export([init/4]).
 -export([handle/4]).
 -export([update_flow/4]).
 -export([closing/4]).
@@ -77,13 +79,15 @@ do_check_options([Opt|_]) ->
 	{error, {options, {ws, Opt}}}.
 
 name() -> ws.
+opts_name() -> ws_opts.
+has_keepalive() -> false.
 
-init(Owner, Socket, Transport, StreamRef, Headers, Extensions, InitialFlow, Handler, Opts) ->
-	Owner ! {gun_upgrade, self(), StreamRef, [<<"websocket">>], Headers},
+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),
-	{switch_protocol, ?MODULE, #ws_state{owner=Owner, stream_ref=StreamRef,
+	#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.