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.

5 files changed, 211 insertions, 161 deletions

diff --git a/src/gun.erl b/src/gun.erl
index 2af6b94..73779c2 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -297,9 +297,9 @@
 	messages :: {atom(), atom(), atom()},
 	protocol :: module(),
 	protocol_state :: any(),
+	cookie_store :: undefined | {module(), any()},
 	event_handler :: module(),
-	event_handler_state :: any(),
-	cookie_store :: undefined | {module(), any()}
+	event_handler_state :: any()
 }).
 
 %% Connection.
@@ -1251,26 +1251,26 @@ connected(internal, {connected, Socket, NewProtocol},
 %%
 %% @todo It might be better, internally, to pass around a URIMap
 %% containing the target URI, instead of separate Host/Port/PathWithQs.
-connected(cast, {headers, ReplyTo, StreamRef, Method, Path, Headers0, InitialFlow},
-		State0=#state{origin_host=Host, origin_port=Port,
-			protocol=Protocol, protocol_state=ProtoState,
+connected(cast, {headers, ReplyTo, StreamRef, Method, Path, Headers, InitialFlow},
+		State=#state{origin_host=Host, origin_port=Port,
+			protocol=Protocol, protocol_state=ProtoState, cookie_store=CookieStore0,
 			event_handler=EvHandler, event_handler_state=EvHandlerState0}) ->
-	{Headers, State} = add_cookie_header(Path, Headers0, State0),
-	{ProtoState2, EvHandlerState} = Protocol:headers(ProtoState,
+	{ProtoState2, CookieStore, EvHandlerState} = Protocol:headers(ProtoState,
 		dereference_stream_ref(StreamRef, State), ReplyTo,
 		Method, Host, Port, Path, Headers,
-		InitialFlow, EvHandler, EvHandlerState0),
-	{keep_state, State#state{protocol_state=ProtoState2, event_handler_state=EvHandlerState}};
-connected(cast, {request, ReplyTo, StreamRef, Method, Path, Headers0, Body, InitialFlow},
-		State0=#state{origin_host=Host, origin_port=Port,
-			protocol=Protocol, protocol_state=ProtoState,
+		InitialFlow, CookieStore0, EvHandler, EvHandlerState0),
+	{keep_state, State#state{protocol_state=ProtoState2, cookie_store=CookieStore,
+		event_handler_state=EvHandlerState}};
+connected(cast, {request, ReplyTo, StreamRef, Method, Path, Headers, Body, InitialFlow},
+		State=#state{origin_host=Host, origin_port=Port,
+			protocol=Protocol, protocol_state=ProtoState, cookie_store=CookieStore0,
 			event_handler=EvHandler, event_handler_state=EvHandlerState0}) ->
-	{Headers, State} = add_cookie_header(Path, Headers0, State0),
-	{ProtoState2, EvHandlerState} = Protocol:request(ProtoState,
+	{ProtoState2, CookieStore, EvHandlerState} = Protocol:request(ProtoState,
 		dereference_stream_ref(StreamRef, State), ReplyTo,
 		Method, Host, Port, Path, Headers, Body,
-		InitialFlow, EvHandler, EvHandlerState0),
-	{keep_state, State#state{protocol_state=ProtoState2, event_handler_state=EvHandlerState}};
+		InitialFlow, CookieStore0, EvHandler, EvHandlerState0),
+	{keep_state, State#state{protocol_state=ProtoState2, cookie_store=CookieStore,
+		event_handler_state=EvHandlerState}};
 connected(cast, {connect, ReplyTo, StreamRef, Destination, Headers, InitialFlow},
 		State=#state{origin_host=Host, origin_port=Port,
 			protocol=Protocol, protocol_state=ProtoState,
@@ -1279,16 +1279,17 @@ connected(cast, {connect, ReplyTo, StreamRef, Destination, Headers, InitialFlow}
 		dereference_stream_ref(StreamRef, State), ReplyTo,
 		Destination, #{host => Host, port => Port},
 		Headers, InitialFlow, EvHandler, EvHandlerState0),
-	{keep_state, State#state{protocol_state=ProtoState2, event_handler_state=EvHandlerState}};
+	{keep_state, State#state{protocol_state=ProtoState2,
+		event_handler_state=EvHandlerState}};
 %% Public Websocket interface.
 %% @todo Maybe make an interface in the protocol module instead of checking on protocol name.
 %% An interface would also make sure that HTTP/1.0 can't upgrade.
 connected(cast, {ws_upgrade, ReplyTo, StreamRef, Path, Headers}, State=#state{opts=Opts}) ->
 	WsOpts = maps:get(ws_opts, Opts, #{}),
 	connected(cast, {ws_upgrade, ReplyTo, StreamRef, Path, Headers, WsOpts}, State);
-connected(cast, {ws_upgrade, ReplyTo, StreamRef, Path, Headers0, WsOpts},
-		State0=#state{origin_host=Host, origin_port=Port,
-			protocol=Protocol, protocol_state=ProtoState,
+connected(cast, {ws_upgrade, ReplyTo, StreamRef, Path, Headers, WsOpts},
+		State=#state{origin_host=Host, origin_port=Port,
+			protocol=Protocol, protocol_state=ProtoState, cookie_store=CookieStore0,
 			event_handler=EvHandler, event_handler_state=EvHandlerState0}) ->
 	EvHandlerState1 = EvHandler:ws_upgrade(#{
 		stream_ref => StreamRef,
@@ -1296,11 +1297,10 @@ connected(cast, {ws_upgrade, ReplyTo, StreamRef, Path, Headers0, WsOpts},
 		opts => WsOpts
 	}, EvHandlerState0),
 	%% @todo Can fail if HTTP/1.0.
-	{Headers, State} = add_cookie_header(Path, Headers0, State0),
-	{ProtoState2, EvHandlerState} = Protocol:ws_upgrade(ProtoState,
+	{ProtoState2, CookieStore, EvHandlerState} = Protocol:ws_upgrade(ProtoState,
 		dereference_stream_ref(StreamRef, State), ReplyTo,
-		Host, Port, Path, Headers, WsOpts, EvHandler, EvHandlerState1),
-	{keep_state, State#state{protocol_state=ProtoState2,
+		Host, Port, Path, Headers, WsOpts, CookieStore0, EvHandler, EvHandlerState1),
+	{keep_state, State#state{protocol_state=ProtoState2, cookie_store=CookieStore,
 		event_handler_state=EvHandlerState}};
 %% @todo Maybe better standardize the protocol callbacks argument orders.
 connected(cast, {ws_send, ReplyTo, StreamRef, Frames}, State=#state{
@@ -1319,35 +1319,6 @@ connected(cast, {ws_send, ReplyTo, _}, _) ->
 connected(Type, Event, State) ->
 	handle_common_connected(Type, Event, ?FUNCTION_NAME, State).
 
-add_cookie_header(_, Headers, State=#state{cookie_store=undefined}) ->
-	{Headers, State};
-add_cookie_header(PathWithQs, Headers0, State=#state{
-		origin_host=OriginHost, transport=Transport, cookie_store=Store0}) ->
-	Scheme = case Transport of
-		gun_tls -> <<"https">>;
-		gun_tls_proxy -> <<"https">>;
-		gun_tcp -> <<"http">>
-	end,
-	#{path := Path} = uri_string:parse(PathWithQs),
-	URIMap = uri_string:normalize(#{
-		scheme => Scheme,
-		host => case lists:keyfind(<<"host">>, 1, Headers0) of
-			false -> iolist_to_binary(OriginHost); %% @todo Probably not enough for atoms and such.
-			{_, HeaderHost} -> iolist_to_binary(HeaderHost)
-		end,
-		path => iolist_to_binary(Path)
-	}, [return_map]),
-	{ok, Cookies0, Store} = gun_cookies:query(Store0, URIMap),
-	Headers = case Cookies0 of
-		[] ->
-			Headers0;
-		_ ->
-			Cookies = [{Name, Value} || #{name := Name, value := Value} <- Cookies0],
-			%% We put cookies at the end of the headers list as it's the least important header.
-			Headers0 ++ [{<<"cookie">>, cow_cookie:cookie(Cookies)}]
-	end,
-	{Headers, State#state{cookie_store=Store}}.
-
 %% When the origin is using raw we do not dereference the stream_ref
 %% because it expects the full stream_ref to function (there's no
 %% other stream involved for this connection).
@@ -1655,6 +1626,8 @@ commands([{set_cookie, _, _, Status, _}|Tail], State=#state{opts=#{cookie_ignore
 %% @todo Make sure this works for proxied requests too.
 commands([{set_cookie, Authority, PathWithQs, _, Headers}|Tail], State=#state{
 			transport=Transport, cookie_store=Store0}) ->
+	%% @todo This is wrong. Also we should probably not do a command for this.
+	%% We should instead give the CookieStore to all callbacks.
 	Scheme = case Transport of
 		gun_tls -> <<"https">>;
 		gun_tls_proxy -> <<"https">>;
diff --git a/src/gun_cookies.erl b/src/gun_cookies.erl
index ce33b49..965c3a5 100644
--- a/src/gun_cookies.erl
+++ b/src/gun_cookies.erl
@@ -14,6 +14,7 @@
 
 -module(gun_cookies).
 
+-export([add_cookie_header/5]).
 -export([domain_match/2]).
 -export([gc/1]).
 -export([path_match/2]).
@@ -79,6 +80,31 @@
 	-> {ok, State}
 	when State::store_state().
 
+-spec add_cookie_header(binary(), iodata(), iodata(), Headers, Store)
+	-> {Headers, Store} when Headers :: [{binary(), iodata()}], Store :: undefined | store().
+add_cookie_header(_, _, _, Headers, Store=undefined) ->
+	{Headers, Store};
+add_cookie_header(Scheme, Authority, PathWithQs, Headers0, Store0) ->
+	#{
+		host := Host,
+		path := Path
+	} = uri_string:parse([Scheme, <<"://">>, Authority, PathWithQs]),
+	URIMap = uri_string:normalize(#{
+		scheme => Scheme,
+		host => iolist_to_binary(Host),
+		path => iolist_to_binary(Path)
+	}, [return_map]),
+	{ok, Cookies0, Store} = query(Store0, URIMap),
+	Headers = case Cookies0 of
+		[] ->
+			Headers0;
+		_ ->
+			Cookies = [{Name, Value} || #{name := Name, value := Value} <- Cookies0],
+			%% We put cookies at the end of the headers list as it's the least important header.
+			Headers0 ++ [{<<"cookie">>, cow_cookie:cookie(Cookies)}]
+	end,
+	{Headers, Store}.
+
 -spec domain_match(binary(), binary()) -> boolean().
 domain_match(String, String) ->
 	true;
diff --git a/src/gun_http.erl b/src/gun_http.erl
index d877068..62d9490 100644
--- a/src/gun_http.erl
+++ b/src/gun_http.erl
@@ -26,14 +26,14 @@
 -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]).
 -export([stream_info/2]).
 -export([down/1]).
--export([ws_upgrade/10]).
+-export([ws_upgrade/11]).
 
 %% Functions shared with gun_http2.
 -export([host_header/3]).
@@ -541,42 +541,44 @@ keepalive(State=#http_state{socket=Socket, transport=Transport, out=head}, _, Ev
 keepalive(State, _, EvHandlerState) ->
 	{State, EvHandlerState}.
 
-headers(State, StreamRef, ReplyTo, _, _, _, _, _, _, _, EvHandlerState)
+headers(State, StreamRef, ReplyTo, _, _, _, _, _, _, CookieStore, _, EvHandlerState)
 		when is_list(StreamRef) ->
 	ReplyTo ! {gun_error, self(), stream_ref(State, StreamRef),
 		{badstate, "The stream is not a tunnel."}},
-	{State, EvHandlerState};
+	{State, CookieStore, EvHandlerState};
 headers(State=#http_state{opts=Opts, out=head},
 		StreamRef, ReplyTo, Method, Host, Port, Path, Headers,
-		InitialFlow0, EvHandler, EvHandlerState0) ->
-	{Authority, Conn, Out, EvHandlerState} = send_request(State, StreamRef, ReplyTo,
-		Method, Host, Port, Path, Headers, undefined,
-		EvHandler, EvHandlerState0, ?FUNCTION_NAME),
+		InitialFlow0, CookieStore0, EvHandler, EvHandlerState0) ->
+	{Authority, Conn, Out, CookieStore, EvHandlerState} = send_request(State,
+		StreamRef, ReplyTo, Method, Host, Port, Path, Headers, undefined,
+		CookieStore0, EvHandler, EvHandlerState0, ?FUNCTION_NAME),
 	InitialFlow = initial_flow(InitialFlow0, Opts),
 	{new_stream(State#http_state{connection=Conn, out=Out}, StreamRef, ReplyTo,
-		Method, Authority, Path, InitialFlow), EvHandlerState}.
+		Method, Authority, Path, InitialFlow),
+		CookieStore, EvHandlerState}.
 
-request(State, StreamRef, ReplyTo, _, _, _, _, _, _, _, _, EvHandlerState)
+request(State, StreamRef, ReplyTo, _, _, _, _, _, _, _, CookieStore, _, EvHandlerState)
 		when is_list(StreamRef) ->
 	ReplyTo ! {gun_error, self(), stream_ref(State, StreamRef),
 		{badstate, "The stream is not a tunnel."}},
-	{State, EvHandlerState};
+	{State, CookieStore, EvHandlerState};
 request(State=#http_state{opts=Opts, out=head}, StreamRef, ReplyTo,
 		Method, Host, Port, Path, Headers, Body,
-		InitialFlow0, EvHandler, EvHandlerState0) ->
-	{Authority, Conn, Out, EvHandlerState} = send_request(State, StreamRef, ReplyTo,
-		Method, Host, Port, Path, Headers, Body,
-		EvHandler, EvHandlerState0, ?FUNCTION_NAME),
+		InitialFlow0, CookieStore0, EvHandler, EvHandlerState0) ->
+	{Authority, Conn, Out, CookieStore, EvHandlerState} = send_request(State,
+		StreamRef, ReplyTo, Method, Host, Port, Path, Headers, Body,
+		CookieStore0, EvHandler, EvHandlerState0, ?FUNCTION_NAME),
 	InitialFlow = initial_flow(InitialFlow0, Opts),
 	{new_stream(State#http_state{connection=Conn, out=Out}, StreamRef, ReplyTo,
-		Method, Authority, Path, InitialFlow), EvHandlerState}.
+		Method, Authority, Path, InitialFlow),
+		CookieStore, EvHandlerState}.
 
 initial_flow(infinity, #{flow := InitialFlow}) -> InitialFlow;
 initial_flow(InitialFlow, _) -> InitialFlow.
 
 send_request(State=#http_state{socket=Socket, transport=Transport, version=Version},
 		StreamRef, ReplyTo, Method, Host, Port, Path, Headers0, Body,
-		EvHandler, EvHandlerState0, Function) ->
+		CookieStore0, EvHandler, EvHandlerState0, Function) ->
 	Headers1 = lists:keydelete(<<"transfer-encoding">>, 1, Headers0),
 	Headers2 = case Body of
 		undefined -> Headers1;
