Add set_resp_header/3 and set_resp_body/2 to cowboy_http_req

These functions allow to set response headers and body in advance,
before calling any of the reply functions.

Also add has_resp_header/2 and has_resp_body/1 to check if the given
response headers have already been set.

4 files changed, 134 insertions, 23 deletions

diff --git a/include/http.hrl b/include/http.hrl
index 7691966..d4fba18 100644
--- a/include/http.hrl
+++ b/include/http.hrl
@@ -66,5 +66,7 @@
 	buffer     = <<>>      :: binary(),
 
 	%% Response.
-	resp_state = waiting   :: locked | waiting | chunks | done
+	resp_state = waiting   :: locked | waiting | chunks | done,
+	resp_headers = []      :: http_headers(),
+	resp_body  = <<>>      :: binary()
 }).
diff --git a/src/cowboy_http_req.erl b/src/cowboy_http_req.erl
index d03f835..f457d4f 100644
--- a/src/cowboy_http_req.erl
+++ b/src/cowboy_http_req.erl
@@ -37,6 +37,8 @@
 ]). %% Request Body API.
 
 -export([
+	set_resp_header/3, set_resp_body/2,
+	has_resp_header/2, has_resp_body/1,
 	reply/2, reply/3, reply/4,
 	chunked_reply/2, chunked_reply/3, chunk/2,
 	upgrade_reply/3
@@ -360,24 +362,50 @@ body_qs(Req) ->
 
 %% Response API.
 
+%% @doc Add a header to the response.
+-spec set_resp_header(http_header(), binary(), #http_req{})
+	-> {ok, #http_req{}}.
+set_resp_header(Name, Value, Req=#http_req{resp_headers=RespHeaders}) ->
+	NameBin = header_to_binary(Name),
+	{ok, Req#http_req{resp_headers=[{NameBin, Value}|RespHeaders]}}.
+
+%% @doc Add a body to the response.
+%%
+%% The body set here is ignored if the response is later sent using
+%% anything other than reply/2 or reply/3.
+-spec set_resp_body(binary(), #http_req{}) -> {ok, #http_req{}}.
+set_resp_body(Body, Req) ->
+	{ok, Req#http_req{resp_body=Body}}.
+
+%% @doc Return whether the given header has been set for the response.
+-spec has_resp_header(http_header(), #http_req{}) -> boolean().
+has_resp_header(Name, #http_req{resp_headers=RespHeaders}) ->
+	NameBin = header_to_binary(Name),
+	lists:keymember(NameBin, 1, RespHeaders).
+
+%% @doc Return whether a body has been set for the response.
+-spec has_resp_body(#http_req{}) -> boolean().
+has_resp_body(#http_req{resp_body=RespBody}) ->
+	byte_size(RespBody) > 0.
+
 %% @equiv reply(Status, [], [], Req)
 -spec reply(http_status(), #http_req{}) -> {ok, #http_req{}}.
-reply(Status, Req) ->
-	reply(Status, [], [], Req).
+reply(Status, Req=#http_req{resp_body=Body}) ->
+	reply(Status, [], Body, Req).
 
 %% @equiv reply(Status, Headers, [], Req)
 -spec reply(http_status(), http_headers(), #http_req{}) -> {ok, #http_req{}}.
-reply(Status, Headers, Req) ->
-	reply(Status, Headers, [], Req).
+reply(Status, Headers, Req=#http_req{resp_body=Body}) ->
+	reply(Status, Headers, Body, Req).
 
 %% @doc Send a reply to the client.
 -spec reply(http_status(), http_headers(), iodata(), #http_req{})
 	-> {ok, #http_req{}}.
 reply(Status, Headers, Body, Req=#http_req{socket=Socket,
 		transport=Transport, connection=Connection,
-		method=Method, resp_state=waiting}) ->
+		method=Method, resp_state=waiting, resp_headers=RespHeaders}) ->
 	RespConn = response_connection(Headers, Connection),
-	Head = response_head(Status, Headers, [
+	Head = response_head(Status, Headers, RespHeaders, [
 		{<<"Connection">>, atom_to_connection(Connection)},
 		{<<"Content-Length">>,
 			list_to_binary(integer_to_list(iolist_size(Body)))},
@@ -388,7 +416,8 @@ reply(Status, Headers, Body, Req=#http_req{socket=Socket,
 		'HEAD' -> Transport:send(Socket, Head);
 		_ -> Transport:send(Socket, [Head, Body])
 	end,
-	{ok, Req#http_req{connection=RespConn, resp_state=done}}.
+	{ok, Req#http_req{connection=RespConn, resp_state=done,
+		resp_headers=[], resp_body= <<>>}}.
 
 %% @equiv chunked_reply(Status, [], Req)
 -spec chunked_reply(http_status(), #http_req{}) -> {ok, #http_req{}}.
@@ -400,16 +429,17 @@ chunked_reply(Status, Req) ->
 -spec chunked_reply(http_status(), http_headers(), #http_req{})
 	-> {ok, #http_req{}}.
 chunked_reply(Status, Headers, Req=#http_req{socket=Socket, transport=Transport,
-		connection=Connection, resp_state=waiting}) ->
+		connection=Connection, resp_state=waiting, resp_headers=RespHeaders}) ->
 	RespConn = response_connection(Headers, Connection),
-	Head = response_head(Status, Headers, [
+	Head = response_head(Status, Headers, RespHeaders, [
 		{<<"Connection">>, atom_to_connection(Connection)},
 		{<<"Transfer-Encoding">>, <<"chunked">>},
 		{<<"Date">>, cowboy_clock:rfc1123()},
 		{<<"Server">>, <<"Cowboy">>}
 	]),
 	Transport:send(Socket, Head),
-	{ok, Req#http_req{connection=RespConn, resp_state=chunks}}.
+	{ok, Req#http_req{connection=RespConn, resp_state=chunks,
+		resp_headers=[], resp_body= <<>>}}.
 
 %% @doc Send a chunk of data.
 %%
@@ -425,12 +455,12 @@ chunk(Data, #http_req{socket=Socket, transport=Transport, resp_state=chunks}) ->
 -spec upgrade_reply(http_status(), http_headers(), #http_req{})
 	-> {ok, #http_req{}}.
 upgrade_reply(Status, Headers, Req=#http_req{socket=Socket, transport=Transport,
-		resp_state=waiting}) ->
-	Head = response_head(Status, Headers, [
+		resp_state=waiting, resp_headers=RespHeaders}) ->
+	Head = response_head(Status, Headers, RespHeaders, [
 		{<<"Connection">>, <<"Upgrade">>}
 	]),
 	Transport:send(Socket, Head),
-	{ok, Req#http_req{resp_state=done}}.
+	{ok, Req#http_req{resp_state=done, resp_headers=[], resp_body= <<>>}}.
 
 %% Misc API.
 
@@ -478,15 +508,27 @@ response_connection_parse(ReplyConn) ->
 	Tokens = cowboy_http:nonempty_list(ReplyConn, fun cowboy_http:token/2),
 	cowboy_http:connection_to_atom(Tokens).
 
--spec response_head(http_status(), http_headers(), http_headers()) -> iolist().
-response_head(Status, Headers, DefaultHeaders) ->
+-spec response_head(http_status(), http_headers(), http_headers(),
+	http_headers()) -> iolist().
+response_head(Status, Headers, RespHeaders, DefaultHeaders) ->
 	StatusLine = <<"HTTP/1.1 ", (status(Status))/binary, "\r\n">>,
 	Headers2 = [{header_to_binary(Key), Value} || {Key, Value} <- Headers],
-	Headers3 = lists:keysort(1, Headers2),
-	Headers4 = lists:ukeymerge(1, Headers3, DefaultHeaders),
-	Headers5 = [[Key, <<": ">>, Value, <<"\r\n">>]
-		|| {Key, Value} <- Headers4],
-	[StatusLine, Headers5, <<"\r\n">>].
+	Headers3 = merge_headers(
+		merge_headers(Headers2, RespHeaders),
+		DefaultHeaders),
+	Headers4 = [[Key, <<": ">>, Value, <<"\r\n">>]
+		|| {Key, Value} <- Headers3],
+	[StatusLine, Headers4, <<"\r\n">>].
+
+-spec merge_headers(http_headers(), http_headers()) -> http_headers().
+merge_headers(Headers, []) ->
+	Headers;
+merge_headers(Headers, [{Name, Value}|Tail]) ->
+	Headers2 = case lists:keymember(Name, 1, Headers) of
+		true -> Headers;
+		false -> Headers ++ [{Name, Value}]
+	end,
+	merge_headers(Headers2, Tail).
 
 -spec atom_to_connection(keepalive) -> <<_:80>>;
 						(close) -> <<_:40>>.
diff --git a/test/http_SUITE.erl b/test/http_SUITE.erl
index bfccd3b..5f02a7c 100644
--- a/test/http_SUITE.erl
+++ b/test/http_SUITE.erl
@@ -21,7 +21,8 @@
 -export([chunked_response/1, headers_dupe/1, headers_huge/1,
 	keepalive_nl/1, nc_rand/1, nc_zero/1, pipeline/1, raw/1,
 	ws0/1, ws8/1, ws8_single_bytes/1, ws8_init_shutdown/1,
-	ws13/1, ws_timeout_hibernate/1]). %% http.
+	ws13/1, ws_timeout_hibernate/1, set_resp_header/1,
+	set_resp_overwrite/1, set_resp_body/1]). %% http.
 -export([http_200/1, http_404/1]). %% http and https.
 -export([http_10_hostless/1]). %% misc.
 
@@ -35,7 +36,8 @@ groups() ->
 	[{http, [], [chunked_response, headers_dupe, headers_huge,
 		keepalive_nl, nc_rand, nc_zero, pipeline, raw,
 		ws0, ws8, ws8_single_bytes, ws8_init_shutdown, ws13,
-		ws_timeout_hibernate] ++ BaseTests},
+		ws_timeout_hibernate, set_resp_header,
+		set_resp_overwrite, set_resp_body] ++ BaseTests},
 	{https, [], BaseTests}, {misc, [], [http_10_hostless]}].
 
 init_per_suite(Config) ->
@@ -100,6 +102,12 @@ init_http_dispatch() ->
 			{[<<"long_polling">>], http_handler_long_polling, []},
 			{[<<"headers">>, <<"dupe">>], http_handler,
 				[{headers, [{<<"Connection">>, <<"close">>}]}]},
+			{[<<"set_resp">>, <<"header">>], http_handler_set_resp,
+				[{headers, [{<<"Vary">>, <<"Accept">>}]}]},
+			{[<<"set_resp">>, <<"overwrite">>], http_handler_set_resp,
+				[{headers, [{<<"Server">>, <<"DesireDrive/1.0">>}]}]},
+			{[<<"set_resp">>, <<"body">>], http_handler_set_resp,
+				[{body, <<"A flameless dance does not equal a cycle">>}]},
 			{[], http_handler, []}
 		]}
 	].
@@ -479,6 +487,34 @@ websocket_headers({ok, {http_header, _I, Key, _R, Value}, Rest}, Acc) ->
 	websocket_headers(erlang:decode_packet(httph, Rest, []),
 		[{F(Key), Value}|Acc]).
 
+set_resp_header(Config) ->
+	{port, Port} = lists:keyfind(port, 1, Config),
+	{ok, Socket} = gen_tcp:connect("localhost", Port,
+		[binary, {active, false}, {packet, raw}]),
+	ok = gen_tcp:send(Socket, "GET /set_resp/header HTTP/1.1\r\n"
+		"Host: localhost\r\nConnection: close\r\n\r\n"),
+	{ok, Data} = gen_tcp:recv(Socket, 0, 6000),
+	{_Start, _Length} = binary:match(Data, <<"Vary: Accept">>).
+
+set_resp_overwrite(Config) ->
+	{port, Port} = lists:keyfind(port, 1, Config),
+	{ok, Socket} = gen_tcp:connect("localhost", Port,
+		[binary, {active, false}, {packet, raw}]),
+	ok = gen_tcp:send(Socket, "GET /set_resp/overwrite HTTP/1.1\r\n"
+		"Host: localhost\r\nConnection: close\r\n\r\n"),
+	{ok, Data} = gen_tcp:recv(Socket, 0, 6000),
+	{_Start, _Length} = binary:match(Data, <<"Server: DesireDrive/1.0">>).
+
+set_resp_body(Config) ->
+	{port, Port} = lists:keyfind(port, 1, Config),
+	{ok, Socket} = gen_tcp:connect("localhost", Port,
+		[binary, {active, false}, {packet, raw}]),
+	ok = gen_tcp:send(Socket, "GET /set_resp/body HTTP/1.1\r\n"
+		"Host: localhost\r\nConnection: close\r\n\r\n"),
+	{ok, Data} = gen_tcp:recv(Socket, 0, 6000),
+	{_Start, _Length} = binary:match(Data, <<"\r\n\r\n"
+		"A flameless dance does not equal a cycle">>).
+
 %% http and https.
 
 build_url(Path, Config) ->
diff --git a/test/http_handler_set_resp.erl b/test/http_handler_set_resp.erl
new file mode 100644
index 0000000..6aca73d
--- /dev/null
+++ b/test/http_handler_set_resp.erl
@@ -0,0 +1,31 @@
+%% Feel free to use, reuse and abuse the code in this file.
+
+-module(http_handler_set_resp).
+-behaviour(cowboy_http_handler).
+-export([init/3, handle/2, terminate/2]).
+
+init({_Transport, http}, Req, Opts) ->
+	Headers = proplists:get_value(headers, Opts, []),
+	Body = proplists:get_value(body, Opts, <<"http_handler_set_resp">>),
+	{ok, Req2} = lists:foldl(fun({Name, Value}, {ok, R}) ->
+		cowboy_http_req:set_resp_header(Name, Value, R)
+	end, {ok, Req}, Headers),
+	{ok, Req3} = cowboy_http_req:set_resp_body(Body, Req2),
+	{ok, Req4} = cowboy_http_req:set_resp_header(
+		<<"X-Cowboy-Test">>, <<"ok">>, Req3),
+	{ok, Req4, undefined}.
+
+handle(Req, State) ->
+	case cowboy_http_req:has_resp_header(<<"X-Cowboy-Test">>, Req) of
+		false -> {ok, Req, State};
+		true ->
+			case cowboy_http_req:has_resp_body(Req) of
+				false -> {ok, Req, State};
+				true ->
+					{ok, Req2} = cowboy_http_req:reply(200, Req),
+					{ok, Req2, State}
+			end
+	end.
+
+terminate(_Req, _State) ->
+	ok.