Fix cookies for tunnels

There are still small issues left to fix. In particular the
set_cookie command should be replaced with doing the same
in the protocol itself so that the scheme is correct. So
CookieStore must be propagated to all callbacks.

1 files changed, 68 insertions, 42 deletions

diff --git a/src/gun_tunnel.erl b/src/gun_tunnel.erl
index cc58351..a1435f3 100644
--- a/src/gun_tunnel.erl
+++ b/src/gun_tunnel.erl
@@ -24,8 +24,8 @@
 -export([closing/4]).
 -export([close/4]).
 -export([keepalive/3]).
--export([headers/11]).
--export([request/12]).
+-export([headers/12]).
+-export([request/13]).
 -export([data/7]).
 -export([connect/9]).
 -export([cancel/5]).
@@ -33,7 +33,7 @@
 -export([stream_info/2]).
 -export([tunneled_name/2]).
 -export([down/1]).
--export([ws_upgrade/10]).
+-export([ws_upgrade/11]).
 -export([ws_send/6]).
 
 -record(tunnel_state, {
@@ -87,7 +87,10 @@
 	%% We keep the new information to provide it in TunnelInfo of
 	%% the new protocol when the switch occurs.
 	protocol_origin = undefined :: undefined
-		| {origin, binary(), inet:hostname() | inet:ip_address(), inet:port_number(), connect | socks5}
+		| {origin, binary(), inet:hostname() | inet:ip_address(), inet:port_number(), connect | socks5},
+
+	%% We must queue some commands given by the sub-protocol.
+	commands_queue = [] :: [{set_cookie, iodata(), iodata(), cow_http:status(), cow_http:headers()}]
 }).
 
 %% Socket is the "origin socket" and Transport the "origin transport".
@@ -158,7 +161,7 @@ handle(Data, State0=#tunnel_state{transport=gun_tcp_proxy,
 		EvHandler, EvHandlerState0) ->
 	{Commands, EvHandlerState1} = Proto:handle(Data, ProtoState0, EvHandler, EvHandlerState0),
 	{State, EvHandlerState} = commands(Commands, State0, EvHandler, EvHandlerState1),
-	{{state, State}, EvHandlerState};
+	{ret({state, State}, State), EvHandlerState};
 handle(Data, State=#tunnel_state{transport=gun_tls_proxy,
 		socket=ProxyPid, tls_origin_socket=OriginSocket},
 		_EvHandler, EvHandlerState) ->