@@ -596,12 +598,14 @@ send_request(State=#http_state{socket=Socket, transport=Transport, version=Versi
 		{_, Authority1} -> {Authority1, Headers2}
 	end,
 	Headers4 = transform_header_names(State, Headers3),
-	Headers = case {Body, Out} of
+	Headers5 = case {Body, Out} of
 		{undefined, body_chunked} when Version =:= 'HTTP/1.0' -> Headers4;
 		{undefined, body_chunked} -> [{<<"transfer-encoding">>, <<"chunked">>}|Headers4];
 		{undefined, _} -> Headers4;
 		_ -> [{<<"content-length">>, integer_to_binary(iolist_size(Body))}|Headers4]
 	end,
+	{Headers, CookieStore} = gun_cookies:add_cookie_header(
+		scheme(State), Authority, Path, Headers5, CookieStore0),
 	RealStreamRef = stream_ref(State, StreamRef),
 	RequestEvent = #{
 		stream_ref => RealStreamRef,
@@ -627,7 +631,7 @@ send_request(State=#http_state{socket=Socket, transport=Transport, version=Versi
 		_ ->
 			EvHandlerState2
 	end,
-	{Authority, Conn, Out, EvHandlerState}.
+	{Authority, Conn, Out, CookieStore, EvHandlerState}.
 
 host_header(Transport, Host0, Port) ->
 	Host = case Host0 of
@@ -649,6 +653,13 @@ transform_header_names(#http_state{opts=Opts}, Headers) ->
 		Fun -> lists:keymap(Fun, 1, Headers)
 	end.
 
+scheme(#http_state{transport=Transport}) ->
+	case Transport of
+		gun_tls -> <<"https">>;
+		gun_tls_proxy -> <<"https">>;
+		_ -> <<"http">>
+	end.
+
 %% We are expecting a new stream.
 data(State=#http_state{out=head}, StreamRef, ReplyTo, _, _, _, EvHandlerState) ->
 	{error_stream_closed(State, StreamRef, ReplyTo), EvHandlerState};
@@ -897,18 +908,18 @@ end_stream(State=#http_state{streams=[_|Tail]}) ->
 
 %% Websocket upgrade.
 
-ws_upgrade(State, StreamRef, ReplyTo, _, _, _, _, _, _, EvHandlerState)
+ws_upgrade(State, StreamRef, ReplyTo, _, _, _, _, _, CookieStore, _, EvHandlerState)
 		when is_list(StreamRef) ->
 	ReplyTo ! {gun_error, self(), stream_ref(State, StreamRef),
 		{badstate, "The stream is not a tunnel."}},
-	{State, EvHandlerState};
+	{State, CookieStore, EvHandlerState};
 ws_upgrade(State=#http_state{version='HTTP/1.0'},
-		StreamRef, ReplyTo, _, _, _, _, _, _, EvHandlerState) ->
+		StreamRef, ReplyTo, _, _, _, _, _, CookieStore, _, EvHandlerState) ->
 	ReplyTo ! {gun_error, self(), stream_ref(State, StreamRef), {badstate,
 		"Websocket cannot be used over an HTTP/1.0 connection."}},
-	{[], EvHandlerState};
+	{State, CookieStore, EvHandlerState};
 ws_upgrade(State=#http_state{out=head}, StreamRef, ReplyTo,
-		Host, Port, Path, Headers0, WsOpts, EvHandler, EvHandlerState0) ->
+		Host, Port, Path, Headers0, WsOpts, CookieStore0, EvHandler, EvHandlerState0) ->
 	{Headers1, GunExtensions} = case maps:get(compress, WsOpts, false) of
 		true -> {[{<<"sec-websocket-extensions">>,
 				<<"permessage-deflate; client_max_window_bits; server_max_window_bits=15">>}
@@ -930,13 +941,14 @@ ws_upgrade(State=#http_state{out=head}, StreamRef, ReplyTo,
 		{<<"sec-websocket-key">>, Key}
 		|Headers2
 	],
-	{Authority, Conn, Out, EvHandlerState} = send_request(State, StreamRef, ReplyTo,
-		<<"GET">>, Host, Port, Path, Headers, undefined,
-		EvHandler, EvHandlerState0, ?FUNCTION_NAME),
+	{Authority, Conn, Out, CookieStore, EvHandlerState} = send_request(State,
+		StreamRef, ReplyTo, <<"GET">>, Host, Port, Path, Headers, undefined,
+		CookieStore0, EvHandler, EvHandlerState0, ?FUNCTION_NAME),
 	InitialFlow = maps:get(flow, WsOpts, infinity),
 	{new_stream(State#http_state{connection=Conn, out=Out},
 		#websocket{ref=StreamRef, reply_to=ReplyTo, key=Key, extensions=GunExtensions, opts=WsOpts},
-		ReplyTo, <<"GET">>, Authority, Path, InitialFlow), EvHandlerState}.
+		ReplyTo, <<"GET">>, Authority, Path, InitialFlow),
+		CookieStore, EvHandlerState}.
 
 ws_handshake(Buffer, State, Ws=#websocket{key=Key}, Headers) ->
 	%% @todo check upgrade, connection
diff --git a/src/gun_http2.erl b/src/gun_http2.erl
index c5521ce..e8eb3aa 100644
--- a/src/gun_http2.erl
+++ b/src/gun_http2.erl
@@ -27,15 +27,15 @@
 -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]).
 -export([timeout/3]).
 -export([stream_info/2]).
 -export([down/1]).
--export([ws_upgrade/10]).
+-export([ws_upgrade/11]).
 -export([ws_send/6]).
 
 -record(tunnel, {
@@ -622,7 +622,7 @@ handle_continue(ContinueStreamRef, Msg, State0, EvHandler, EvHandlerState0) ->
 			{Commands, EvHandlerState1} = Proto:handle_continue(ContinueStreamRef,
 				Msg, ProtoState0, EvHandler, EvHandlerState0),
 			{State, EvHandlerState} = tunnel_commands(Commands, Stream, State0, EvHandler, EvHandlerState1),
-			{{state, State}, EvHandlerState}
+			{handle_ret({state, State}, State), EvHandlerState}
 		%% The stream may have ended while TLS was being decoded. @todo What should we do?
 %		error ->
 %			{error_stream_not_found(State, StreamRef, ReplyTo), EvHandlerState0}
@@ -750,11 +750,12 @@ keepalive(State=#http2_state{socket=Socket, transport=Transport}, _, EvHandlerSt
 
 headers(State=#http2_state{socket=Socket, transport=Transport, opts=Opts,
 		http2_machine=HTTP2Machine0}, StreamRef, ReplyTo, Method, Host, Port,
-		Path, Headers0, InitialFlow0, EvHandler, EvHandlerState0)
+		Path, Headers0, InitialFlow0, CookieStore0, EvHandler, EvHandlerState0)
 		when is_reference(StreamRef) ->
 	{ok, StreamID, HTTP2Machine1} = cow_http2_machine:init_stream(
 		iolist_to_binary(Method), HTTP2Machine0),
-	{ok, PseudoHeaders, Headers} = prepare_headers(State, Method, Host, Port, Path, Headers0),
+	{ok, PseudoHeaders, Headers, CookieStore} = prepare_headers(
+		State, Method, Host, Port, Path, Headers0, CookieStore0),
 	Authority = maps:get(authority, PseudoHeaders),
 	RequestEvent = #{
 		stream_ref => stream_ref(State, StreamRef),
@@ -773,36 +774,39 @@ headers(State=#http2_state{socket=Socket, transport=Transport, opts=Opts,
 	InitialFlow = initial_flow(InitialFlow0, Opts),
 	Stream = #stream{id=StreamID, ref=StreamRef, reply_to=ReplyTo, flow=InitialFlow,
 		authority=Authority, path=Path},
-	{create_stream(State#http2_state{http2_machine=HTTP2Machine}, Stream), EvHandlerState};
+	{create_stream(State#http2_state{http2_machine=HTTP2Machine}, Stream),
+		CookieStore, EvHandlerState};
 %% Tunneled request.
 headers(State, RealStreamRef=[StreamRef|_], ReplyTo, Method, _Host, _Port,
-		Path, Headers, InitialFlow, EvHandler, EvHandlerState0) ->
+		Path, Headers, InitialFlow, CookieStore0, EvHandler, EvHandlerState0) ->
 	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, info=#{
 				origin_host := OriginHost, origin_port := OriginPort}}} ->
-			{ProtoState, EvHandlerState} = Proto:headers(ProtoState0, RealStreamRef,
+			{ProtoState, CookieStore, EvHandlerState} = Proto:headers(ProtoState0, RealStreamRef,
 				ReplyTo, Method, OriginHost, OriginPort, Path, Headers,
-				InitialFlow, EvHandler, EvHandlerState0),
+				InitialFlow, CookieStore0, EvHandler, EvHandlerState0),
 			{store_stream(State, Stream#stream{tunnel=Tunnel#tunnel{protocol_state=ProtoState}}),
-				EvHandlerState};
+				CookieStore, EvHandlerState};
 		#stream{tunnel=undefined} ->
 			ReplyTo ! {gun_error, self(), stream_ref(State, StreamRef), {badstate,
 				"The stream is not a tunnel."}},
-			{State, EvHandlerState0};
+			{State, CookieStore0, EvHandlerState0};
 		error ->
-			{error_stream_not_found(State, StreamRef, ReplyTo), EvHandlerState0}
+			{error_stream_not_found(State, StreamRef, ReplyTo),
+				CookieStore0, EvHandlerState0}
 	end.
 
 request(State0=#http2_state{socket=Socket, transport=Transport, opts=Opts,
 		http2_machine=HTTP2Machine0}, StreamRef, ReplyTo, Method, Host, Port,
-		Path, Headers0, Body, InitialFlow0, EvHandler, EvHandlerState0)
+		Path, Headers0, Body, InitialFlow0, CookieStore0, EvHandler, EvHandlerState0)
 		when is_reference(StreamRef) ->
 	Headers1 = lists:keystore(<<"content-length">>, 1, Headers0,
 		{<<"content-length">>, integer_to_binary(iolist_size(Body))}),
 	{ok, StreamID, HTTP2Machine1} = cow_http2_machine:init_stream(
 		iolist_to_binary(Method), HTTP2Machine0),
-	{ok, PseudoHeaders, Headers} = prepare_headers(State0, Method, Host, Port, Path, Headers1),
+	{ok, PseudoHeaders, Headers, CookieStore} = prepare_headers(
+		State0, Method, Host, Port, Path, Headers1, CookieStore0),
 	Authority = maps:get(authority, PseudoHeaders),
 	RealStreamRef = stream_ref(State0, StreamRef),
 	RequestEvent = #{
@@ -833,60 +837,66 @@ request(State0=#http2_state{socket=Socket, transport=Transport, opts=Opts,
 				stream_ref => RealStreamRef,
 				reply_to => ReplyTo
 			},
-			{State, EvHandler:request_end(RequestEndEvent, EvHandlerState)};
+			{State, CookieStore, EvHandler:request_end(RequestEndEvent, EvHandlerState)};
 		nofin ->
-			maybe_send_data(State, StreamID, fin, Body, EvHandler, EvHandlerState)
+			{StateRet, EvHandlerStateRet} = maybe_send_data(
+				State, StreamID, fin, Body, EvHandler, EvHandlerState),
+			{StateRet, CookieStore, EvHandlerStateRet}
 	end;
 %% Tunneled request.
 request(State, RealStreamRef=[StreamRef|_], ReplyTo, Method, _Host, _Port,
-		Path, Headers, Body, InitialFlow, EvHandler, EvHandlerState0) ->
+		Path, Headers, Body, InitialFlow, CookieStore0, EvHandler, EvHandlerState0) ->
 	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, info=#{
 				origin_host := OriginHost, origin_port := OriginPort}}} ->
