Add cowboy_http_req:set_resp_body_fun/3.

4 files changed, 73 insertions, 10 deletions

diff --git a/include/http.hrl b/include/http.hrl
index 84b9489..c47a244 100644
--- a/include/http.hrl
+++ b/include/http.hrl
@@ -36,6 +36,8 @@
 -type http_headers() :: list({http_header(), iodata()}).
 -type http_cookies() :: list({binary(), binary()}).
 -type http_status() :: non_neg_integer() | binary().
+-type http_resp_body() :: iodata() | {non_neg_integer(),
+		fun(() -> {sent, non_neg_integer()})}.
 
 -record(http_req, {
 	%% Transport.
@@ -69,7 +71,7 @@
 	%% Response.
 	resp_state = waiting   :: locked | waiting | chunks | done,
 	resp_headers = []      :: http_headers(),
-	resp_body  = <<>>      :: iodata(),
+	resp_body  = <<>>      :: http_resp_body(),
 
 	%% Functions.
 	urldecode :: {fun((binary(), T) -> binary()), T}
diff --git a/src/cowboy_http_req.erl b/src/cowboy_http_req.erl
index 758346d..cf1790f 100644
--- a/src/cowboy_http_req.erl
+++ b/src/cowboy_http_req.erl
@@ -39,7 +39,7 @@
 
 -export([
 	set_resp_cookie/4, set_resp_header/3, set_resp_body/2,
-	has_resp_header/2, has_resp_body/1,
+	set_resp_body_fun/3, 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
@@ -419,11 +419,33 @@ set_resp_header(Name, Value, Req=#http_req{resp_headers=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.
+%% anything other than reply/2 or reply/3. The response body is expected
+%% to be a binary or an iolist.
 -spec set_resp_body(iodata(), #http_req{}) -> {ok, #http_req{}}.
 set_resp_body(Body, Req) ->
 	{ok, Req#http_req{resp_body=Body}}.
 
+
+%% @doc Add a body function to the response.
+%%
+%% The response body may also be set to a content-length - stream-function pair.
+%% If the response body is of this type normal response headers will be sent.
+%% After the response headers has been sent the body function is applied.
+%% The body function is expected to write the response body directly to the
+%% socket using the transport module.
+%%
+%% If the body function crashes while writing the response body or writes fewer
+%% bytes than declared the behaviour is undefined. The body set here is ignored
+%% if the response is later sent using anything other than `reply/2' or
+%% `reply/3'.
+%%
+%% @see cowboy_http_req:transport/1.
+-spec set_resp_body_fun(non_neg_integer(), fun(() -> {sent, non_neg_integer()}),
+		#http_req{}) -> {ok, #http_req{}}.
+set_resp_body_fun(StreamLen, StreamFun, Req) ->
+	{ok, Req#http_req{resp_body={StreamLen, StreamFun}}}.
+
+
 %% @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}) ->
@@ -432,6 +454,8 @@ has_resp_header(Name, #http_req{resp_headers=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={Length, _}}) ->
+	Length > 0;
 has_resp_body(#http_req{resp_body=RespBody}) ->
 	iolist_size(RespBody) > 0.
 
@@ -452,16 +476,17 @@ reply(Status, Headers, Body, Req=#http_req{socket=Socket,
 		transport=Transport, connection=Connection,
 		method=Method, resp_state=waiting, resp_headers=RespHeaders}) ->
 	RespConn = response_connection(Headers, Connection),
+	ContentLen = case Body of {CL, _} -> CL; _ -> iolist_size(Body) end,
 	Head = response_head(Status, Headers, RespHeaders, [
 		{<<"Connection">>, atom_to_connection(Connection)},
-		{<<"Content-Length">>,
-			list_to_binary(integer_to_list(iolist_size(Body)))},
+		{<<"Content-Length">>, integer_to_list(ContentLen)},
 		{<<"Date">>, cowboy_clock:rfc1123()},
 		{<<"Server">>, <<"Cowboy">>}
 	]),
-	case Method of
-		'HEAD' -> Transport:send(Socket, Head);
-		_ -> Transport:send(Socket, [Head, Body])
+	case {Method, Body} of
+		{'HEAD', _} -> Transport:send(Socket, Head);
+		{_, {_, StreamFun}} -> Transport:send(Socket, Head), StreamFun();
+		{_, _} -> Transport:send(Socket, [Head, Body])
 	end,
 	{ok, Req#http_req{connection=RespConn, resp_state=done,
 		resp_headers=[], resp_body= <<>>}}.
diff --git a/test/http_SUITE.erl b/test/http_SUITE.erl
index 3c4af28..b7fd551 100644
--- a/test/http_SUITE.erl
+++ b/test/http_SUITE.erl
@@ -21,7 +21,7 @@
 -export([chunked_response/1, headers_dupe/1, headers_huge/1,
 	keepalive_nl/1, max_keepalive/1, nc_rand/1, nc_zero/1,
 	pipeline/1, raw/1, set_resp_header/1, set_resp_overwrite/1,
-	set_resp_body/1, response_as_req/1]). %% http.
+	set_resp_body/1, stream_body_set_resp/1, response_as_req/1]). %% http.
 -export([http_200/1, http_404/1]). %% http and https.
 -export([http_10_hostless/1]). %% misc.
 -export([rest_simple/1, rest_keepalive/1]). %% rest.
@@ -36,7 +36,7 @@ groups() ->
 	[{http, [], [chunked_response, headers_dupe, headers_huge,
 		keepalive_nl, max_keepalive, nc_rand, nc_zero, pipeline, raw,
 		set_resp_header, set_resp_overwrite,
-		set_resp_body, response_as_req] ++ BaseTests},
+		set_resp_body, response_as_req, stream_body_set_resp] ++ BaseTests},
 	{https, [], BaseTests},
 	{misc, [], [http_10_hostless]},
 	{rest, [], [rest_simple, rest_keepalive]}].
@@ -115,6 +115,8 @@ init_http_dispatch() ->
 				[{headers, [{<<"Server">>, <<"DesireDrive/1.0">>}]}]},
 			{[<<"set_resp">>, <<"body">>], http_handler_set_resp,
 				[{body, <<"A flameless dance does not equal a cycle">>}]},
+			{[<<"stream_body">>, <<"set_resp">>], http_handler_stream_body,
+				[{reply, set_resp}, {body, <<"stream_body_set_resp">>}]},
 			{[], http_handler, []}
 		]}
 	].
