Replace gun_tunnel_up/3 message with /4 variant

Also fixes all the tests. Lots of work remain around protocols
(how best to pass the base stream_ref to them? maybe the current
solution, maybe a new argument to Protocol:init) and strengthen
the concept of stream_ref, at least with its own type.

3 files changed, 23 insertions, 12 deletions

diff --git a/test/raw_SUITE.erl b/test/raw_SUITE.erl
index bd34ebb..d909dce 100644
--- a/test/raw_SUITE.erl
+++ b/test/raw_SUITE.erl
@@ -87,7 +87,7 @@ do_socks5_raw(OriginTransport, ProxyTransport) ->
 	}),
 	%% We receive a gun_up and a gun_tunnel_up.
 	{ok, socks} = gun:await_up(ConnPid),
-	{ok, raw} = gun:await_up(ConnPid),
+	{up, raw} = gun:await(ConnPid, undefined),
 	%% The proxy received two packets.
 	{auth_methods, 1, [none]} = receive_from(ProxyPid),
 	{connect, <<"localhost">>, OriginPort} = receive_from(ProxyPid),
@@ -141,9 +141,9 @@ do_connect_raw(OriginTransport, ProxyTransport) ->
 	{request, <<"CONNECT">>, Authority, 'HTTP/1.1', _} = receive_from(ProxyPid),
 	{response, fin, 200, _} = gun:await(ConnPid, StreamRef), %% @todo Why fin?
 	handshake_completed = receive_from(OriginPid),
-	%% When we take over the entire connection there is no stream reference.
-	gun:data(ConnPid, undefined, nofin, <<"Hello world!">>),
-	{data, nofin, <<"Hello world!">>} = gun:await(ConnPid, undefined),
+	{up, raw} = gun:await(ConnPid, StreamRef),
+	gun:data(ConnPid, StreamRef, nofin, <<"Hello world!">>),
+	{data, nofin, <<"Hello world!">>} = gun:await(ConnPid, StreamRef),
 	#{
 		transport := OriginTransport,
 		protocol := raw,
@@ -166,8 +166,9 @@ connect_raw_reply_to(_) ->
 	ReplyTo = spawn(fun() ->
 		{ConnPid, StreamRef} = receive Msg -> Msg after 1000 -> error(timeout) end,
 		{response, fin, 200, _} = gun:await(ConnPid, StreamRef),
+		{up, raw} = gun:await(ConnPid, StreamRef),
 		Self ! {self(), ready},
-		{data, nofin, <<"Hello world!">>} = gun:await(ConnPid, undefined),
+		{data, nofin, <<"Hello world!">>} = gun:await(ConnPid, StreamRef),
 		Self ! {self(), ok}
 	end),
 	{ok, OriginPid, OriginPort} = init_origin(tcp, raw, fun do_echo/3),
@@ -183,7 +184,7 @@ connect_raw_reply_to(_) ->
 	{request, <<"CONNECT">>, _, 'HTTP/1.1', _} = receive_from(ProxyPid),
 	handshake_completed = receive_from(OriginPid),
 	receive {ReplyTo, ready} -> ok after 1000 -> error(timeout) end,
-	gun:data(ConnPid, undefined, nofin, <<"Hello world!">>),
+	gun:data(ConnPid, StreamRef, nofin, <<"Hello world!">>),
 	receive {ReplyTo, ok} -> gun:close(ConnPid) after 1000 -> error(timeout) end.
 
 http11_upgrade_raw_tcp(_) ->
diff --git a/test/rfc7231_SUITE.erl b/test/rfc7231_SUITE.erl
index c2f3da6..5ef37bd 100644
--- a/test/rfc7231_SUITE.erl
+++ b/test/rfc7231_SUITE.erl
@@ -168,7 +168,9 @@ do_connect_http(OriginScheme, OriginTransport, ProxyTransport) ->
 	}),
 	{request, <<"CONNECT">>, Authority, 'HTTP/1.1', _} = receive_from(ProxyPid),
 	{response, fin, 200, _} = gun:await(ConnPid, StreamRef),
+	%% @todo Do we still need these handshake_completed messages?
 	handshake_completed = receive_from(OriginPid),
+	{up, http} = gun:await(ConnPid, StreamRef),
 	_ = gun:get(ConnPid, "/proxied"),
 	Data = receive_from(OriginPid),
 	Lines = binary:split(Data, <<"\r\n">>, [global]),
@@ -223,6 +225,7 @@ do_connect_h2(OriginScheme, OriginTransport, ProxyTransport) ->
 	{request, <<"CONNECT">>, Authority, 'HTTP/1.1', _} = receive_from(ProxyPid),
 	{response, fin, 200, _} = gun:await(ConnPid, StreamRef),
 	handshake_completed = receive_from(OriginPid),
+	{up, http2} = gun:await(ConnPid, StreamRef),
 	_ = gun:get(ConnPid, "/proxied"),
 	<<_:24, 1:8, _/bits>> = receive_from(OriginPid),
 	#{
@@ -268,15 +271,17 @@ do_connect_through_multiple_proxies(OriginScheme, OriginTransport, ProxiesTransp
 	}),
 	{request, <<"CONNECT">>, Authority1, 'HTTP/1.1', _} = receive_from(Proxy1Pid),
 	{response, fin, 200, _} = gun:await(ConnPid, StreamRef1),