-			{ProtoState, EvHandlerState} = Proto:request(ProtoState0, RealStreamRef,
+			{ProtoState, CookieStore, EvHandlerState} = Proto:request(ProtoState0, RealStreamRef,
 				ReplyTo, Method, OriginHost, OriginPort, Path, Headers, Body,
-				InitialFlow, EvHandler, EvHandlerState0),
+				InitialFlow, CookieStore0, EvHandler, EvHandlerState0),
 			{store_stream(State, Stream#stream{tunnel=Tunnel#tunnel{protocol_state=ProtoState}}),
-				EvHandlerState};
+				CookieStore, EvHandlerState};
 		#stream{tunnel=undefined} ->
 			ReplyTo ! {gun_error, self(), stream_ref(State, StreamRef), {badstate,
 				"The stream is not a tunnel."}},
-			{State, EvHandlerState0};
+			{State, CookieStore0, EvHandlerState0};
 		error ->
-			{error_stream_not_found(State, StreamRef, ReplyTo), EvHandlerState0}
+			{error_stream_not_found(State, StreamRef, ReplyTo),
+				CookieStore0, EvHandlerState0}
 	end.
 
 initial_flow(infinity, #{flow := InitialFlow}) -> InitialFlow;
 initial_flow(InitialFlow, _) -> InitialFlow.
 
-prepare_headers(#http2_state{transport=Transport}, Method, Host0, Port, Path, Headers0) ->
+prepare_headers(#http2_state{transport=Transport}, Method, Host0, Port, Path, Headers0, CookieStore0) ->
+	Scheme = case Transport of
+		gun_tls -> <<"https">>;
+		gun_tls_proxy -> <<"https">>;
+		gun_tcp -> <<"http">>;
+		gun_tcp_proxy -> <<"http">>;
+		gun_tls_proxy_http2_connect -> <<"http">>
+	end,
 	Authority = case lists:keyfind(<<"host">>, 1, Headers0) of
 		{_, Host} -> Host;
 		_ -> gun_http:host_header(Transport, Host0, Port)
 	end,
 	%% @todo We also must remove any header found in the connection header.
 	%% @todo Much of this is duplicated in cow_http2_machine; sort things out.
-	Headers =
+	Headers1 =
 		lists:keydelete(<<"host">>, 1,
 		lists:keydelete(<<"connection">>, 1,
 		lists:keydelete(<<"keep-alive">>, 1,
 		lists:keydelete(<<"proxy-connection">>, 1,
 		lists:keydelete(<<"transfer-encoding">>, 1,
 		lists:keydelete(<<"upgrade">>, 1, Headers0)))))),
+	{Headers, CookieStore} = gun_cookies:add_cookie_header(
+		Scheme, Authority, Path, Headers1, CookieStore0),
 	PseudoHeaders = #{
 		method => Method,
-		scheme => case Transport of
-			gun_tls -> <<"https">>;
-			gun_tls_proxy -> <<"https">>;
-			gun_tcp -> <<"http">>;
-			gun_tcp_proxy -> <<"http">>;
-			gun_tls_proxy_http2_connect -> <<"http">>
-		end,
+		scheme => Scheme,
 		authority => Authority,
 		path => Path
 	},
-	{ok, PseudoHeaders, Headers}.
+	{ok, PseudoHeaders, Headers, CookieStore}.
 
 %% @todo Make all calls go through this clause.
 data(State=#http2_state{http2_machine=HTTP2Machine}, StreamRef, ReplyTo, IsFin, Data,
@@ -1181,13 +1191,16 @@ down(#http2_state{stream_refs=Refs}) ->
 
 %% Websocket upgrades are currently only accepted when tunneled.
 ws_upgrade(State, RealStreamRef=[StreamRef|_], ReplyTo,
-		Host, Port, Path, Headers, WsOpts, EvHandler, EvHandlerState0) ->
+		Host, Port, Path, Headers, WsOpts, CookieStore0, EvHandler, EvHandlerState0) ->
 	case get_stream_by_ref(State, StreamRef) of
 		Stream=#stream{tunnel=Tunnel=#tunnel{protocol=Proto, protocol_state=ProtoState0}} ->
-			{ProtoState, EvHandlerState} = Proto:ws_upgrade(ProtoState0, RealStreamRef, ReplyTo,
-				Host, Port, Path, Headers, WsOpts, EvHandler, EvHandlerState0),
+			{ProtoState, CookieStore, EvHandlerState} = Proto:ws_upgrade(
+				ProtoState0, RealStreamRef, ReplyTo,
+				Host, Port, Path, Headers, WsOpts,
+				CookieStore0, EvHandler, EvHandlerState0),
 			{store_stream(State, Stream#stream{
-				tunnel=Tunnel#tunnel{protocol_state=ProtoState}}), EvHandlerState}
+				tunnel=Tunnel#tunnel{protocol_state=ProtoState}}),
+				CookieStore, EvHandlerState}
 		%% @todo Error conditions?
 	end.
 
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.