Add an 'onresponse' hook

This new protocol option is a fun.

It expects 3 args: the Status code used in the reply (this is the
cowboy_http:status() type, it can be an integer or a binary), the
headers that will be sent in the reply, and the Req. It should
only return a possibly modified Req. This can be used for many
things like error logging or custom error pages.

If a reply is sent inside the hook, then Cowboy will discard the
reply initially sent. Extra caution must be used in the handlers
making use of inline chunked replies as they will throw an error.

This fun cannot be used as a filter, you can either observe the
reply sent or discard it to send a different one instead.

The hook will not be called for replies sent from inside the hook.

3 files changed, 57 insertions, 31 deletions

diff --git a/src/cowboy_client.erl b/src/cowboy_client.erl
index 21931e1..e46619f 100644
--- a/src/cowboy_client.erl
+++ b/src/cowboy_client.erl
@@ -158,7 +158,9 @@ response_body_loop(Client, Acc) ->
 		{ok, Data, Client2} ->
 			response_body_loop(Client2, << Acc/binary, Data/binary >>);
 		{done, Client2} ->
-			{ok, Acc, Client2}
+			{ok, Acc, Client2};
+		{error, Reason} ->
+			{error, Reason}
 	end.
 
 skip_body(Client=#client{state=response_body}) ->
diff --git a/src/cowboy_http_protocol.erl b/src/cowboy_http_protocol.erl
index 816c825..9fad2d0 100644
--- a/src/cowboy_http_protocol.erl
+++ b/src/cowboy_http_protocol.erl
@@ -48,6 +48,8 @@
 	dispatch :: cowboy_dispatcher:dispatch_rules(),
 	handler :: {module(), any()},
 	onrequest :: undefined | fun((#http_req{}) -> #http_req{}),
+	onresponse = undefined :: undefined | fun((cowboy_http:status(),
+		cowboy_http:headers(), #http_req{}) -> #http_req{}),
 	urldecode :: {fun((binary(), T) -> binary()), T},
 	req_empty_lines = 0 :: integer(),
 	max_empty_lines :: integer(),
@@ -79,6 +81,7 @@ init(ListenerPid, Socket, Transport, Opts) ->
 	MaxKeepalive = proplists:get_value(max_keepalive, Opts, infinity),
 	MaxLineLength = proplists:get_value(max_line_length, Opts, 4096),
 	OnRequest = proplists:get_value(onrequest, Opts),
+	OnResponse = proplists:get_value(onresponse, Opts),
 	Timeout = proplists:get_value(timeout, Opts, 5000),
 	URLDecDefault = {fun cowboy_http:urldecode/2, crash},
 	URLDec = proplists:get_value(urldecode, Opts, URLDecDefault),
@@ -86,7 +89,8 @@ init(ListenerPid, Socket, Transport, Opts) ->
 	wait_request(#state{listener=ListenerPid, socket=Socket, transport=Transport,
 		dispatch=Dispatch, max_empty_lines=MaxEmptyLines,
 		max_keepalive=MaxKeepalive, max_line_length=MaxLineLength,
-		timeout=Timeout, onrequest=OnRequest, urldecode=URLDec}).
+		timeout=Timeout, onrequest=OnRequest, onresponse=OnResponse,
+		urldecode=URLDec}).
 
 %% @private
 -spec parse_request(#state{}) -> ok.
@@ -122,7 +126,7 @@ request({http_request, Method, {absoluteURI, _Scheme, _Host, _Port, Path},
 request({http_request, Method, {abs_path, AbsPath}, Version},
 		State=#state{socket=Socket, transport=Transport,
 		req_keepalive=Keepalive, max_keepalive=MaxKeepalive,
-		urldecode={URLDecFun, URLDecArg}=URLDec}) ->
+		onresponse=OnResponse, urldecode={URLDecFun, URLDecArg}=URLDec}) ->
 	URLDecode = fun(Bin) -> URLDecFun(Bin, URLDecArg) end,
 	{Path, RawPath, Qs} = cowboy_dispatcher:split_path(AbsPath, URLDecode),
 	ConnAtom = if Keepalive < MaxKeepalive -> version_to_connection(Version);
@@ -130,16 +134,19 @@ request({http_request, Method, {abs_path, AbsPath}, Version},
 	end,
 	parse_header(#http_req{socket=Socket, transport=Transport,
 		connection=ConnAtom, pid=self(), method=Method, version=Version,
-		path=Path, raw_path=RawPath, raw_qs=Qs, urldecode=URLDec}, State);
+		path=Path, raw_path=RawPath, raw_qs=Qs, onresponse=OnResponse,
+		urldecode=URLDec}, State);
 request({http_request, Method, '*', Version},
 		State=#state{socket=Socket, transport=Transport,
-		req_keepalive=Keepalive, max_keepalive=MaxKeepalive, urldecode=URLDec}) ->
+		req_keepalive=Keepalive, max_keepalive=MaxKeepalive,
+		onresponse=OnResponse, urldecode=URLDec}) ->
 	ConnAtom = if Keepalive < MaxKeepalive -> version_to_connection(Version);
 		true -> close
 	end,
 	parse_header(#http_req{socket=Socket, transport=Transport,
 		connection=ConnAtom, pid=self(), method=Method, version=Version,
-		path='*', raw_path= <<"*">>, raw_qs= <<>>, urldecode=URLDec}, State);
+		path='*', raw_path= <<"*">>, raw_qs= <<>>, onresponse=OnResponse,
+		urldecode=URLDec}, State);
 request({http_request, _Method, _URI, _Version}, State) ->
 	error_terminate(501, State);
 request({http_error, <<"\r\n">>},
diff --git a/src/cowboy_http_req.erl b/src/cowboy_http_req.erl
index cdb44ba..f25c5eb 100644
--- a/src/cowboy_http_req.erl
+++ b/src/cowboy_http_req.erl
@@ -696,7 +696,7 @@ reply(Status, Headers, Req=#http_req{resp_body=Body}) ->
 -spec reply(cowboy_http:status(), cowboy_http:headers(), iodata(), #http_req{})
 	-> {ok, #http_req{}}.
 reply(Status, Headers, Body, Req=#http_req{socket=Socket, transport=Transport,
-		version=Version, connection=Connection, pid=ReqPid,
+		version=Version, 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,
@@ -704,18 +704,20 @@ reply(Status, Headers, Body, Req=#http_req{socket=Socket, transport=Transport,
 		{1, 1} -> [{<<"Connection">>, atom_to_connection(Connection)}];
 		_ -> []
 	end,
-	response(Status, Headers, RespHeaders,  [
+	{ReplyType, Req2} = response(Status, Headers, RespHeaders,  [
 		{<<"Content-Length">>, integer_to_list(ContentLen)},
 		{<<"Date">>, cowboy_clock:rfc1123()},
 		{<<"Server">>, <<"Cowboy">>}
 	|HTTP11Headers], Req),
-	case {Method, Body} of
-		{'HEAD', _} -> ok;
-		{_, {_, StreamFun}} -> StreamFun();
-		{_, _} -> Transport:send(Socket, Body)
+	if	Method =:= 'HEAD' -> ok;
+		ReplyType =:= hook -> ok; %% Hook replied for us, stop there.
+		true ->
+			case Body of
+				{_, StreamFun} -> StreamFun();
+				_ -> Transport:send(Socket, Body)
+			end
 	end,
-	ReqPid ! {?MODULE, resp_sent},
-	{ok, Req#http_req{connection=RespConn, resp_state=done,
+	{ok, Req2#http_req{connection=RespConn, resp_state=done,
 		resp_headers=[], resp_body= <<>>}}.
 
 %% @equiv chunked_reply(Status, [], Req)
@@ -729,7 +731,7 @@ chunked_reply(Status, Req) ->
 	-> {ok, #http_req{}}.
 chunked_reply(Status, Headers, Req=#http_req{
 		version=Version, connection=Connection,
-		pid=ReqPid, resp_state=waiting, resp_headers=RespHeaders}) ->
+		resp_state=waiting, resp_headers=RespHeaders}) ->
 	RespConn = response_connection(Headers, Connection),
 	HTTP11Headers = case Version of
 		{1, 1} -> [
@@ -737,12 +739,11 @@ chunked_reply(Status, Headers, Req=#http_req{
 			{<<"Transfer-Encoding">>, <<"chunked">>}];
 		_ -> []
 	end,
-	response(Status, Headers, RespHeaders, [
+	{_, Req2} = response(Status, Headers, RespHeaders, [
 		{<<"Date">>, cowboy_clock:rfc1123()},
 		{<<"Server">>, <<"Cowboy">>}
 	|HTTP11Headers], Req),
-	ReqPid ! {?MODULE, resp_sent},
-	{ok, Req#http_req{connection=RespConn, resp_state=chunks,
+	{ok, Req2#http_req{connection=RespConn, resp_state=chunks,
 		resp_headers=[], resp_body= <<>>}}.
 
 %% @doc Send a chunk of data.
@@ -762,12 +763,11 @@ chunk(Data, #http_req{socket=Socket, transport=Transport, resp_state=chunks}) ->
 -spec upgrade_reply(cowboy_http:status(), cowboy_http:headers(), #http_req{})
 	-> {ok, #http_req{}}.
 upgrade_reply(Status, Headers, Req=#http_req{
-		pid=ReqPid, resp_state=waiting, resp_headers=RespHeaders}) ->
-	response(Status, Headers, RespHeaders, [
+		resp_state=waiting, resp_headers=RespHeaders}) ->
+	{_, Req2} = response(Status, Headers, RespHeaders, [
 		{<<"Connection">>, <<"Upgrade">>}
 	], Req),
-	ReqPid ! {?MODULE, resp_sent},
-	{ok, Req#http_req{resp_state=done, resp_headers=[], resp_body= <<>>}}.
+	{ok, Req2#http_req{resp_state=done, resp_headers=[], resp_body= <<>>}}.
 
 %% Misc API.
 
@@ -798,16 +798,33 @@ transport(#http_req{transport=Transport, socket=Socket}) ->
 %% Internal.
 
 -spec response(cowboy_http:status(), cowboy_http:headers(),
-	cowboy_http:headers(), cowboy_http:headers(), #http_req{}) -> ok.
-response(Status, Headers, RespHeaders, DefaultHeaders, #http_req{
-		socket=Socket, transport=Transport, version=Version}) ->
+	cowboy_http:headers(), cowboy_http:headers(), #http_req{})
+	-> {normal | hook, #http_req{}}.
+response(Status, Headers, RespHeaders, DefaultHeaders, Req=#http_req{
+		socket=Socket, transport=Transport, version=Version,
+		pid=ReqPid, onresponse=OnResponse}) ->
 	FullHeaders = response_merge_headers(Headers, RespHeaders, DefaultHeaders),
-	%% @todo 'onresponse' hook here.
-	HTTPVer = cowboy_http:version_to_binary(Version),
-	StatusLine = << HTTPVer/binary, " ", (status(Status))/binary, "\r\n" >>,
-	HeaderLines = [[Key, <<": ">>, Value, <<"\r\n">>]
-		|| {Key, Value} <- FullHeaders],
-	Transport:send(Socket, [StatusLine, HeaderLines, <<"\r\n">>]).
+	Req2 = case OnResponse of
+		undefined -> Req;
+		OnResponse -> OnResponse(Status, FullHeaders,
+			%% Don't call 'onresponse' from the hook itself.
+			Req#http_req{resp_headers=[], resp_body= <<>>,
+				onresponse=undefined})
+	end,
+	ReplyType = case Req2#http_req.resp_state of
+		waiting ->
+			HTTPVer = cowboy_http:version_to_binary(Version),
+			StatusLine = << HTTPVer/binary, " ",
+				(status(Status))/binary, "\r\n" >>,
+			HeaderLines = [[Key, <<": ">>, Value, <<"\r\n">>]
+				|| {Key, Value} <- FullHeaders],
+			Transport:send(Socket, [StatusLine, HeaderLines, <<"\r\n">>]),
+			ReqPid ! {?MODULE, resp_sent},
+			normal;
+		_ ->
+			hook
+	end,
+	{ReplyType, Req2}.
 
 -spec response_connection(cowboy_http:headers(), keepalive | close)
 	-> keepalive | close.