Make HTTP/2 CONNECT to a SOCKS server work

4 files changed, 147 insertions, 27 deletions

diff --git a/src/gun.erl b/src/gun.erl
index 8b0b78d..987b806 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -115,19 +115,6 @@
 	| #{binary() | string() | atom() => iodata()}.
 -export_type([req_headers/0]).
 
--type tunnel_info() :: #{
-	stream_ref := reference() | [reference()],
-
-	%% Tunnel.
-	host := inet:hostname() | inet:ip_address(),
-	port := inet:port_number(),
-
-	%% Origin.
-	origin_host => inet:hostname() | inet:ip_address(),
-	origin_port => inet:port_number()
-}.
--export_type([tunnel_info/0]).
-
 -type ws_close_code() :: 1000..4999.
 
 -type ws_frame() :: close | ping | pong
@@ -183,6 +170,20 @@
 	protocol := http | socks
 }.
 
+-type tunnel_info() :: #{
+	%% Tunnel.
+	host := inet:hostname() | inet:ip_address(),
+	port := inet:port_number(),
+
+	%% Origin.
+	origin_host => inet:hostname() | inet:ip_address(),
+	origin_port => inet:port_number(),
+
+	%% Non-stream intermediaries (for example SOCKS).
+	intermediaries => [intermediary()]
+}.
+-export_type([tunnel_info/0]).
+
 -type raw_opts() :: #{}.
 -export_type([raw_opts/0]).
 