@@ -168,7 +171,7 @@ handle(Data, State=#tunnel_state{transport=gun_tls_proxy,
 	%% message and forward it to the right stream via the handle_continue
 	%% callback.
 	ProxyPid ! {tls_proxy_http2_connect, OriginSocket, Data},
-	{{state, State}, EvHandlerState}.
+	{ret({state, State}, State), EvHandlerState}.
 
 %% This callback will only be called for TLS.
 %%
@@ -199,10 +202,11 @@ handle_continue(ContinueStreamRef, {gun_tls_proxy, ProxyPid, {ok, Negotiated},
 	{_, ProtoState} = Proto:init(ReplyTo, OriginSocket, gun_tcp_proxy,
 		ProtoOpts#{stream_ref => StreamRef, tunnel_transport => tls}),
 	ReplyTo ! {gun_tunnel_up, self(), StreamRef, Proto:name()},
-	{{state, State#tunnel_state{protocol=Proto, protocol_state=ProtoState}}, EvHandlerState};
+	{ret({state, State#tunnel_state{protocol=Proto, protocol_state=ProtoState}}, State),
+		EvHandlerState};
 handle_continue(ContinueStreamRef, {gun_tls_proxy, ProxyPid, {error, Reason},
 		{handle_continue, _, HandshakeEvent, _}},
-		#tunnel_state{socket=ProxyPid}, EvHandler, EvHandlerState0)
+		State=#tunnel_state{socket=ProxyPid}, EvHandler, EvHandlerState0)
 		when is_reference(ContinueStreamRef) ->
 	EvHandlerState = EvHandler:tls_handshake_end(HandshakeEvent#{
 		error => Reason
@@ -217,29 +221,29 @@ handle_continue(ContinueStreamRef, {gun_tls_proxy, ProxyPid, {error, Reason},
 %%   receives a TCP segment with the FIN bit set sends a DATA frame with
 %%   the END_STREAM flag set.  Note that the final TCP segment or DATA
 %%   frame could be empty.
-	{{error, Reason}, EvHandlerState};
+	{ret({error, Reason}, State), EvHandlerState};
 %% Send the data. This causes TLS to encrypt the data and send it to the inner layer.
 handle_continue(ContinueStreamRef, {data, _ReplyTo, _StreamRef, IsFin, Data},
-		#tunnel_state{}, _EvHandler, EvHandlerState)
+		State=#tunnel_state{}, _EvHandler, EvHandlerState)
 		when is_reference(ContinueStreamRef) ->
-	{{send, IsFin, Data}, EvHandlerState};
+	{ret({send, IsFin, Data}, State), EvHandlerState};
 handle_continue(ContinueStreamRef, {tls_proxy, ProxyPid, Data},
 		State0=#tunnel_state{socket=ProxyPid, protocol=Proto, protocol_state=ProtoState},
 		EvHandler, EvHandlerState0)
 		when is_reference(ContinueStreamRef) ->
 	{Commands, EvHandlerState1} = Proto:handle(Data, ProtoState, EvHandler, EvHandlerState0),
 	{State, EvHandlerState} = commands(Commands, State0, EvHandler, EvHandlerState1),
-	{{state, State}, EvHandlerState};
+	{ret({state, State}, State), EvHandlerState};
 handle_continue(ContinueStreamRef, {tls_proxy_closed, ProxyPid},
-		#tunnel_state{socket=ProxyPid}, _EvHandler, EvHandlerState0)
+		State=#tunnel_state{socket=ProxyPid}, _EvHandler, EvHandlerState0)
 		when is_reference(ContinueStreamRef) ->
 	%% @todo All sub-streams must produce a stream_error.
-	{{error, closed}, EvHandlerState0};
+	{ret({error, closed}, State), EvHandlerState0};
 handle_continue(ContinueStreamRef, {tls_proxy_error, ProxyPid, Reason},
-		#tunnel_state{socket=ProxyPid}, _EvHandler, EvHandlerState0)
+		State=#tunnel_state{socket=ProxyPid}, _EvHandler, EvHandlerState0)
 		when is_reference(ContinueStreamRef) ->
 	%% @todo All sub-streams must produce a stream_error.
-	{{error, Reason}, EvHandlerState0};
+	{ret({error, Reason}, State), EvHandlerState0};
 %% We always dereference the ContinueStreamRef because it includes a
 %% reference() for Socks layers too.
 %%
@@ -254,7 +258,7 @@ handle_continue([_StreamRef|ContinueStreamRef0], Msg,
 	{Commands, EvHandlerState1} = Proto:handle_continue(ContinueStreamRef,
 		Msg, ProtoState, EvHandler, EvHandlerState0),
 	{State, EvHandlerState} = commands(Commands, State0, EvHandler, EvHandlerState1),
-	{{state, State}, EvHandlerState}.
+	{ret({state, State}, State), EvHandlerState}.
 
 %% @todo This function will need EvHandler/EvHandlerState?
 update_flow(State0=#tunnel_state{protocol=Proto, protocol_state=ProtoState},
@@ -262,11 +266,11 @@ update_flow(State0=#tunnel_state{protocol=Proto, protocol_state=ProtoState},
 	StreamRef = maybe_dereference(State0, StreamRef0),
 	Commands = Proto:update_flow(ProtoState, ReplyTo, StreamRef, Inc),
 	{State, undefined} = commands(Commands, State0, undefined, undefined),
-	{state, State}.
+	ret({state, State}, State).
 
-closing(_Reason, _State, _EvHandler, EvHandlerState) ->
+closing(_Reason, State, _EvHandler, EvHandlerState) ->
 	%% @todo Graceful shutdown must be propagated to tunnels.
-	{[], EvHandlerState}.
+	{ret([], State), EvHandlerState}.
 
 close(_Reason, _State, _EvHandler, EvHandlerState) ->
 	%% @todo Closing must be propagated to tunnels.
@@ -279,23 +283,25 @@ keepalive(State, _EvHandler, EvHandlerState) ->
 %% 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,
-		InitialFlow, EvHandler, EvHandlerState0) ->
+		InitialFlow, CookieStore0, EvHandler, EvHandlerState0) ->
 	StreamRef = maybe_dereference(State, StreamRef0),
-	{ProtoState, EvHandlerState} = Proto:headers(ProtoState0, StreamRef,
+	{ProtoState, CookieStore, EvHandlerState} = Proto:headers(ProtoState0, StreamRef,
 		ReplyTo, Method, Host, Port, Path, Headers,
-		InitialFlow, EvHandler, EvHandlerState0),
-	{State#tunnel_state{protocol_state=ProtoState}, EvHandlerState}.
+		InitialFlow, CookieStore0, EvHandler, EvHandlerState0),
+	{State#tunnel_state{protocol_state=ProtoState},
+		CookieStore, EvHandlerState}.
 
 %% We pass the request forward and optionally dereference StreamRef.
 request(State=#tunnel_state{protocol=Proto, protocol_state=ProtoState0,
 		info=#{origin_host := OriginHost, origin_port := OriginPort}},
 		StreamRef0, ReplyTo, Method, _Host, _Port, Path, Headers, Body,
-		InitialFlow, EvHandler, EvHandlerState0) ->
+		InitialFlow, CookieStore0, EvHandler, EvHandlerState0) ->
 	StreamRef = maybe_dereference(State, StreamRef0),
-	{ProtoState, EvHandlerState} = Proto:request(ProtoState0, StreamRef,
+	{ProtoState, CookieStore, EvHandlerState} = Proto:request(ProtoState0, StreamRef,
 		ReplyTo, Method, OriginHost, OriginPort, Path, Headers, Body,
-		InitialFlow, EvHandler, EvHandlerState0),
-	{State#tunnel_state{protocol_state=ProtoState}, EvHandlerState}.
+		InitialFlow, CookieStore0, EvHandler, EvHandlerState0),
+	{State#tunnel_state{protocol_state=ProtoState},
+		CookieStore, EvHandlerState}.
 
 %% When the next tunnel is SOCKS we pass the data forward directly.
 %% This is needed because SOCKS has no StreamRef and the data cannot
@@ -340,14 +346,14 @@ cancel(State0=#tunnel_state{protocol=Proto, protocol_state=ProtoState},
 	StreamRef = maybe_dereference(State0, StreamRef0),
 	{Commands, EvHandlerState1} = Proto:cancel(ProtoState, StreamRef, ReplyTo, EvHandler, EvHandlerState0),
 	{State, EvHandlerState} = commands(Commands, State0, EvHandler, EvHandlerState1),
-	{{state, State}, EvHandlerState}.
+	{ret({state, State}, State), EvHandlerState}.
 
 timeout(State=#tunnel_state{protocol=Proto, protocol_state=ProtoState0}, Msg, TRef) ->
 	case Proto:timeout(ProtoState0, Msg, TRef) of
 		{state, ProtoState} ->
-			{state, State#tunnel_state{protocol_state=ProtoState}};
+			ret({state, State#tunnel_state{protocol_state=ProtoState}}, State);
 		Other ->
-			Other
+			ret(Other, State)
 	end.
 
 stream_info(#tunnel_state{transport=Transport0, stream_ref=TunnelStreamRef, reply_to=ReplyTo,
@@ -409,20 +415,23 @@ tunneled_name(#tunnel_state{tunnel_protocol=TunnelProto}, false) ->
 tunneled_name(#tunnel_state{protocol=Proto}, _) ->
 	Proto:name().
 
-down(_State) ->
+down(State) ->
 	%% @todo Tunnels must be included in the gun_down message.
-	[].
+	ret([], State).
 
 ws_upgrade(State=#tunnel_state{info=TunnelInfo, protocol=Proto, protocol_state=ProtoState0},
-		StreamRef0, ReplyTo, _, _, Path, Headers, WsOpts, EvHandler, EvHandlerState0) ->
+		StreamRef0, ReplyTo, _, _, Path, Headers, WsOpts,
+		CookieStore0, EvHandler, EvHandlerState0) ->
 	StreamRef = maybe_dereference(State, StreamRef0),
 	#{
 		origin_host := Host,
 		origin_port := Port
 	} = TunnelInfo,
-	{ProtoState, EvHandlerState} = Proto:ws_upgrade(ProtoState0, StreamRef, ReplyTo,
-		Host, Port, Path, Headers, WsOpts, EvHandler, EvHandlerState0),
-	{State#tunnel_state{protocol_state=ProtoState}, EvHandlerState}.
+	{ProtoState, CookieStore, EvHandlerState} = Proto:ws_upgrade(ProtoState0, StreamRef, ReplyTo,
+		Host, Port, Path, Headers, WsOpts,
+		CookieStore0, EvHandler, EvHandlerState0),
+	{State#tunnel_state{protocol_state=ProtoState},
+		CookieStore, EvHandlerState}.
 
 ws_send(Frames, State0=#tunnel_state{protocol=Proto, protocol_state=ProtoState},
 		StreamRef0, ReplyTo, EvHandler, EvHandlerState0) ->
@@ -430,7 +439,7 @@ ws_send(Frames, State0=#tunnel_state{protocol=Proto, protocol_state=ProtoState},
 	{Commands, EvHandlerState1} = Proto:ws_send(Frames,
 		ProtoState, StreamRef, ReplyTo, EvHandler, EvHandlerState0),
 	{State, EvHandlerState} = commands(Commands, State0, EvHandler, EvHandlerState1),
-	{{state, State}, EvHandlerState}.
+	{ret({state, State}, State), EvHandlerState}.
 
 %% Internal.
 
@@ -440,11 +449,14 @@ commands([], State, _, EvHandlerState) ->
 	{State, EvHandlerState};
 commands([{state, ProtoState}|Tail], State, EvHandler, EvHandlerState) ->
 	commands(Tail, State#tunnel_state{protocol_state=ProtoState}, EvHandler, EvHandlerState);
-%% @todo We must pass down the set_cookie commands. Have a commands_queue.
-commands([_SetCookie={set_cookie, _, _, _, _}|Tail], State=#tunnel_state{}, EvHandler, EvHandlerState) ->
-	commands(Tail, State, EvHandler, EvHandlerState);
+commands([SetCookie={set_cookie, _, _, _, _}|Tail],
+		State=#tunnel_state{commands_queue=Queue},
+		EvHandler, EvHandlerState) ->
+	commands(Tail, State#tunnel_state{commands_queue=[SetCookie|Queue]},
+		EvHandler, EvHandlerState);
 %% @todo What to do about IsFin?
-commands([{send, _IsFin, Data}|Tail], State=#tunnel_state{socket=Socket, transport=Transport},
+commands([{send, _IsFin, Data}|Tail],
+		State=#tunnel_state{socket=Socket, transport=Transport},
 		EvHandler, EvHandlerState) ->
 	Transport:send(Socket, Data),
 	commands(Tail, State, EvHandler, EvHandlerState);
@@ -592,3 +604,17 @@ outer_stream_ref(StreamRef) when is_list(StreamRef) ->
 	lists:last(StreamRef);
 outer_stream_ref(StreamRef) ->
 	StreamRef.
+
+%% This function is called before returning commands.
+ret(CommandOrCommands, #tunnel_state{commands_queue=[]}) ->
+	empty_commands_queue(CommandOrCommands);
+ret(Commands, #tunnel_state{commands_queue=Queue}) when is_list(Commands) ->
+	lists:reverse(Queue, empty_commands_queue(Commands));
+ret(Command, #tunnel_state{commands_queue=Queue}) ->
+	lists:reverse([empty_commands_queue(Command)|Queue]).
+
+empty_commands_queue([{state, State}|Tail]) -> [{state, State#tunnel_state{commands_queue=[]}}|Tail];
+empty_commands_queue([Command|Tail]) -> [Command|empty_commands_queue(Tail)];
+empty_commands_queue([]) -> [];
+empty_commands_queue({state, State}) -> {state, State#tunnel_state{commands_queue=[]}};
+empty_commands_queue(Command) -> Command.