Remove duplicate code for ensure_response

3 files changed, 27 insertions, 37 deletions

diff --git a/src/cowboy_protocol.erl b/src/cowboy_protocol.erl
index bc56e76..bb94a40 100644
--- a/src/cowboy_protocol.erl
+++ b/src/cowboy_protocol.erl
@@ -403,12 +403,12 @@ terminate_request(HandlerState, Req, State) ->
 -spec next_request(cowboy_req:req(), #state{}, any()) -> ok.
 next_request(Req=#http_req{connection=Conn}, State=#state{
 		req_keepalive=Keepalive}, HandlerRes) ->
-	RespRes = ensure_response(Req),
+	cowboy_req:ensure_response(Req, 204),
 	{BodyRes, Buffer} = ensure_body_processed(Req),
 	%% Flush the resp_sent message before moving on.
 	receive {cowboy_req, resp_sent} -> ok after 0 -> ok end,
-	case {HandlerRes, BodyRes, RespRes, Conn} of
-		{ok, ok, ok, keepalive} ->
+	case {HandlerRes, BodyRes, Conn} of
+		{ok, ok, keepalive} ->
 			?MODULE:parse_request(State#state{
 				buffer=Buffer, host_tokens=undefined, path_tokens=undefined,
 				req_empty_lines=0, req_keepalive=Keepalive + 1});
@@ -428,25 +428,6 @@ ensure_body_processed(Req=#http_req{body_state={stream, _, _, _}}) ->
 	{ok, Req2} = cowboy_req:multipart_skip(Req),
 	ensure_body_processed(Req2).
 
--spec ensure_response(cowboy_req:req()) -> ok.
-%% The handler has already fully replied to the client.
-ensure_response(#http_req{resp_state=done}) ->
-	ok;
-%% No response has been sent but everything apparently went fine.
-%% Reply with 204 No Content to indicate this.
-ensure_response(Req=#http_req{resp_state=waiting}) ->
-	_ = cowboy_req:reply(204, [], [], Req),
-	ok;
-%% Terminate the chunked body for HTTP/1.1 only.
-ensure_response(#http_req{method='HEAD', resp_state=chunks}) ->
-	ok;
-ensure_response(#http_req{version={1, 0}, resp_state=chunks}) ->
-	ok;
-ensure_response(#http_req{socket=Socket, transport=Transport,
-		resp_state=chunks}) ->
-	Transport:send(Socket, <<"0\r\n\r\n">>),
-	ok.
-
 %% Only send an error reply if there is no resp_sent message.
 -spec error_terminate(cowboy_http:status(), #state{}) -> ok.
 error_terminate(Code, State=#state{socket=Socket, transport=Transport,
diff --git a/src/cowboy_req.erl b/src/cowboy_req.erl
index 51b1874..ad6c1a4 100644
--- a/src/cowboy_req.erl
+++ b/src/cowboy_req.erl
@@ -76,6 +76,7 @@
 -export([chunked_reply/3]).
 -export([chunk/2]).
 -export([upgrade_reply/3]).
+-export([ensure_response/2]).
 
 %% Misc API.
 -export([compact/1]).
@@ -826,6 +827,27 @@ upgrade_reply(Status, Headers, Req=#http_req{
 	], Req),
 	{ok, Req2#http_req{resp_state=done, resp_headers=[], resp_body= <<>>}}.
 
+%% @doc Ensure the response has been sent fully.
+%% @private
+-spec ensure_response(req(), cowboy_http:status()) -> ok.
+%% The response has already been fully sent to the client.
+ensure_response(#http_req{resp_state=done}, _) ->
+	ok;
+%% No response has been sent but everything apparently went fine.
+%% Reply with the status code found in the second argument.
+ensure_response(Req=#http_req{resp_state=waiting}, Status) ->
+	_ = reply(Status, [], [], Req),
+	ok;
+%% Terminate the chunked body for HTTP/1.1 only.
+ensure_response(#http_req{method='HEAD', resp_state=chunks}, _) ->
+	ok;
+ensure_response(#http_req{version={1, 0}, resp_state=chunks}, _) ->
+	ok;
+ensure_response(#http_req{socket=Socket, transport=Transport,
+		resp_state=chunks}, _) ->
+	Transport:send(Socket, <<"0\r\n\r\n">>),
+	ok.
+
 %% Misc API.
 
 %% @doc Compact the request data by removing all non-system information.
diff --git a/src/cowboy_websocket.erl b/src/cowboy_websocket.erl
index 2db0faa..55855c7 100644
--- a/src/cowboy_websocket.erl
+++ b/src/cowboy_websocket.erl
@@ -126,7 +126,8 @@ handler_init(State=#state{transport=Transport, handler=Handler, opts=Opts},
 			websocket_handshake(State#state{timeout=Timeout,
 				hibernate=true}, Req2, HandlerState);
 		{shutdown, Req2} ->
-			upgrade_denied(Req2)
+			cowboy_req:ensure_response(Req2, 400),
+			closed
 	catch Class:Reason ->
 		upgrade_error(Req),
 		PLReq = cowboy_req:to_list(Req),
@@ -146,20 +147,6 @@ upgrade_error(Req) ->
 		closed
 	end.
 
-%% @see cowboy_protocol:ensure_response/1
--spec upgrade_denied(cowboy_req:req()) -> closed.
-upgrade_denied(#http_req{resp_state=done}) ->
-	closed;
-upgrade_denied(Req=#http_req{resp_state=waiting}) ->
-	{ok, _Req2} = cowboy_req:reply(400, [], [], Req),
-	closed;
-upgrade_denied(#http_req{method='HEAD', resp_state=chunks}) ->
-	closed;
-upgrade_denied(#http_req{socket=Socket, transport=Transport,
-		resp_state=chunks}) ->
-	Transport:send(Socket, <<"0\r\n\r\n">>),
-	closed.
-
 -spec websocket_handshake(#state{}, cowboy_req:req(), any()) -> closed.
 websocket_handshake(State=#state{socket=Socket, transport=Transport,
 		version=0, origin=Origin, challenge={Key1, Key2}},