From 64a40cb479e45226c3498133c4e198a6dc35a3f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Hoguin?= Date: Mon, 28 Nov 2011 09:09:41 +0100 Subject: 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. --- include/http.hrl | 4 ++- src/cowboy_http_req.erl | 82 +++++++++++++++++++++++++++++++----------- test/http_SUITE.erl | 40 +++++++++++++++++++-- test/http_handler_set_resp.erl | 31 ++++++++++++++++ 4 files changed, 134 insertions(+), 23 deletions(-) create mode 100644 test/http_handler_set_resp.erl 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. -- cgit v1.2.3