+	{up, http} = gun:await(ConnPid, StreamRef1),
 	Authority2 = iolist_to_binary(["localhost:", integer_to_binary(OriginPort)]),
 	StreamRef2 = gun:connect(ConnPid, #{
 		host => "localhost",
 		port => OriginPort,
 		transport => OriginTransport
-	}),
+	}, [], #{tunnel => StreamRef1}),
 	{request, <<"CONNECT">>, Authority2, 'HTTP/1.1', _} = receive_from(Proxy2Pid),
 	{response, fin, 200, _} = gun:await(ConnPid, StreamRef2),
 	handshake_completed = receive_from(OriginPid),
+	{up, http} = gun:await(ConnPid, StreamRef2),
 	_ = gun:get(ConnPid, "/proxied"),
 	Data = receive_from(OriginPid),
 	Lines = binary:split(Data, <<"\r\n">>, [global]),
@@ -317,6 +322,7 @@ connect_delay(_) ->
 	{request, <<"CONNECT">>, Authority, 'HTTP/1.1', _} = receive_from(ProxyPid, 3000),
 	{response, fin, 201, _} = gun:await(ConnPid, StreamRef),
 	handshake_completed = receive_from(OriginPid),
+	{up, http} = gun:await(ConnPid, StreamRef),
 	_ = gun:get(ConnPid, "/proxied"),
 	Data = receive_from(OriginPid),
 	Lines = binary:split(Data, <<"\r\n">>, [global]),
@@ -351,6 +357,7 @@ connect_response_201(_) ->
 	{request, <<"CONNECT">>, Authority, 'HTTP/1.1', _} = receive_from(ProxyPid),
 	{response, fin, 201, _} = gun:await(ConnPid, StreamRef),
 	handshake_completed = receive_from(OriginPid),
+	{up, http} = gun:await(ConnPid, StreamRef),
 	_ = gun:get(ConnPid, "/proxied"),
 	Data = receive_from(OriginPid),
 	Lines = binary:split(Data, <<"\r\n">>, [global]),
@@ -492,6 +499,7 @@ connect_response_ignore_transfer_encoding(_) ->
 	{request, <<"CONNECT">>, Authority, 'HTTP/1.1', _} = receive_from(ProxyPid),
 	{response, fin, 200, Headers} = gun:await(ConnPid, StreamRef),
 	handshake_completed = receive_from(OriginPid),
+	{up, http} = gun:await(ConnPid, StreamRef),
 	_ = gun:get(ConnPid, "/proxied"),
 	Data = receive_from(OriginPid),
 	Lines = binary:split(Data, <<"\r\n">>, [global]),
@@ -514,6 +522,7 @@ connect_response_ignore_content_length(_) ->
 	{request, <<"CONNECT">>, Authority, 'HTTP/1.1', _} = receive_from(ProxyPid),
 	{response, fin, 200, Headers} = gun:await(ConnPid, StreamRef),
 	handshake_completed = receive_from(OriginPid),
+	{up, http} = gun:await(ConnPid, StreamRef),
 	_ = gun:get(ConnPid, "/proxied"),
 	Data = receive_from(OriginPid),
 	Lines = binary:split(Data, <<"\r\n">>, [global]),
diff --git a/test/socks_SUITE.erl b/test/socks_SUITE.erl
index 436ea3f..3b8c822 100644
--- a/test/socks_SUITE.erl
+++ b/test/socks_SUITE.erl
@@ -224,7 +224,7 @@ do_socks5(OriginScheme, OriginTransport, OriginProtocol, ProxyTransport, SocksAu
 	}),
 	%% We receive a gun_up and a gun_tunnel_up.
 	{ok, socks} = gun:await_up(ConnPid),
-	{ok, OriginProtocol} = gun:await_up(ConnPid),
+	{up, OriginProtocol} = gun:await(ConnPid, undefined),
 	%% The proxy received two packets.
 	AuthMethod = do_auth_method(SocksAuth),
 	{auth_methods, 1, [AuthMethod]} = receive_from(ProxyPid),
@@ -302,8 +302,8 @@ do_socks5_through_multiple_proxies(OriginScheme, OriginTransport, ProxyTransport
 	}),
 	%% We receive a gun_up and two gun_tunnel_up.
 	{ok, socks} = gun:await_up(ConnPid),
-	{ok, socks} = gun:await_up(ConnPid),
-	{ok, http} = gun:await_up(ConnPid),
+	{up, socks} = gun:await(ConnPid, undefined),
+	{up, http} = gun:await(ConnPid, undefined),
 	%% The first proxy received two packets.
 	{auth_methods, 1, [none]} = receive_from(Proxy1Pid),
 	{connect, <<"localhost">>, Proxy2Port} = receive_from(Proxy1Pid),
@@ -382,8 +382,9 @@ do_socks5_through_connect_proxy(OriginScheme, OriginTransport, ProxyTransport) -
 	}),
 	{request, <<"CONNECT">>, Authority1, 'HTTP/1.1', _} = receive_from(Proxy1Pid),
 	{response, fin, 200, _} = gun:await(ConnPid, StreamRef),
-	%% We receive a gun_tunnel_up afterwards. This is the origin HTTP server.
-	{ok, http} = gun:await_up(ConnPid),
+	%% We receive two gun_tunnel_up messages. First the SOCKS server and then the origin HTTP server.
+	{up, socks} = gun:await(ConnPid, StreamRef),
+	{up, http} = gun:await(ConnPid, StreamRef),
 	%% The second proxy receives a Socks5 auth/connect request.
 	{auth_methods, 1, [none]} = receive_from(Proxy2Pid),
 	{connect, <<"localhost">>, OriginPort} = receive_from(Proxy2Pid),