From d7481d5f593f48327d16c29b48a74e49c7581e0a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Hoguin?= <essen@ninenines.eu>
Date: Fri, 17 Jul 2020 14:05:43 +0200
Subject: Improve HTTP/2 CONNECT to non-HTTP origin

Now has a proper StreamRef given to it by the CONNECT stream.
---
 src/gun_http2.erl  |  4 +++-
 src/gun_raw.erl    | 10 ++++++----
 test/raw_SUITE.erl | 43 +++++++++++++++++++++++++------------------
 3 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/src/gun_http2.erl b/src/gun_http2.erl
index 0dad3d9..e1a9873 100644
--- a/src/gun_http2.erl
+++ b/src/gun_http2.erl
@@ -407,7 +407,9 @@ headers_frame(State0=#http2_state{content_handlers=Handlers0, commands_queue=Com
 				stream_ref => StreamRef
 			},
 			OriginTransport = gun_tcp_proxy,
-			{_, ProtoState} = Protocol:init(ReplyTo, OriginSocket, OriginTransport, ProtoOpts),
+			{_, ProtoState} = Protocol:init(ReplyTo, OriginSocket, OriginTransport,
+				%% @todo We are giving the wrong StreamRef we need to give the list (if any).
+				ProtoOpts#{stream_ref => StreamRef}),
 			%% @todo EvHandlerState = EvHandler:protocol_changed(#{protocol => Protocol:name()}, EvHandlerState0),
 			%% @todo What about keepalive?
 			{store_stream(State, Stream#stream{tunnel={Protocol, ProtoState,
diff --git a/src/gun_raw.erl b/src/gun_raw.erl
index bef2c1c..369353e 100644
--- a/src/gun_raw.erl
+++ b/src/gun_raw.erl
@@ -26,6 +26,7 @@
 %% @todo down
 
 -record(raw_state, {
+	ref :: reference() | undefined,
 	reply_to :: pid(),
 	socket :: inet:socket() | ssl:sslsocket(),
 	transport :: module()
@@ -39,12 +40,13 @@ name() -> raw.
 opts_name() -> raw_opts.
 has_keepalive() -> false.
 
-init(ReplyTo, Socket, Transport, _Opts) ->
-	{connected_data_only, #raw_state{reply_to=ReplyTo, socket=Socket, transport=Transport}}.
+init(ReplyTo, Socket, Transport, Opts) ->
+	StreamRef = maps:get(stream_ref, Opts, undefined),
+	{connected_data_only, #raw_state{ref=StreamRef, reply_to=ReplyTo, socket=Socket, transport=Transport}}.
 
-handle(Data, State=#raw_state{reply_to=ReplyTo}, _, EvHandlerState) ->
+handle(Data, State=#raw_state{ref=StreamRef, reply_to=ReplyTo}, _, EvHandlerState) ->
 	%% When we take over the entire connection there is no stream reference.
-	ReplyTo ! {gun_data, self(), undefined, nofin, Data},
+	ReplyTo ! {gun_data, self(), StreamRef, nofin, Data},
 	{{state, State}, EvHandlerState}.
 
 %% We can always close immediately.
diff --git a/test/raw_SUITE.erl b/test/raw_SUITE.erl
index c3eec1d..29717e6 100644
--- a/test/raw_SUITE.erl
+++ b/test/raw_SUITE.erl
@@ -188,13 +188,13 @@ connect_raw_reply_to(_) ->
 
 h2_connect_tcp_raw_tcp(_) ->
 	doc("Use CONNECT over clear HTTP/2 to connect to a remote endpoint using the raw protocol over TCP."),
-	do_h2_connect_raw(tcp, tcp).
+	do_h2_connect_raw(tcp, <<"http">>, tcp).
 
 h2_connect_tls_raw_tcp(_) ->
 	doc("Use CONNECT over secure HTTP/2 to connect to a remote endpoint using the raw protocol over TCP."),
-	do_h2_connect_raw(tcp, tls).
+	do_h2_connect_raw(tcp, <<"https">>, tls).
 
-do_h2_connect_raw(OriginTransport, ProxyTransport) ->
+do_h2_connect_raw(OriginTransport, ProxyScheme, ProxyTransport) ->
 	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw, fun do_echo/3),
 	{ok, ProxyPid, ProxyPort} = rfc7540_SUITE:do_proxy_start(ProxyTransport, [
 		{proxy_stream, 1, 200, [], 0, undefined}
@@ -219,21 +219,28 @@ do_h2_connect_raw(OriginTransport, ProxyTransport) ->
 	{response, nofin, 200, _} = gun:await(ConnPid, StreamRef),
 	handshake_completed = receive_from(OriginPid),
 	gun:data(ConnPid, StreamRef, nofin, <<"Hello world!">>),
-	{data, nofin, <<"Hello world!">>} = gun:await(ConnPid, undefined),
-%% @todo
-%	#{
-%		transport := OriginTransport,
-%		protocol := raw,
-%		origin_scheme := _, %% @todo This should be 'undefined'.
-%		origin_host := "localhost",
-%		origin_port := OriginPort,
-%		intermediaries := [#{
-%			type := connect,
-%			host := "localhost",
-%			port := ProxyPort,
-%			transport := ProxyTransport,
-%			protocol := http
-%	}]} = gun:info(ConnPid),
+	{data, nofin, <<"Hello world!">>} = gun:await(ConnPid, StreamRef),
+	#{
+		transport := ProxyTransport,
+		protocol := http2,
+		origin_scheme := ProxyScheme,
+		origin_host := "localhost",
+		origin_port := ProxyPort,
+		intermediaries := []
+	} = gun:info(ConnPid),
+	Self = self(),
+	{ok, #{
+		ref := StreamRef,
+		reply_to := Self,
+		state := running,
+		tunnel := #{
+			transport := OriginTransport,
+			protocol := raw,
+			origin_scheme := _, %% @todo This should be 'undefined'.
+			origin_host := "localhost",
+			origin_port := OriginPort
+		}
+	}} = gun:stream_info(ConnPid, StreamRef),
 	gun:close(ConnPid).
 
 http11_upgrade_raw_tcp(_) ->
-- 
cgit v1.2.3