From 638638a841f22eacbebd827c0fb40002358bf1ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Hoguin?= <essen@ninenines.eu>
Date: Tue, 29 Jan 2013 14:35:26 +0100
Subject: Fix {cowboy_req, resp_sent} potentially leaking in loop handlers

---
 src/cowboy_handler.erl   | 21 +++++++++++++++++----
 src/cowboy_websocket.erl |  1 +
 2 files changed, 18 insertions(+), 4 deletions(-)

(limited to 'src')

diff --git a/src/cowboy_handler.erl b/src/cowboy_handler.erl
index 4acb876..7ed7db3 100644
--- a/src/cowboy_handler.erl
+++ b/src/cowboy_handler.erl
@@ -29,7 +29,8 @@
 	env :: cowboy_middleware:env(),
 	hibernate = false :: boolean(),
 	loop_timeout = infinity :: timeout(),
-	loop_timeout_ref :: undefined | reference()
+	loop_timeout_ref :: undefined | reference(),
+	resp_sent = false :: boolean()
 }).
 
 %% @private
@@ -80,7 +81,7 @@ handler_init(Req, State, Handler, HandlerOpts) ->
 			"** Stacktrace: ~p~n~n",
 			[Handler, init, 3, Class, Reason, HandlerOpts,
 				cowboy_req:to_list(Req), erlang:get_stacktrace()]),
-		{error, 500, Req}
+		error_terminate(Req, State)
 	end.
 
 -spec upgrade_protocol(Req, #state{}, module(), any(), module())
@@ -112,7 +113,7 @@ handler_handle(Req, State, Handler, HandlerState) ->
 			[Handler, handle, 2, Class, Reason, HandlerState,
 				cowboy_req:to_list(Req), erlang:get_stacktrace()]),
 		handler_terminate(Req, Handler, HandlerState, Reason),
-		{error, 500, Req}
+		error_terminate(Req, State)
 	end.
 
 %% We don't listen for Transport closes because that would force us
@@ -147,6 +148,9 @@ handler_loop_timeout(State=#state{loop_timeout=Timeout,
 	when Req::cowboy_req:req().
 handler_loop(Req, State=#state{loop_timeout_ref=TRef}, Handler, HandlerState) ->
 	receive
+		{cowboy_req, resp_sent} ->
+			handler_loop(Req, State#state{resp_sent=true},
+				Handler, HandlerState);
 		{timeout, TRef, ?MODULE} ->
 			terminate_request(Req, State, Handler, HandlerState,
 				{normal, timeout});
@@ -180,7 +184,7 @@ handler_call(Req, State, Handler, HandlerState, Message) ->
 			[Handler, info, 3, Class, Reason, HandlerState,
 				cowboy_req:to_list(Req), erlang:get_stacktrace()]),
 		handler_terminate(Req, Handler, HandlerState, Reason),
-		{error, 500, Req}
+		error_terminate(Req, State)
 	end.
 
 -spec terminate_request(Req, #state{}, module(), any(),
@@ -205,3 +209,12 @@ handler_terminate(Req, Handler, HandlerState, Reason) ->
 			[Handler, terminate, 3, Class, Reason2, HandlerState,
 				cowboy_req:to_list(Req), erlang:get_stacktrace()])
 	end.
+
+%% Only send an error reply if there is no resp_sent message.
+-spec error_terminate(Req, #state{})
+	-> {error, 500, Req} | {halt, Req} when Req::cowboy_req:req().
+error_terminate(Req, #state{resp_sent=true}) ->
+	%% Close the connection, but do not attempt sending a reply.
+	{halt, cowboy_req:set([{connection, close}, {resp_state, done}], Req)};
+error_terminate(Req, _) ->
+	{error, 500, Req}.
diff --git a/src/cowboy_websocket.erl b/src/cowboy_websocket.erl
index 1b2ad3b..debb69f 100644
--- a/src/cowboy_websocket.erl
+++ b/src/cowboy_websocket.erl
@@ -121,6 +121,7 @@ handler_init(State=#state{env=Env, transport=Transport,
 		upgrade_error(Req, Env)
 	end.
 
+%% Only send an error reply if there is no resp_sent message.
 -spec upgrade_error(Req, Env) -> {ok, Req, Env} | {error, 400, Req}
 	when Req::cowboy_req:req(), Env::cowboy_middleware:env().
 upgrade_error(Req, Env) ->
-- 
cgit v1.2.3