Fix an issue for replies in onresponse mishandling headers

Reported and fixed over email by Adrian Roe.

2 files changed, 41 insertions, 4 deletions

diff --git a/src/cowboy_req.erl b/src/cowboy_req.erl
index 849831d..7f7ef32 100644
--- a/src/cowboy_req.erl
+++ b/src/cowboy_req.erl
@@ -163,7 +163,8 @@
 		| {non_neg_integer(), resp_body_fun()},
 
 	%% Functions.
-	onresponse = undefined :: undefined | cowboy_protocol:onresponse_fun()
+	onresponse = undefined :: undefined | already_called
+		| cowboy_protocol:onresponse_fun()
 }).
 
 -opaque req() :: #http_req{}.
@@ -1162,13 +1163,17 @@ to_list(Req) ->
 response(Status, Headers, RespHeaders, DefaultHeaders, Body, Req=#http_req{
 		socket=Socket, transport=Transport, version=Version,
 		pid=ReqPid, onresponse=OnResponse}) ->
-	FullHeaders = response_merge_headers(Headers, RespHeaders, DefaultHeaders),
+	FullHeaders = case OnResponse of
+		already_called -> Headers;
+		_ -> response_merge_headers(Headers, RespHeaders, DefaultHeaders)
+	end,
 	Req2 = case OnResponse of
+		already_called -> Req;
 		undefined -> Req;
 		OnResponse -> OnResponse(Status, FullHeaders, Body,
 			%% Don't call 'onresponse' from the hook itself.
 			Req#http_req{resp_headers=[], resp_body= <<>>,
-				onresponse=undefined})
+				onresponse=already_called})
 	end,
 	ReplyType = case Req2#http_req.resp_state of
 		waiting ->
diff --git a/test/http_SUITE.erl b/test/http_SUITE.erl
index 3200188..67920a8 100644
--- a/test/http_SUITE.erl
+++ b/test/http_SUITE.erl
@@ -45,6 +45,7 @@
 -export([nc_zero/1]).
 -export([onrequest/1]).
 -export([onrequest_reply/1]).
+-export([onresponse_capitalize/1]).
 -export([onresponse_crash/1]).
 -export([onresponse_reply/1]).
 -export([pipeline/1]).
@@ -84,7 +85,8 @@ all() ->
 		{group, http_compress},
 		{group, https_compress},
 		{group, onrequest},
-		{group, onresponse}
+		{group, onresponse},
+		{group, onresponse_capitalize}
 	].
 
 groups() ->
@@ -146,6 +148,9 @@ groups() ->
 		{onresponse, [], [
 			onresponse_crash,
 			onresponse_reply
+		]},
+		{onresponse_capitalize, [], [
+			onresponse_capitalize
 		]}
 	].
 
@@ -250,6 +255,18 @@ init_per_group(onresponse, Config) ->
 	]),
 	{ok, Client} = cowboy_client:init([]),
 	[{scheme, <<"http">>}, {port, Port}, {opts, []},
+		{transport, Transport}, {client, Client}|Config];
+init_per_group(onresponse_capitalize, Config) ->
+	Port = 33086,
+	Transport = ranch_tcp,
+	{ok, _} = cowboy:start_http(onresponse_capitalize, 100, [{port, Port}], [
+		{env, [{dispatch, init_dispatch(Config)}]},
+		{max_keepalive, 50},
+		{onresponse, fun onresponse_capitalize_hook/4},
+		{timeout, 500}
+	]),
+	{ok, Client} = cowboy_client:init([]),
+	[{scheme, <<"http">>}, {port, Port}, {opts, []},
 		{transport, Transport}, {client, Client}|Config].
 
 end_per_group(Group, Config) when Group =:= https; Group =:= https_compress ->
@@ -671,6 +688,21 @@ onrequest_hook(Req) ->
 			Req3
 	end.
 
+onresponse_capitalize(Config) ->
+	Client = ?config(client, Config),
+	{ok, Client2} = cowboy_client:request(<<"GET">>,
+		build_url("/", Config), Client),
+	{ok, Transport, Socket} = cowboy_client:transport(Client2),
+	{ok, Data} = Transport:recv(Socket, 0, 1000),
+	false = nomatch =:= binary:match(Data, <<"Content-Length">>).
+
+%% Hook for the above onresponse_capitalize test.
+onresponse_capitalize_hook(Status, Headers, Body, Req) ->
+	Headers2 = [{cowboy_bstr:capitalize_token(N), V}
+		|| {N, V} <- Headers],
+	{ok, Req2} = cowboy_req:reply(Status, Headers2, Body, Req),
+	Req2.
+
 onresponse_crash(Config) ->
 	Client = ?config(client, Config),
 	{ok, Client2} = cowboy_client:request(<<"GET">>,