@@ -1229,7 +1230,7 @@ connected(cast, {connect, ReplyTo, StreamRef, Destination, Headers, InitialFlow}
 			protocol=Protocol, protocol_state=ProtoState}) ->
 	%% @todo No events are currently handled for the CONNECT request?
 	ProtoState2 = Protocol:connect(ProtoState, StreamRef, ReplyTo,
-		Destination, #{stream_ref => StreamRef, host => Host, port => Port},
+		Destination, #{host => Host, port => Port},
 		Headers, InitialFlow),
 	{keep_state, State#state{protocol_state=ProtoState2}};
 %% Public Websocket interface.
diff --git a/src/gun_http2.erl b/src/gun_http2.erl
index b1dc1e6..b8ae033 100644
--- a/src/gun_http2.erl
+++ b/src/gun_http2.erl
@@ -345,6 +345,42 @@ tunnel_commands([SetCookie={set_cookie, _, _, _, _}|Tail], Stream, Protocol, Tun
 		State=#http2_state{commands_queue=Queue}) ->
 	tunnel_commands(Tail, Stream, Protocol, TunnelInfo,
 		State#http2_state{commands_queue=[SetCookie|Queue]});
+tunnel_commands([{origin, _, NewHost, NewPort, Type}|Tail], Stream, Protocol, TunnelInfo, State) ->
+%% @todo Event?
+	tunnel_commands(Tail, Stream, Protocol, TunnelInfo#{
+		origin_host => NewHost,
+		origin_port => NewPort,
+		intermediaries => [#{
+			type => Type,
+			host => maps:get(origin_host, TunnelInfo),
+			port => maps:get(origin_port, TunnelInfo),
+			transport => tcp, %% @todo
+			protocol => Protocol:name()
+		}|maps:get(intermediaries, TunnelInfo, [])]
+	}, State);
+tunnel_commands([{switch_protocol, Protocol0, ReplyTo}|Tail], Stream=#stream{ref=StreamRef},
+		CurrentProtocol, TunnelInfo, State=#http2_state{opts=Opts}) ->
+	{Protocol, ProtoOpts} = case Protocol0 of
+		{P, PO} -> {gun:protocol_handler(P), PO};
+		P ->
+			Protocol1 = gun:protocol_handler(P),
+			%% @todo We need to allow other protocol opts in http2_opts too.
+			{Protocol1, maps:get(Protocol1:opts_name(), Opts, #{})}
+	end,
+	%% When we switch_protocol from socks we must send a gun_socks_up message.
+%% @todo OK but perhaps we should give the StreamRef!!
+	_ = case CurrentProtocol of
+		gun_socks -> ReplyTo ! {gun_socks_up, self(), Protocol:name()};
+		_ -> ok
+	end,
+	OriginSocket = #{
+		reply_to => ReplyTo,
+		stream_ref => StreamRef
+	},
+	OriginTransport = gun_tcp_proxy,
+	{_, ProtoState} = Protocol:init(ReplyTo, OriginSocket, OriginTransport, ProtoOpts),
+%% @todo	EvHandlerState = EvHandler:protocol_changed(#{protocol => Protocol:name()}, EvHandlerState0),
+	tunnel_commands([{state, ProtoState}|Tail], Stream, Protocol, TunnelInfo, State);
 tunnel_commands([{active, true}|Tail], Stream, Protocol, TunnelInfo, State) ->
 	tunnel_commands(Tail, Stream, Protocol, TunnelInfo, State).
 
@@ -410,8 +446,7 @@ headers_frame(State0=#http2_state{content_handlers=Handlers0, commands_queue=Com
 			{setup, Destination=#{host := DestHost, port := DestPort}, TunnelInfo} = Tunnel,
 			%% In the case of CONNECT responses the RealStreamRef is found in TunnelInfo.
 			%% We therefore do not need to call stream_ref/2.
-			%% @todo Maybe we don't need it in TunnelInfo anymore?
-			#{stream_ref := RealStreamRef} = TunnelInfo,
+			RealStreamRef = stream_ref(State, StreamRef),
 			ReplyTo ! {gun_response, self(), RealStreamRef, IsFin, Status, Headers},
 			EvHandlerState = EvHandler:response_headers(#{
 				stream_ref => RealStreamRef,
@@ -428,19 +463,12 @@ headers_frame(State0=#http2_state{content_handlers=Handlers0, commands_queue=Com
 				{P, PO} -> {gun:protocol_handler(P), PO};
 				P -> {gun:protocol_handler(P), #{}}
 			end,
-			%% @todo What about gun_socks_up?
 			%% @todo What about the StateName returned?
 			OriginSocket = #{
 				reply_to => ReplyTo,
 				stream_ref => RealStreamRef
 			},
 			OriginTransport = gun_tcp_proxy,
-			%% @todo Depending on protocol:
-			%% - HTTP/1.1 will need to add the stream_ref in Opts to its StreamRef in messages.
-			%% - HTTP/2 as well
-			%% - raw already uses it
-			%% - ws already uses it (but it's passed slightly differently)
-			%% - socks might not need it? what about gun_socks_up?
 			{_, ProtoState} = Protocol:init(ReplyTo, OriginSocket, OriginTransport,
 				ProtoOpts#{stream_ref => RealStreamRef}),
 			%% @todo EvHandlerState = EvHandler:protocol_changed(#{protocol => Protocol:name()}, EvHandlerState0),
@@ -993,7 +1021,7 @@ stream_info(State, StreamRef) when is_reference(StreamRef) ->
 %% Tunneled streams.
 stream_info(State=#http2_state{transport=Transport}, StreamRefList=[StreamRef|Tail]) ->
 	case get_stream_by_ref(State, StreamRef) of
-		#stream{tunnel={Protocol, ProtoState, #{host := TunnelHost, port := TunnelPort}}} ->
+		#stream{tunnel={Protocol, ProtoState, TunnelInfo=#{host := TunnelHost, port := TunnelPort}}} ->
 			%% We must return the real StreamRef as seen by the user.
 			%% We therefore set it on return, with the outer layer "winning".
 			%%
@@ -1004,7 +1032,8 @@ stream_info(State=#http2_state{transport=Transport}, StreamRefList=[StreamRef|Ta
 				{ok, undefined} ->
 					{ok, undefined};
 				{ok, Info} ->
-					Intermediaries = maps:get(intermediaries, Info, []),
+					Intermediaries1 = maps:get(intermediaries, TunnelInfo, []),
+					Intermediaries2 = maps:get(intermediaries, Info, []),
 					{ok, Info#{
 						ref => StreamRefList,
 						intermediaries => [#{
@@ -1017,7 +1046,7 @@ stream_info(State=#http2_state{transport=Transport}, StreamRefList=[StreamRef|Ta
 								TransportName -> TransportName
 							end,
 							protocol => http2
-						}|Intermediaries]
+						}|Intermediaries1 ++ Intermediaries2]
 					}}
 			end;
 		error ->
diff --git a/test/rfc7540_SUITE.erl b/test/rfc7540_SUITE.erl
index 879c342..ebc5392 100644
--- a/test/rfc7540_SUITE.erl
+++ b/test/rfc7540_SUITE.erl
@@ -559,7 +559,8 @@ do_connect_cowboy(_OriginScheme, OriginTransport, OriginProtocol, _ProxyScheme,
 	{ok, Ref, OriginPort} = do_cowboy_origin(OriginTransport, OriginProtocol),
 	try
 		{ok, ProxyPid, ProxyPort} = do_proxy_start(ProxyTransport, [
-			#proxy_stream{id=1, status=200}
+			#proxy_stream{id=1, status=200},
+			#proxy_stream{id=3, status=299}
 		]),
 		Authority = iolist_to_binary(["localhost:", integer_to_binary(OriginPort)]),
 		{ok, ConnPid} = gun:open("localhost", ProxyPort, #{
@@ -581,6 +582,9 @@ do_connect_cowboy(_OriginScheme, OriginTransport, OriginProtocol, _ProxyScheme,
 		{response, nofin, 200, _} = gun:await(ConnPid, StreamRef),
 		ProxiedStreamRef = gun:get(ConnPid, "/proxied", #{}, #{tunnel => StreamRef}),
 		{response, nofin, 200, _} = gun:await(ConnPid, ProxiedStreamRef),
+		%% We can create more requests on the proxy as well.
+		ProxyStreamRef = gun:get(ConnPid, "/"),
+		{response, fin, 299, _} = gun:await(ConnPid, ProxyStreamRef),
 		gun:close(ConnPid)
 	after
 		cowboy:stop_listener(Ref)
diff --git a/test/socks_SUITE.erl b/test/socks_SUITE.erl
index e692a82..72b038c 100644
--- a/test/socks_SUITE.erl
+++ b/test/socks_SUITE.erl
@@ -413,3 +413,89 @@ do_socks5_through_connect_proxy(OriginScheme, OriginTransport, ProxyTransport) -
 			protocol := socks
 	}]} = gun:info(ConnPid),
 	gun:close(ConnPid).
+
+socks5_tcp_through_h2_connect_tcp_to_tcp_origin(_) ->
+	doc("CONNECT can be used to establish a TCP connection "
+		"to an HTTP/1.1 server via a tunnel going through "
+		"a TCP HTTP/2 proxy followed by a Socks5 proxy."),
+	do_socks5_through_h2_connect_proxy(<<"http">>, tcp, <<"http">>, tcp).
+
+do_socks5_through_h2_connect_proxy(OriginScheme, OriginTransport, ProxyScheme, ProxyTransport) ->
+	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, http),
+	{ok, Proxy1Pid, Proxy1Port} = rfc7540_SUITE:do_proxy_start(ProxyTransport, [
+		{proxy_stream, 1, 200, [], 0, undefined}
+	]),
+	{ok, Proxy2Pid, Proxy2Port} = do_proxy_start(ProxyTransport, none),
+	{ok, ConnPid} = gun:open("localhost", Proxy1Port, #{
+		transport => ProxyTransport,
+		protocols => [http2]
+	}),
+	%% We receive a gun_up first. This is the HTTP proxy.
+	{ok, http2} = gun:await_up(ConnPid),
+	handshake_completed = receive_from(Proxy1Pid),
+	Authority1 = iolist_to_binary(["localhost:", integer_to_binary(Proxy2Port)]),
+	StreamRef = gun:connect(ConnPid, #{
+		host => "localhost",
+		port => Proxy2Port,
+		transport => ProxyTransport,
+		protocols => [{socks, #{
+			host => "localhost",
+			port => OriginPort,
+			transport => OriginTransport
+		}}]
+	}),
+	{request, #{
+		<<":method">> := <<"CONNECT">>,
+		<<":authority">> := Authority1
+	}} = receive_from(Proxy1Pid),
+	{response, nofin, 200, _} = gun:await(ConnPid, StreamRef),
+	%% We receive a gun_socks_up afterwards. This is the origin HTTP server.
+	{ok, http} = gun:await_up(ConnPid),
+	%% The second proxy receives a Socks5 auth/connect request.
+	{auth_methods, 1, [none]} = receive_from(Proxy2Pid),
+	{connect, <<"localhost">>, OriginPort} = receive_from(Proxy2Pid),
+	handshake_completed = receive_from(OriginPid),
+	ProxiedStreamRef = gun:get(ConnPid, "/proxied", #{}, #{tunnel => StreamRef}),
+	Authority2 = iolist_to_binary(["localhost:", integer_to_binary(OriginPort)]),
+	Data = receive_from(OriginPid),
+	Lines = binary:split(Data, <<"\r\n">>, [global]),
+	[<<"host: ", Authority2/bits>>] = [L || <<"host: ", _/bits>> = L <- Lines],
+	#{
+		transport := ProxyTransport,
+		protocol := http2,
+		origin_scheme := ProxyScheme,
+		origin_host := "localhost",
+		origin_port := Proxy1Port,
+		intermediaries := [] %% Intermediaries are specific to the CONNECT stream.
+	} = gun:info(ConnPid),
+	{ok, #{
+		ref := StreamRef,
+		reply_to := Self,
+		state := running,
+		tunnel := #{
+			transport := OriginTransport,
+			protocol := http,
+			origin_scheme := OriginScheme,
+			origin_host := "localhost",
+			origin_port := OriginPort
+		}
+	}} = gun:stream_info(ConnPid, StreamRef),
+	{ok, #{
+		ref := ProxiedStreamRef,
+		reply_to := Self,
+		state := running,
+		intermediaries := [#{
+			type := connect,
+			host := "localhost",
+			port := Proxy1Port,
+			transport := ProxyTransport,
+			protocol := http2
+		}, #{
+			type := socks5,
+			host := "localhost",
+			port := Proxy2Port,
+			transport := ProxyTransport,
+			protocol := socks
+		}]
+	}} = gun:stream_info(ConnPid, ProxiedStreamRef),
+	gun:close(ConnPid).