@@ -328,6 +330,16 @@ The document has moved
 </BODY></HTML>",
 	{Packet, 400} = raw_req(Packet, Config).
 
+stream_body_set_resp(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 /stream_body/set_resp 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, <<"stream_body_set_resp">>).
+
+
 %% http and https.
 
 build_url(Path, Config) ->
diff --git a/test/http_handler_stream_body.erl b/test/http_handler_stream_body.erl
new file mode 100644
index 0000000..c90f746
--- /dev/null
+++ b/test/http_handler_stream_body.erl
@@ -0,0 +1,24 @@
+%% Feel free to use, reuse and abuse the code in this file.
+
+-module(http_handler_stream_body).
+-behaviour(cowboy_http_handler).
+-export([init/3, handle/2, terminate/2]).
+
+-record(state, {headers, body, reply}).
+
+init({_Transport, http}, Req, Opts) ->
+	Headers = proplists:get_value(headers, Opts, []),
+	Body = proplists:get_value(body, Opts, "http_handler_stream_body"),
+	Reply = proplists:get_value(reply, Opts),
+	{ok, Req, #state{headers=Headers, body=Body, reply=Reply}}.
+
+handle(Req, State=#state{headers=_Headers, body=Body, reply=set_resp}) ->
+	{ok, Transport, Socket} = cowboy_http_req:transport(Req),
+	SFun = fun() -> Transport:send(Socket, Body), sent end,
+	SLen = iolist_size(Body),
+	{ok, Req2} = cowboy_http_req:set_resp_body_fun(SLen, SFun, Req),
+	{ok, Req3} = cowboy_http_req:reply(200, Req2),
+	{ok, Req3, State}.
+
+terminate(_Req, _State) ->
+	ok.