Fix crash when TLS connection closes very early

And ensure that we don't infinite loop when retries are
enabled, by decrementing the retry count instead of using
a new one.

Also check for ssl:negotiated_protocol {error,closed}
which was possible but was not documented in OTP before
this change.

Thanks @voluntas for the help.

9 files changed, 213 insertions, 77 deletions

diff --git a/src/gun.erl b/src/gun.erl
index c29bd13..db6965b 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -105,6 +105,7 @@
 -export([initial_tls_handshake/3]).
 -export([ensure_alpn_sni/3]).
 -export([tls_handshake/3]).
+-export([connected_protocol_init/3]).
 -export([connected/3]).
 -export([connected_data_only/3]).
 -export([connected_no_input/3]).
@@ -1075,8 +1076,9 @@ connecting(_, {retries, Retries, LookupInfo}, State=#state{opts=Opts,
 				socket => Socket,
 				protocol => ProtocolName
 			}, EvHandlerState1),
-			{next_state, connected, State#state{event_handler_state=EvHandlerState},
-				{next_event, internal, {connected, Socket, Protocol}}};
+			{next_state, connected_protocol_init,
+				State#state{event_handler_state=EvHandlerState},
+				{next_event, internal, {connected, Retries, Socket, Protocol}}};
 		{ok, Socket} when Transport =:= gun_tls ->
 			EvHandlerState = EvHandler:connect_end(ConnectEvent#{
 				socket => Socket
@@ -1099,8 +1101,8 @@ initial_tls_handshake(_, {retries, Retries, Socket}, State0=#state{opts=Opts, or
 	},
 	case normal_tls_handshake(Socket, State0, HandshakeEvent, Protocols) of
 		{ok, TLSSocket, Protocol, State} ->
-			{next_state, connected, State,
-				{next_event, internal, {connected, TLSSocket, Protocol}}};
+			{next_state, connected_protocol_init, State,
+				{next_event, internal, {connected, Retries, TLSSocket, Protocol}}};
 		{error, Reason, State} ->
 			{next_state, not_connected, State,
 				{next_event, internal, {retries, Retries, Reason}}}
@@ -1207,13 +1209,26 @@ normal_tls_handshake(Socket, State=#state{
 	EvHandlerState1 = EvHandler:tls_handshake_start(HandshakeEvent, EvHandlerState0),
 	case gun_tls:connect(Socket, TLSOpts, TLSTimeout) of
 		{ok, TLSSocket} ->
-			NewProtocol = gun_protocols:negotiated(ssl:negotiated_protocol(TLSSocket), Protocols),
-			Protocol = gun_protocols:handler(NewProtocol),
-			EvHandlerState = EvHandler:tls_handshake_end(HandshakeEvent#{
-				socket => TLSSocket,
-				protocol => Protocol:name()
-			}, EvHandlerState1),
-			{ok, TLSSocket, NewProtocol, State#state{event_handler_state=EvHandlerState}};
+			%% This call may return {error,closed} when the socket has
+			%% been closed by the peer. This should be very rare (due to
+			%% timing) but can happen for example when client certificates
+			%% were required but not sent or invalid with some servers.
+			case ssl:negotiated_protocol(TLSSocket) of
+				{error, Reason = closed} ->
+					EvHandlerState = EvHandler:tls_handshake_end(HandshakeEvent#{
+						error => Reason
+					}, EvHandlerState1),
+					{error, Reason, State#state{event_handler_state=EvHandlerState}};
+				NegotiatedProtocol ->
+					NewProtocol = gun_protocols:negotiated(NegotiatedProtocol, Protocols),
+					Protocol = gun_protocols:handler(NewProtocol),
+					EvHandlerState = EvHandler:tls_handshake_end(HandshakeEvent#{
+						socket => TLSSocket,
+						protocol => Protocol:name()
+					}, EvHandlerState1),
+					{ok, TLSSocket, NewProtocol,
+						State#state{event_handler_state=EvHandlerState}}
+			end;
 		{error, Reason} ->
 			EvHandlerState = EvHandler:tls_handshake_end(HandshakeEvent#{
 				error => Reason
@@ -1221,6 +1236,29 @@ normal_tls_handshake(Socket, State=#state{
 			{error, Reason, State#state{event_handler_state=EvHandlerState}}
 	end.
 
+connected_protocol_init(internal, {connected, Retries, Socket, NewProtocol},
+		State0=#state{owner=Owner, opts=Opts, transport=Transport}) ->
+	{Protocol, ProtoOpts} = gun_protocols:handler_and_opts(NewProtocol, Opts),
+	case Protocol:init(Owner, Socket, Transport, ProtoOpts) of
+		{error, Reason} ->
+			{next_state, not_connected, State0,
+				{next_event, internal, {retries, Retries, Reason}}};
+		{ok, StateName, ProtoState} ->
+			%% @todo Don't send gun_up and gun_down if active/1 fails here.
+			Owner ! {gun_up, self(), Protocol:name()},
+			State1 = State0#state{socket=Socket, protocol=Protocol, protocol_state=ProtoState},
+			case active(State1) of
+				{ok, State2} ->
+					State = case Protocol:has_keepalive() of
+						true -> keepalive_timeout(State2);
+						false -> State2
+					end,
+					{next_state, StateName, State};
+				Disconnect ->
+					Disconnect
+			end
+	end.
+
 connected_no_input(Type, Event, State) ->
 	handle_common_connected_no_input(Type, Event, ?FUNCTION_NAME, State).
 
@@ -1253,29 +1291,6 @@ connected_ws_only(cast, Msg, _)
 connected_ws_only(Type, Event, State) ->
 	handle_common_connected_no_input(Type, Event, ?FUNCTION_NAME, State).
 
-connected(internal, {connected, Socket, NewProtocol},
-		State0=#state{owner=Owner, opts=Opts, transport=Transport}) ->
-	{Protocol, ProtoOpts} = gun_protocols:handler_and_opts(NewProtocol, Opts),
-	case Protocol:init(Owner, Socket, Transport, ProtoOpts) of
-		Error={error, _} ->
-			%% @todo Don't send gun_up and gun_down if Protocol:init/4 failes here.
-			Owner ! {gun_up, self(), Protocol:name()},
-			disconnect(State0, Error);
-		{ok, StateName, ProtoState} ->
-			%% @todo Don't send gun_up and gun_down if active/1 failes here.
-			Owner ! {gun_up, self(), Protocol:name()},
-			State1 = State0#state{socket=Socket, protocol=Protocol, protocol_state=ProtoState},
-			case active(State1) of
-				{ok, State2} ->
-					State = case Protocol:has_keepalive() of
-						true -> keepalive_timeout(State2);
-						false -> State2
-					end,
-					{next_state, StateName, State};
-				Disconnect ->
-					Disconnect
-			end
-	end;
 %% Public HTTP interface.
 %%
 %% @todo It might be better, internally, to pass around a URIMap
diff --git a/test/gun_SUITE.erl b/test/gun_SUITE.erl
index 6d1b00b..656158e 100644
--- a/test/gun_SUITE.erl
+++ b/test/gun_SUITE.erl
@@ -193,7 +193,7 @@ keepalive_infinity(_) ->
 killed_streams_http(_) ->
 	doc("Ensure completed responses with a connection: close are not considered killed streams."),
 	{ok, _, OriginPort} = init_origin(tcp, http,
-		fun (_, ClientSocket, ClientTransport) ->
+		fun (_, _, ClientSocket, ClientTransport) ->
 			{ok, _} = ClientTransport:recv(ClientSocket, 0, 1000),
 			ClientTransport:send(ClientSocket,
 				"HTTP/1.1 200 OK\r\n"
@@ -272,7 +272,7 @@ reply_to_http2(_) ->
 
 do_reply_to(Protocol) ->
 	{ok, OriginPid, OriginPort} = init_origin(tcp, Protocol,
-		fun(_, ClientSocket, ClientTransport) ->
+		fun(_, _, ClientSocket, ClientTransport) ->
 			{ok, _} = ClientTransport:recv(ClientSocket, 0, infinity),
 			ResponseData = case Protocol of
 				http ->
@@ -474,7 +474,7 @@ server_name_indication_default(_) ->
 do_server_name_indication(Host, Expected, GunOpts) ->
 	Self = self(),
 	{ok, OriginPid, OriginPort} = init_origin(tls, http,
-		fun(_, ClientSocket, _) ->
+		fun(_, _, ClientSocket, _) ->
 			{ok, Info} = ssl:connection_information(ClientSocket),
 			Msg = {sni_hostname, _} = lists:keyfind(sni_hostname, 1, Info),
 			Self ! Msg
@@ -527,7 +527,7 @@ do_shutdown_reason() ->
 stream_info_http(_) ->
 	doc("Ensure the function gun:stream_info/2 works as expected for HTTP/1.1."),
 	{ok, OriginPid, OriginPort} = init_origin(tcp, http,
-		fun(_, ClientSocket, ClientTransport) ->
+		fun(_, _, ClientSocket, ClientTransport) ->
 			%% Wait for the cancel signal.
 			receive cancel -> ok end,
 			%% Then terminate the stream.
@@ -574,7 +574,7 @@ stream_info_http(_) ->
 stream_info_http2(_) ->
 	doc("Ensure the function gun:stream_info/2 works as expected for HTTP/2."),
 	{ok, OriginPid, OriginPort} = init_origin(tcp, http2,
-		fun(_, _, _) -> receive disconnect -> ok end end),
+		fun(_, _, _, _) -> receive disconnect -> ok end end),
 	{ok, Pid} = gun:open("localhost", OriginPort, #{
 		event_handler => {gun_test_event_h, self()},
 		protocols => [http2]
@@ -604,6 +604,72 @@ supervise_false(_) ->
 	[] = [P || {_, P, _, _} <- supervisor:which_children(gun_sup), P =:= Pid],
 	ok.
 
+tls_handshake_error_gun_http2_init_retry_0(_) ->
+	doc("Ensure an early TLS connection close is propagated "
+		"to the user of the connection."),
+	%% The server will immediately close the connection upon
+	%% establishment so that the client's socket is down
+	%% before it attempts to initialise HTTP/2.
+	%%
+	%% We use 'http' for the server to skip the HTTP/2 init
+	%% but the client is connecting using 'http2'.
+	{ok, _, OriginPort} = init_origin(tls, http,
+		fun(_, _, ClientSocket, ClientTransport) ->
+			ClientTransport:close(ClientSocket)
+		end),
+	{ok, ConnPid} = gun:open("localhost", OriginPort, #{
+		event_handler => {gun_test_fun_event_h, #{
+			tls_handshake_end => fun(_, #{socket := _}) ->
+				%% We sleep right after the gun_tls:connect succeeds to make
+				%% sure the next call to the socket fails (as the socket
+				%% should be disconnected at that point by the server because
+				%% we are not sending a certificate). The call we want to fail
+				%% is the sending of the HTTP/2 preface in gun_http2:init.
+				timer:sleep(1000)
+			end
+		}},
+		protocols => [http2],
+		retry => 0,
+		transport => tls
+	}),
+	{error, {down, {shutdown, closed}}} = gun:await_up(ConnPid),
+	gun:close(ConnPid).
+
+tls_handshake_error_gun_http2_init_retry_1(_) ->
+	doc("Ensure an early TLS connection close is propagated "
+		"to the user of the connection and does not result "
+		"in an infinite connect loop."),
+	%% The server will immediately close the connection upon
+	%% establishment so that the client's socket is down
+	%% before it attempts to initialise HTTP/2.
+	%%
+	%% We use 'http' for the server to skip the HTTP/2 init
+	%% but the client is connecting using 'http2'.
+	%%
+	%% We immediately accept another connection to handle
+	%% the coming retry and close that connection again.
+	{ok, _, OriginPort} = init_origin(tls, http,
+		fun(_, ListenSocket, ClientSocket1, ClientTransport) ->
+			ClientTransport:close(ClientSocket1),
+			%% We immediately accept a second connection and close it again.
+			{ok, ClientSocket2t} = ssl:transport_accept(ListenSocket, 5000),
+			{ok, ClientSocket2} = ssl:handshake(ClientSocket2t, 5000),
+			ClientTransport:close(ClientSocket2)
+		end),
+	{ok, ConnPid} = gun:open("localhost", OriginPort, #{
+		event_handler => {gun_test_fun_event_h, #{
+			tls_handshake_end => fun(_, #{socket := _}) ->
+				%% See tls_handshake_error_gun_http2_init_retry_0 for details.
+				timer:sleep(1000)
+			end
+		}},
+		protocols => [http2],
+		retry => 1,
+		transport => tls
+	}),
+	{error, {down, {shutdown, closed}}} = gun:await_up(ConnPid),
+	gun:close(ConnPid).
+
 tls_handshake_timeout(_) ->
 	doc("Ensure an integer value for tls_handshake_timeout is accepted."),
 	do_timeout(tls_handshake_timeout, 1000).
diff --git a/test/gun_test.erl b/test/gun_test.erl
index 5142898..18fcfbf 100644
--- a/test/gun_test.erl
+++ b/test/gun_test.erl
@@ -33,7 +33,7 @@ init_origin(Transport) ->
 	init_origin(Transport, http).
 
 init_origin(Transport, Protocol) ->
-	init_origin(Transport, Protocol, fun loop_origin/3).
+	init_origin(Transport, Protocol, fun loop_origin/4).
 
 init_origin(Transport, Protocol, Fun) ->
 	Pid = spawn_link(?MODULE, init_origin, [self(), Transport, Protocol, Fun]),
@@ -55,7 +55,7 @@ init_origin(Parent, Transport, Protocol, Fun)
 		_ -> ok
 	end,
 	Parent ! {self(), handshake_completed},
-	Fun(Parent, ClientSocket, gen_tcp);
+	Fun(Parent, ListenSocket, ClientSocket, gen_tcp);
 init_origin(Parent, tls, Protocol, Fun) ->
 	Opts0 = ct_helper:get_certs_from_ets(),
 	Opts1 = case Protocol of
@@ -77,7 +77,7 @@ init_origin(Parent, tls, Protocol, Fun) ->
 			ok
 	end,
 	Parent ! {self(), handshake_completed},
-	Fun(Parent, ClientSocket, ssl).
+	Fun(Parent, ListenSocket, ClientSocket, ssl).
 
 http2_handshake(Socket, Transport) ->
 	%% Send a valid preface.
@@ -96,11 +96,11 @@ http2_handshake(Socket, Transport) ->
 	{ok, <<0:24, 4:8, 1:8, 0:32>>} = Transport:recv(Socket, 9, 5000),
 	ok.
 
-loop_origin(Parent, ClientSocket, ClientTransport) ->
+loop_origin(Parent, ListenSocket, ClientSocket, ClientTransport) ->
 	case ClientTransport:recv(ClientSocket, 0, 5000) of
 		{ok, Data} ->
 			Parent ! {self(), Data},
-			loop_origin(Parent, ClientSocket, ClientTransport);
+			loop_origin(Parent, ListenSocket, ClientSocket, ClientTransport);
 		{error, closed} ->
 			ok
 	end.
diff --git a/test/gun_test_fun_event_h.erl b/test/gun_test_fun_event_h.erl
new file mode 100644
index 0000000..4006a69
--- /dev/null
+++ b/test/gun_test_fun_event_h.erl
@@ -0,0 +1,55 @@
+%% Copyright (c) 2023, Loïc Hoguin <[email protected]>
+%%
+%% Permission to use, copy, modify, and/or distribute this software for any
+%% purpose with or without fee is hereby granted, provided that the above
+%% copyright notice and this permission notice appear in all copies.
+%%
+%% THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+%% WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+%% MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+%% ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+%% WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+%% ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+%% OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+-module(gun_test_fun_event_h).
+-compile(export_all).
+-compile(nowarn_export_all).
+
+init(Event, State) -> common(?FUNCTION_NAME, Event, State).
+domain_lookup_start(Event, State) -> common(?FUNCTION_NAME, Event, State).
+domain_lookup_end(Event, State) -> common(?FUNCTION_NAME, Event, State).
+connect_start(Event, State) -> common(?FUNCTION_NAME, Event, State).
+connect_end(Event, State) -> common(?FUNCTION_NAME, Event, State).
+tls_handshake_start(Event, State) -> common(?FUNCTION_NAME, Event, State).
+tls_handshake_end(Event, State) -> common(?FUNCTION_NAME, Event, State).
+request_start(Event, State) -> common(?FUNCTION_NAME, Event, State).
+request_headers(Event, State) -> common(?FUNCTION_NAME, Event, State).
+request_end(Event, State) -> common(?FUNCTION_NAME, Event, State).
+push_promise_start(Event, State) -> common(?FUNCTION_NAME, Event, State).
+push_promise_end(Event, State) -> common(?FUNCTION_NAME, Event, State).
+response_start(Event, State) -> common(?FUNCTION_NAME, Event, State).
+response_inform(Event, State) -> common(?FUNCTION_NAME, Event, State).
+response_headers(Event, State) -> common(?FUNCTION_NAME, Event, State).
+response_trailers(Event, State) -> common(?FUNCTION_NAME, Event, State).
+response_end(Event, State) -> common(?FUNCTION_NAME, Event, State).
+ws_upgrade(Event, State) -> common(?FUNCTION_NAME, Event, State).
+ws_recv_frame_start(Event, State) -> common(?FUNCTION_NAME, Event, State).
+ws_recv_frame_header(Event, State) -> common(?FUNCTION_NAME, Event, State).
+ws_recv_frame_end(Event, State) -> common(?FUNCTION_NAME, Event, State).
+ws_send_frame_start(Event, State) -> common(?FUNCTION_NAME, Event, State).
+ws_send_frame_end(Event, State) -> common(?FUNCTION_NAME, Event, State).
+protocol_changed(Event, State) -> common(?FUNCTION_NAME, Event, State).
+origin_changed(Event, State) -> common(?FUNCTION_NAME, Event, State).
+cancel(Event, State) -> common(?FUNCTION_NAME, Event, State).
+disconnect(Event, State) -> common(?FUNCTION_NAME, Event, State).
+terminate(Event, State) -> common(?FUNCTION_NAME, Event, State).
+
+common(EventType, Event, State=EventFunsMap) ->
+	case EventFunsMap of
+		#{EventType := Fun} ->
+			Fun(EventType, Event),
+			State;
+		_ ->
+			State
+	end.
diff --git a/test/raw_SUITE.erl b/test/raw_SUITE.erl
index 8597674..57c04b9 100644
--- a/test/raw_SUITE.erl
+++ b/test/raw_SUITE.erl
@@ -57,7 +57,7 @@ direct_raw_tls_with_server_side_close(_) ->
 	do_direct_raw(tls, flow_control_disabled, server_side_close).
 
 do_direct_raw(OriginTransport, FlowControl, CloseSide) ->
-	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw, fun do_echo/3),
+	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw, fun do_echo/4),
 	Opts0 = #{
 		transport => OriginTransport,
 		tls_opts => [{verify, verify_none}, {versions, ['tlsv1.2']}],
@@ -121,7 +121,7 @@ socks5_tls_raw_tls(_) ->
 	do_socks5_raw(tls, tls).
 
 do_socks5_raw(OriginTransport, ProxyTransport) ->
-	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw, fun do_echo/3),
+	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw, fun do_echo/4),
 	{ok, ProxyPid, ProxyPort} = socks_SUITE:do_proxy_start(ProxyTransport, none),
 	{ok, ConnPid} = gun:open("localhost", ProxyPort, #{
 		transport => ProxyTransport,
@@ -176,7 +176,7 @@ connect_tls_raw_tls(_) ->
 	do_connect_raw(tls, tls).
 
 do_connect_raw(OriginTransport, ProxyTransport) ->
-	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw, fun do_echo/3),
+	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw, fun do_echo/4),
 	{ok, ProxyPid, ProxyPort} = rfc7231_SUITE:do_proxy_start(ProxyTransport),
 	Authority = iolist_to_binary(["localhost:", integer_to_binary(OriginPort)]),
 	{ok, ConnPid} = gun:open("localhost", ProxyPort, #{
@@ -224,7 +224,7 @@ connect_raw_reply_to(_) ->
 		{data, nofin, <<"Hello world!">>} = gun:await(ConnPid, StreamRef),
 		Self ! {self(), ok}
 	end),
-	{ok, OriginPid, OriginPort} = init_origin(tcp, raw, fun do_echo/3),
+	{ok, OriginPid, OriginPort} = init_origin(tcp, raw, fun do_echo/4),
 	{ok, ProxyPid, ProxyPort} = rfc7231_SUITE:do_proxy_start(tcp),
 	{ok, ConnPid} = gun:open("localhost", ProxyPort),
 	{ok, http} = gun:await_up(ConnPid),
@@ -250,7 +250,7 @@ http11_upgrade_raw_tls(_) ->
 
 do_http11_upgrade_raw(OriginTransport) ->
 	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw,
-		fun (Parent, ClientSocket, ClientTransport) ->
+		fun (Parent, ListenSocket, ClientSocket, ClientTransport) ->
 			%% We skip the request and send a 101 response unconditionally.
 			{ok, _} = ClientTransport:recv(ClientSocket, 0, 5000),
 			ClientTransport:send(ClientSocket,
@@ -258,7 +258,7 @@ do_http11_upgrade_raw(OriginTransport) ->
 				"Connection: upgrade\r\n"
 				"Upgrade: custom/1.0\r\n"
 				"\r\n"),
-			do_echo(Parent, ClientSocket, ClientTransport)
+			do_echo(Parent, ListenSocket, ClientSocket, ClientTransport)
 		end),
 	{ok, ConnPid} = gun:open("localhost", OriginPort, #{
 		transport => OriginTransport,
@@ -296,7 +296,7 @@ http11_upgrade_raw_reply_to(_) ->
 		Self ! {self(), ok}
 	end),
 	{ok, OriginPid, OriginPort} = init_origin(tcp, raw,
-		fun (Parent, ClientSocket, ClientTransport) ->
+		fun (Parent, ListenSocket, ClientSocket, ClientTransport) ->
 			%% We skip the request and send a 101 response unconditionally.
 			{ok, _} = ClientTransport:recv(ClientSocket, 0, 5000),
 			ClientTransport:send(ClientSocket,
@@ -304,7 +304,7 @@ http11_upgrade_raw_reply_to(_) ->
 				"Connection: upgrade\r\n"
 				"Upgrade: custom/1.0\r\n"
 				"\r\n"),
-			do_echo(Parent, ClientSocket, ClientTransport)
+			do_echo(Parent, ListenSocket, ClientSocket, ClientTransport)
 		end),
 	{ok, ConnPid} = gun:open("localhost", OriginPort),
 	{ok, http} = gun:await_up(ConnPid),
@@ -327,7 +327,7 @@ http2_connect_tls_raw_tcp(_) ->
 	do_http2_connect_raw(tcp, <<"https">>, tls).
 
 do_http2_connect_raw(OriginTransport, ProxyScheme, ProxyTransport) ->
-	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw, fun do_echo/3),
+	{ok, OriginPid, OriginPort} = init_origin(OriginTransport, raw, fun do_echo/4),
 	{ok, ProxyPid, ProxyPort} = rfc7540_SUITE:do_proxy_start(ProxyTransport, [
 		{proxy_stream, 1, 200, [], 0, undefined}
 	]),
@@ -379,13 +379,13 @@ do_http2_connect_raw(OriginTransport, ProxyScheme, ProxyTransport) ->
 
 %% The origin server will echo everything back.
 
-do_echo(Parent, ClientSocket, ClientTransport) ->
+do_echo(Parent, ListenSocket, ClientSocket, ClientTransport) ->
 	case ClientTransport:recv(ClientSocket, 0, 5000) of
 		{ok, <<"close">>} ->
 			ok = ClientTransport:close(ClientSocket);
 		{ok, Data} ->
 			ClientTransport:send(ClientSocket, Data),
-			do_echo(Parent, ClientSocket, ClientTransport);
+			do_echo(Parent, ListenSocket, ClientSocket, ClientTransport);
 		{error, closed} ->
 			ok
 	end.
diff --git a/test/rfc7230_SUITE.erl b/test/rfc7230_SUITE.erl
index 95c1eba..4d1e902 100644
--- a/test/rfc7230_SUITE.erl
+++ b/test/rfc7230_SUITE.erl
@@ -79,7 +79,7 @@ transfer_encoding_overrides_content_length(_) ->
 	doc("When both transfer-encoding and content-length are provided, "
 		"content-length must be ignored. (RFC7230 3.3.3)"),
 	{ok, _, OriginPort} = init_origin(tcp, http,
-		fun(_, ClientSocket, ClientTransport) ->
+		fun(_, _, ClientSocket, ClientTransport) ->
 			{ok, _} = ClientTransport:recv(ClientSocket, 0, 1000),
 			ClientTransport:send(ClientSocket,
 				"HTTP/1.1 200 OK\r\n"
diff --git a/test/rfc7540_SUITE.erl b/test/rfc7540_SUITE.erl
index 10dd875..ac88469 100644
--- a/test/rfc7540_SUITE.erl
+++ b/test/rfc7540_SUITE.erl
@@ -176,7 +176,7 @@ authority_default_port_https(_) ->
 authority_ipv6(_) ->
 	doc("When connecting to a server using an IPv6 address the :authority "
 		"pseudo-header must wrap the address with brackets. (RFC7540 8.1.2.3, RFC3986 3.2.2)"),
-	{ok, OriginPid, OriginPort} = init_origin(tcp6, http2, fun(Parent, Socket, Transport) ->
+	{ok, OriginPid, OriginPort} = init_origin(tcp6, http2, fun(Parent, _, Socket, Transport) ->
 		%% Receive the HEADERS frame and send the headers decoded.
 		{ok, <<Len:24, 1:8, _:8, 1:32>>} = Transport:recv(Socket, 9, 1000),
 		{ok, ReqHeadersBlock} = Transport:recv(Socket, Len, 1000),
@@ -205,7 +205,7 @@ authority_other_port_https(_) ->
 	do_authority_port(tls, 80, <<":80">>).
 
 do_authority_port(Transport0, DefaultPort, AuthorityHeaderPort) ->
-	{ok, OriginPid, OriginPort} = init_origin(Transport0, http2, fun(Parent, Socket, Transport) ->
+	{ok, OriginPid, OriginPort} = init_origin(Transport0, http2, fun(Parent, _, Socket, Transport) ->
 		%% Receive the HEADERS frame and send the headers decoded.
 		{ok, <<Len:24, 1:8, _:8, 1:32>>} = Transport:recv(Socket, 9, 1000),
 		{ok, ReqHeadersBlock} = Transport:recv(Socket, Len, 1000),
@@ -235,7 +235,7 @@ prior_knowledge_preface_garbage(_) ->
 		"an invalid preface in the form of garbage when connecting "
 		"using the prior knowledge method. (RFC7540 3.4, RFC7540 3.5)"),
 	%% We use 'http' here because we are going to do the handshake manually.
-	{ok, OriginPid, Port} = init_origin(tcp, http, fun(_, Socket, Transport) ->
+	{ok, OriginPid, Port} = init_origin(tcp, http, fun(_, _, Socket, Transport) ->
 		ok = Transport:send(Socket, <<0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15>>),
 		timer:sleep(100)
 	end),
@@ -257,7 +257,7 @@ prior_knowledge_preface_http1(_) ->
 		"an invalid preface in the form of an HTTP/1.1 response when connecting "
 		"using the prior knowledge method. (RFC7540 3.4, RFC7540 3.5)"),
 	%% We use 'http' here because we are going to do the handshake manually.
-	{ok, OriginPid, Port} = init_origin(tcp, http, fun(_, Socket, Transport) ->
+	{ok, OriginPid, Port} = init_origin(tcp, http, fun(_, _, Socket, Transport) ->
 		ok = Transport:send(Socket, <<
 			"HTTP/1.1 400 Bad Request\r\n"
 			"Connection: close\r\n"
@@ -284,7 +284,7 @@ prior_knowledge_preface_http1_await(_) ->
 		"an invalid preface in the form of an HTTP/1.1 response when connecting "
 		"using the prior knowledge method. (RFC7540 3.4, RFC7540 3.5)"),
 	%% We use 'http' here because we are going to do the handshake manually.
-	{ok, OriginPid, Port} = init_origin(tcp, http, fun(_, Socket, Transport) ->
+	{ok, OriginPid, Port} = init_origin(tcp, http, fun(_, _, Socket, Transport) ->
 		timer:sleep(100),
 		ok = Transport:send(Socket, <<
 			"HTTP/1.1 400 Bad Request\r\n"
@@ -307,7 +307,7 @@ prior_knowledge_preface_other_frame(_) ->
 		"an invalid preface in the form of a non-SETTINGS frame when connecting "
 		"using the prior knowledge method. (RFC7540 3.4, RFC7540 3.5)"),
 	%% We use 'http' here because we are going to do the handshake manually.
-	{ok, OriginPid, Port} = init_origin(tcp, http, fun(_, Socket, Transport) ->
+	{ok, OriginPid, Port} = init_origin(tcp, http, fun(_, _, Socket, Transport) ->
 		ok = Transport:send(Socket, cow_http2:window_update(1)),
 		timer:sleep(100)
 	end),
@@ -327,7 +327,7 @@ prior_knowledge_preface_other_frame(_) ->
 lingering_data_counts_toward_connection_window(_) ->
 	doc("DATA frames received after sending RST_STREAM must be counted "
 		"toward the connection flow-control window. (RFC7540 5.1)"),
-	{ok, OriginPid, Port} = init_origin(tcp, http2, fun(_, Socket, Transport) ->
+	{ok, OriginPid, Port} = init_origin(tcp, http2, fun(_, _, Socket, Transport) ->
 		%% Step 2.
 		%% Receive a HEADERS frame.
 		{ok, <<SkipLen:24, 1:8, _:8, 1:32>>} = Transport:recv(Socket, 9, 1000),
@@ -378,7 +378,7 @@ lingering_data_counts_toward_connection_window(_) ->
 headers_priority_flag(_) ->
 	doc("HEADERS frames may include a PRIORITY flag indicating "
 		"that stream dependency information is attached. (RFC7540 6.2)"),
-	{ok, OriginPid, Port} = init_origin(tcp, http2, fun(_, Socket, Transport) ->
+	{ok, OriginPid, Port} = init_origin(tcp, http2, fun(_, _, Socket, Transport) ->
 		%% Receive a HEADERS frame.
 		{ok, <<_:24, 1:8, _:8, 1:32>>} = Transport:recv(Socket, 9, 1000),
 		%% Send a HEADERS frame with PRIORITY back.
@@ -414,7 +414,7 @@ settings_ack_timeout(_) ->
 	doc("Failure to acknowledge the client's SETTINGS frame "
 		"results in a SETTINGS_TIMEOUT connection error. (RFC7540 6.5.3)"),
 	%% We use 'http' here because we are going to do the handshake manually.
-	{ok, _, Port} = init_origin(tcp, http, fun(_, Socket, Transport) ->
+	{ok, _, Port} = init_origin(tcp, http, fun(_, _, Socket, Transport) ->
 		%% Send a valid preface.
 		ok = Transport:send(Socket, cow_http2:settings(#{})),
 		%% Receive the fixed sequence from the preface.
@@ -471,7 +471,7 @@ keepalive_tolerance_ping_ack_timeout(_) ->
 
 do_ping_ack_loop_fun() ->
 	%% Receive ping, sync with parent, send ping ack, loop.
-	fun Loop(Parent, Socket, Transport) ->
+	fun Loop(Parent, ListenSocket, Socket, Transport) ->
 		{ok, Data} = Transport:recv(Socket, 9, infinity),
 		<<Len:24, 6:8, %% PING
 			0:8, %% Flags
@@ -486,7 +486,7 @@ do_ping_ack_loop_fun() ->
 					0:1, 0:31, Payload/binary>>,
 				ok = Transport:send(Socket, Ack)
 		end,
-		Loop(Parent, Socket, Transport)
+		Loop(Parent, ListenSocket, Socket, Transport)
 	end.
 
 connect_http_via_h2c(_) ->
@@ -530,7 +530,7 @@ connect_h2_via_h2(_) ->
 	do_connect_http(<<"https">>, tls, http2, <<"https">>, tls).
 
 do_origin_fun(http) ->
-	fun(Parent, Socket, Transport) ->
+	fun(Parent, ListenSocket, Socket, Transport) ->
 		%% Receive the request-line and headers, parse and send them.
 		{ok, Data} = Transport:recv(Socket, 0, 5000),
 		{Method, Target, 'HTTP/1.1', Rest} = cow_http:parse_request_line(Data),
@@ -542,16 +542,16 @@ do_origin_fun(http) ->
 			<<":method">> => Method,
 			<<":path">> => Target
 		}},
-		gun_test:loop_origin(Parent, Socket, Transport)
+		gun_test:loop_origin(Parent, ListenSocket, Socket, Transport)
 	end;
 do_origin_fun(http2) ->
-	fun(Parent, Socket, Transport) ->
+	fun(Parent, ListenSocket, Socket, Transport) ->
 		%% Receive the HEADERS frame and send the headers decoded.
 		{ok, <<Len:24, 1:8, _:8, 1:32>>} = Transport:recv(Socket, 9, 1000),
 		{ok, ReqHeadersBlock} = Transport:recv(Socket, Len, 1000),
 		{ReqHeaders, _} = cow_hpack:decode(ReqHeadersBlock),
 		Parent ! {self(), maps:from_list(ReqHeaders)},
-		gun_test:loop_origin(Parent, Socket, Transport)
+		gun_test:loop_origin(Parent, ListenSocket, Socket, Transport)
 	end.
 
 do_connect_http(OriginScheme, OriginTransport, OriginProtocol, ProxyScheme, ProxyTransport) ->
@@ -716,7 +716,7 @@ do_cowboy_origin(OriginTransport, OriginProtocol) ->
 connect_handshake_timeout(_) ->
 	doc("HTTP/2 timeouts are properly routed to the appropriate "
 		"tunnel layer. (RFC7540 3.5, RFC7540 8.3)"),
-	{ok, _, OriginPort} = init_origin(tcp, raw, fun(_, _, _) ->
+	{ok, _, OriginPort} = init_origin(tcp, raw, fun(_, _, _, _) ->
 		timer:sleep(5000)
 	end),
 	{ok, ProxyPid, ProxyPort} = do_proxy_start(tcp, [
diff --git a/test/shutdown_SUITE.erl b/test/shutdown_SUITE.erl
index 20ca879..77a2903 100644
--- a/test/shutdown_SUITE.erl
+++ b/test/shutdown_SUITE.erl
@@ -301,7 +301,7 @@ http1_response_connection_close_delayed_body(_) ->
 	doc("HTTP/1.1: Confirm that requests initiated when Gun has received a "
 		"connection: close response header fail immediately if retry "
 		"is disabled, without waiting for the response body."),
-	ServerFun = fun(_Parent, ClientSocket, gen_tcp) ->
+	ServerFun = fun(_, _, ClientSocket, gen_tcp) ->
 		try
 			{ok, Req} = gen_tcp:recv(ClientSocket, 0, 5000),
 			<<"GET / HTTP/1.1\r\n", _/binary>> = Req,
@@ -406,7 +406,7 @@ http2_server_goaway_no_streams(_) ->
 	doc("HTTP/2: Confirm that the Gun process shuts down gracefully "
 		"when receiving a GOAWAY frame with no active streams and "
 		"retry is disabled."),
-	{ok, OriginPid, Port} = init_origin(tcp, http2, fun(_, Socket, Transport) ->
+	{ok, OriginPid, Port} = init_origin(tcp, http2, fun(_, _, Socket, Transport) ->
 		receive go_away -> ok end,
 		Transport:send(Socket, cow_http2:goaway(0, no_error, <<>>)),
 		timer:sleep(500)
@@ -425,7 +425,7 @@ http2_server_goaway_one_stream(_) ->
 	doc("HTTP/2: Confirm that the Gun process shuts down gracefully "
 		"when receiving a GOAWAY frame with one active stream and "
 		"retry is disabled."),
-	{ok, OriginPid, OriginPort} = init_origin(tcp, http2, fun(_, Socket, Transport) ->
+	{ok, OriginPid, OriginPort} = init_origin(tcp, http2, fun(_, _, Socket, Transport) ->
 		%% Receive a HEADERS frame.
 		{ok, <<SkipLen:24, 1:8, _:8, 1:32>>} = Transport:recv(Socket, 9, 1000),
 		%% Skip the header.
@@ -459,7 +459,7 @@ http2_server_goaway_many_streams(_) ->
 	doc("HTTP/2: Confirm that the Gun process shuts down gracefully "
 		"when receiving a GOAWAY frame with many active streams and "
 		"retry is disabled."),
-	{ok, OriginPid, OriginPort} = init_origin(tcp, http2, fun(_, Socket, Transport) ->
+	{ok, OriginPid, OriginPort} = init_origin(tcp, http2, fun(_, _, Socket, Transport) ->
 		%% Stream 1.
 		%% Receive a HEADERS frame.
 		{ok, <<SkipLen1:24, 1:8, _:8, 1:32>>} = Transport:recv(Socket, 9, 1000),
diff --git a/test/tunnel_SUITE.erl b/test/tunnel_SUITE.erl
index 39f7e45..14fb2ef 100644
--- a/test/tunnel_SUITE.erl
+++ b/test/tunnel_SUITE.erl
@@ -812,7 +812,7 @@ do_doc(Type, Endpoint) ->
 
 do_origin_start(Type) when Type =:= raw; Type =:= rawtls ->
 	{Transport, Protocol} = do_type(Type),
-	gun_test:init_origin(Transport, Protocol, fun raw_SUITE:do_echo/3);
+	gun_test:init_origin(Transport, Protocol, fun raw_SUITE:do_echo/4);
 do_origin_start(Type) ->
 	{Transport, Protocol} = do_type(Type),
 	rfc7540_SUITE:do_cowboy_origin(Transport, Protocol).