Make CONNECT responses produce a response_end event

While the stream has not ended, the response has.

3 files changed, 23 insertions, 22 deletions

diff --git a/src/gun_http.erl b/src/gun_http.erl
index 0eeca05..a2c3303 100644
--- a/src/gun_http.erl
+++ b/src/gun_http.erl
@@ -324,6 +324,10 @@ handle_connect(Rest, State=#http_state{
 		status => Status,
 		headers => Headers
 	}, EvHandlerState0),
+	EvHandlerState = EvHandler:response_end(#{
+		stream_ref => RealStreamRef,
+		reply_to => ReplyTo
+	}, EvHandlerState1),
 	%% We expect there to be no additional data after the CONNECT response.
 	%% @todo That's probably wrong.
 	<<>> = Rest,
@@ -342,7 +346,7 @@ handle_connect(Rest, State=#http_state{
 			{[
 				{origin, <<"https">>, NewHost, NewPort, connect},
 				{tls_handshake, HandshakeEvent, Protocols, ReplyTo}
-			], CookieStore, EvHandlerState1};
+			], CookieStore, EvHandlerState};
 		_ ->
 			[NewProtocol0] = maps:get(protocols, Destination, [http]),
 			NewProtocol = gun_protocols:add_stream_ref(NewProtocol0, RealStreamRef),
@@ -351,7 +355,7 @@ handle_connect(Rest, State=#http_state{
 			{[
 				{origin, <<"http">>, NewHost, NewPort, connect},
 				{switch_protocol, NewProtocol, ReplyTo}
-			], CookieStore, EvHandlerState1}
+			], CookieStore, EvHandlerState}
 	end.
 
 %% @todo We probably shouldn't send info messages if the stream is not alive.
diff --git a/src/gun_http2.erl b/src/gun_http2.erl
index 8f0632b..04d4de5 100644
--- a/src/gun_http2.erl
+++ b/src/gun_http2.erl
@@ -517,7 +517,11 @@ headers_frame_connect(State=#http2_state{transport=Transport, opts=Opts, tunnel_
 		status => Status,
 		headers => Headers
 	}, EvHandlerState0),
-	EvHandlerState2 = EvHandler:origin_changed(#{
+	EvHandlerState2 = EvHandler:response_end(#{
+		stream_ref => RealStreamRef,
+		reply_to => ReplyTo
+	}, EvHandlerState1),
+	EvHandlerState3 = EvHandler:origin_changed(#{
 		stream_ref => RealStreamRef,
 		type => connect,
 		origin_scheme => case Destination of
@@ -526,7 +530,7 @@ headers_frame_connect(State=#http2_state{transport=Transport, opts=Opts, tunnel_
 		end,
 		origin_host => DestHost,
 		origin_port => DestPort
-	}, EvHandlerState1),
+	}, EvHandlerState2),
 	ContinueStreamRef = continue_stream_ref(State, StreamRef),
 	OriginSocket = #{
 		gun_pid => self(),
@@ -576,7 +580,7 @@ headers_frame_connect(State=#http2_state{transport=Transport, opts=Opts, tunnel_
 			}
 	end,
 	{tunnel, ProtoState, EvHandlerState} = Proto:init(
-		ReplyTo, OriginSocket, gun_tcp_proxy, ProtoOpts, EvHandler, EvHandlerState2),
+		ReplyTo, OriginSocket, gun_tcp_proxy, ProtoOpts, EvHandler, EvHandlerState3),
 	{store_stream(State, Stream#stream{tunnel=Tunnel#tunnel{state=established,
 		info=TunnelInfo, protocol=Proto, protocol_state=ProtoState}}),
 		EvHandlerState}.
diff --git a/test/event_SUITE.erl b/test/event_SUITE.erl
index ca7c55d..236938c 100644
--- a/test/event_SUITE.erl
+++ b/test/event_SUITE.erl
@@ -1142,11 +1142,10 @@ do_response_end_connect(Config, EventName, Path) ->
 		port => OriginPort,
 		protocols => [Protocol]
 	}, []),
-	%% @todo Figure out whether the response should end when the tunnel is established.
-%	#{
-%		stream_ref := StreamRef1,
-%		reply_to := ReplyTo
-%	} = do_receive_event(EventName),
+	#{
+		stream_ref := StreamRef1,
+		reply_to := ReplyTo
+	} = do_receive_event(EventName),
 	{response, fin, 200, _} = gun:await(ConnPid, StreamRef1),
 	{up, Protocol} = gun:await(ConnPid, StreamRef1),
 	StreamRef2 = gun:get(ConnPid, Path, [{<<"te">>, <<"trailers">>}], #{tunnel => StreamRef1}),
@@ -1195,11 +1194,10 @@ http1_response_end_body_close(Config) ->
 %		port => OriginPort,
 %		protocols => [{http, #{version => 'HTTP/1.0'}}]
 %	}, []),
-%	%% @todo Figure out whether the response should end when the tunnel is established.
-%%	#{
-%%		stream_ref := StreamRef1,
-%%		reply_to := ReplyTo
-%%	} = do_receive_event(EventName),
+%	#{
+%		stream_ref := StreamRef1,
+%		reply_to := ReplyTo
+%	} = do_receive_event(response_end),
 %	{response, fin, 200, _} = gun:await(ConnPid, StreamRef1),
 %	{up, http} = gun:await(ConnPid, StreamRef1),
 %	StreamRef2 = gun:get(ConnPid, "/stream", [], #{tunnel => StreamRef1}),
@@ -1367,13 +1365,8 @@ do_ws_upgrade_all_events_connect(Config, ProxyProtocol) ->
 	_ = do_receive_event(request_headers),
 	_ = do_receive_event(request_end),
 	_ = do_receive_event(response_start),
-	case OriginProtocol of
-		http -> ok;
-		http2 ->
-			_ = do_receive_event(response_headers),
-%			_ = do_receive_event(response_end), @todo Probably should response_end CONNECT responses for both protocols.
-			ok
-	end,
+	_ = do_receive_event(response_headers),
+	_ = do_receive_event(response_end),
 	_ = do_receive_event(protocol_changed),
 	%% Check the Websocket events.
 	StreamRef2 = gun:ws_upgrade(ConnPid, "/ws", [], #{tunnel => StreamRef1}),