Implement user pings for tunnels

8 files changed, 61 insertions, 31 deletions

diff --git a/src/gun.erl b/src/gun.erl
index ead6152..39b3b95 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -739,17 +739,17 @@ data(ServerPid, StreamRef, IsFin, Data) ->
 
 %% User pings.
 
--spec ping(pid()) -> stream_ref().
+-spec ping(pid()) -> reference().
 ping(ServerPid) ->
 	ping(ServerPid, #{}).
 
--spec ping(pid(), req_opts()) -> stream_ref().
+-spec ping(pid(), req_opts()) -> reference().
 ping(ServerPid, ReqOpts) ->
 	Tunnel = get_tunnel(ReqOpts),
-	StreamRef = make_stream_ref(Tunnel),
+	PingRef = make_ref(),
 	ReplyTo = maps:get(reply_to, ReqOpts, self()),
-	gen_statem:cast(ServerPid, {ping, ReplyTo, StreamRef}),
-	StreamRef.
+	gen_statem:cast(ServerPid, {ping, ReplyTo, Tunnel, PingRef}),
+	PingRef.
 
 %% Tunneling.
 
@@ -1397,9 +1397,13 @@ connected_ws_only(Type, Event, State) ->
 %%
 %% @todo It might be better, internally, to pass around a URIMap
 %% containing the target URI, instead of separate Host/Port/PathWithQs.
-connected(cast, {ping, ReplyTo, StreamRef},
+connected(cast, {ping, ReplyTo, Tunnel0, PingRef},
 		State=#state{protocol=Protocol, protocol_state=ProtoState}) ->
-	Commands = Protocol:ping(ProtoState, dereference_stream_ref(StreamRef, State), ReplyTo),
+	Tunnel = case dereference_stream_ref(Tunnel0, State) of
+		[] -> undefined;
+		Tunnel1 -> Tunnel1
+	end,
+	Commands = Protocol:ping(ProtoState, Tunnel, ReplyTo, PingRef),
 	commands(Commands, State);
 connected(cast, {headers, ReplyTo, StreamRef, Method, Path, Headers, InitialFlow},
 		State=#state{origin_host=Host, origin_port=Port,
diff --git a/src/gun_http.erl b/src/gun_http.erl
index d77d9b6..005f04b 100644
--- a/src/gun_http.erl
+++ b/src/gun_http.erl
@@ -26,7 +26,7 @@
 -export([closing/4]).
 -export([close/4]).
 -export([keepalive/3]).
--export([ping/3]).
+-export([ping/4]).
 -export([headers/12]).
 -export([request/13]).
 -export([data/7]).
@@ -561,7 +561,7 @@ keepalive(#http_state{socket=Socket, transport=Transport, out=head}, _, EvHandle
 keepalive(_State, _, EvHandlerState) ->
 	{[], EvHandlerState}.
 
-ping(_State, _PingRef, _ReplyTo) ->
+ping(_State, undefined, _ReplyTo, _PingRef) ->
 	{error, unsupported_by_protocol}.
 
 headers(State, StreamRef, ReplyTo, _, _, _, _, _, _, CookieStore, _, EvHandlerState)
diff --git a/src/gun_http2.erl b/src/gun_http2.erl
index bc100f5..a1d717e 100644
--- a/src/gun_http2.erl
+++ b/src/gun_http2.erl
@@ -27,7 +27,7 @@
 -export([closing/4]).
 -export([close/4]).
 -export([keepalive/3]).
--export([ping/3]).
+-export([ping/4]).
 -export([headers/12]).
 -export([request/13]).
 -export([data/7]).
@@ -389,9 +389,8 @@ maybe_ack_or_notify(State=#http2_state{reply_to=ReplyTo, socket=Socket,
 		{ping_ack, Payload} ->
 			%% User ping.
 			case lists:keytake(Payload, #user_ping.payload, UserPings0) of
-				{value, #user_ping{ref=StreamRef, reply_to=PingReplyTo}, UserPings} ->
-					RealStreamRef = stream_ref(State, StreamRef),
-					PingReplyTo ! {gun_notify, self(), ping_ack, RealStreamRef},
+				{value, #user_ping{ref=PingRef, reply_to=PingReplyTo}, UserPings} ->
+					PingReplyTo ! {gun_notify, self(), ping_ack, PingRef},
 					{state, State#http2_state{user_pings=UserPings}};
 				false ->
 					%% Ignore unexpected ping ack. RFC 7540
@@ -971,15 +970,36 @@ keepalive(State=#http2_state{socket=Socket, transport=Transport, pings_unack=Pin
 			{Error, EvHandlerState}
 	end.
 
-ping(State=#http2_state{socket=Socket, transport=Transport, user_pings=UserPings}, StreamRef, ReplyTo) ->
+ping(State=#http2_state{socket=Socket, transport=Transport, user_pings=UserPings},
+		undefined, ReplyTo, PingRef) ->
 	%% Use non-zero 64-bit payload for user pings. 0 is reserved for keepalive.
 	Payload = erlang:unique_integer([monotonic, positive]),
 	case Transport:send(Socket, cow_http2:ping(Payload)) of
 		ok ->
-			UserPing = #user_ping{ref = StreamRef, reply_to = ReplyTo, payload = Payload},
-			{state, State#http2_state{user_pings = UserPings ++ [UserPing]}};
-		Error={error, _} ->
-			{ok, Error}
+			UserPing = #user_ping{ref = PingRef, reply_to = ReplyTo, payload = Payload},
+			{state, State#http2_state{user_pings = [UserPing|UserPings]}};
+		Error = {error, _} ->
+			Error
+	end;
+%% Tunneled ping.
+ping(State, TunnelRef=[StreamRef|_], ReplyTo, PingRef) ->
+	case get_stream_by_ref(State, StreamRef) of
+		%% @todo We should send an error to the user if the stream isn't ready.
+		Stream=#stream{tunnel=Tunnel=#tunnel{protocol=Proto, protocol_state=ProtoState0}} ->
+			case Proto:ping(ProtoState0, TunnelRef, ReplyTo, PingRef) of
+				{state, ProtoState} ->
+					{state, store_stream(State, Stream#stream{
+						tunnel=Tunnel#tunnel{protocol_state=ProtoState}})};
+				Error = {error, _} ->
+					Error
+			end;
+		#stream{tunnel=undefined} ->
+			gun:reply(ReplyTo, {gun_error, self(), stream_ref(State, StreamRef), {badstate,
+				"The stream is not a tunnel."}}),
+			{state, State};
+		error ->
+			error_stream_not_found(State, StreamRef, ReplyTo),
+			{state, State}
 	end.
 
 headers(State, StreamRef, ReplyTo, Method, Host, Port, Path,
diff --git a/src/gun_http3.erl b/src/gun_http3.erl
index 5b95d4a..99f14cc 100644
--- a/src/gun_http3.erl
+++ b/src/gun_http3.erl
@@ -28,7 +28,7 @@
 -export([closing/4]).
 -export([close/4]).
 -export([keepalive/3]).
--export([ping/3]).
+-export([ping/4]).
 -export([headers/12]).
 -export([request/13]).
 -export([data/7]).
@@ -487,9 +487,9 @@ close(_Reason, _State, _, EvHandlerState) ->
 keepalive(_State, _, _EvHandlerState) ->
 	error(todo).
 
--spec ping(_, _, _) -> {error, not_implemented}.
+-spec ping(_, _, _, _) -> {error, not_implemented}.
 
-ping(_State, _PingRef, _ReplyTo) ->
+ping(_State, _Tunnel, _ReplyTo, _PingRef) ->
 	{error, not_implemented}.
 
 headers(State0=#http3_state{conn=Conn, transport=Transport,
diff --git a/src/gun_raw.erl b/src/gun_raw.erl
index 3c506b6..a3f1ef5 100644
--- a/src/gun_raw.erl
+++ b/src/gun_raw.erl
@@ -23,7 +23,6 @@
 -export([update_flow/4]).
 -export([closing/4]).
 -export([close/4]).
--export([ping/3]).
 -export([data/7]).
 -export([down/1]).
 
@@ -85,9 +84,6 @@ closing(_, _, _, EvHandlerState) ->
 close(_, _, _, EvHandlerState) ->
 	EvHandlerState.
 
-ping(_State, _PingRef, _ReplyTo) ->
-	{error, unsupported_by_protocol}.
-
 %% @todo Initiate closing on IsFin=fin.
 data(#raw_state{ref=StreamRef, socket=Socket, transport=Transport}, StreamRef,
 		_ReplyTo, _IsFin, Data, _EvHandler, EvHandlerState) ->
diff --git a/src/gun_socks.erl b/src/gun_socks.erl
index 7a2bdd7..d29f906 100644
--- a/src/gun_socks.erl
+++ b/src/gun_socks.erl
@@ -23,7 +23,6 @@
 -export([handle/5]).
 -export([closing/4]).
 -export([close/4]).
--export([ping/3]).
 %% @todo down
 
 -record(socks_state, {
@@ -206,6 +205,3 @@ closing(_, _, _, EvHandlerState) ->
 
 close(_, _, _, EvHandlerState) ->
 	EvHandlerState.
-
-ping(_State, _PingRef, _ReplyTo) ->
-	{error, unsupported_by_protocol}.
diff --git a/src/gun_tunnel.erl b/src/gun_tunnel.erl
index a6d51e3..d89b665 100644
--- a/src/gun_tunnel.erl
+++ b/src/gun_tunnel.erl
@@ -24,6 +24,7 @@
 -export([closing/4]).
 -export([close/4]).
 -export([keepalive/3]).
+-export([ping/4]).
 -export([headers/12]).
 -export([request/13]).
 -export([data/7]).
@@ -285,6 +286,19 @@ keepalive(_State, _EvHandler, EvHandlerState) ->
 	%% @todo Need to figure out how to handle keepalive for tunnels.
 	{[], EvHandlerState}.
 
+ping(State=#tunnel_state{protocol=Proto, protocol_state=ProtoState0},
+		TunnelRef0, ReplyTo, PingRef) ->
+	TunnelRef = case maybe_dereference(State, TunnelRef0) of
+		[] -> undefined;
+		TunnelRef1 -> TunnelRef1
+	end,
+	case Proto:ping(ProtoState0, TunnelRef, ReplyTo, PingRef) of
+		{state, ProtoState} ->
+			{state, State#tunnel_state{protocol_state=ProtoState}};
+		Error = {error, _} ->
+			Error
+	end.
+
 %% We pass the headers forward and optionally dereference StreamRef.
 headers(State=#tunnel_state{protocol=Proto, protocol_state=ProtoState0},
 		StreamRef0, ReplyTo, Method, Host, Port, Path, Headers,
diff --git a/src/gun_ws.erl b/src/gun_ws.erl
index 09e29aa..07df766 100644
--- a/src/gun_ws.erl
+++ b/src/gun_ws.erl
@@ -28,7 +28,7 @@
 -export([closing/4]).
 -export([close/4]).
 -export([keepalive/3]).
--export([ping/3]).
+-export([ping/4]).
 -export([ws_send/6]).
 -export([down/1]).
 
@@ -308,7 +308,7 @@ close(_, _, _, EvHandlerState) ->
 keepalive(State=#ws_state{reply_to=ReplyTo}, EvHandler, EvHandlerState0) ->
 	send(ping, State, ReplyTo, EvHandler, EvHandlerState0).
 
-ping(_State, _PingRef, _ReplyTo) ->
+ping(_State, undefined, _ReplyTo, _PingRef) ->
 	{error, not_implemented}.
 
 %% Send one frame.