Make the default 204 response go through stream handlers

2 files changed, 13 insertions, 11 deletions

diff --git a/src/cowboy_http.erl b/src/cowboy_http.erl
index 983d95f..002944d 100644
--- a/src/cowboy_http.erl
+++ b/src/cowboy_http.erl
@@ -913,19 +913,19 @@ stream_reset(State, StreamID, StreamError={internal_error, _, _}) ->
 %	stream_terminate(State#state{out_state=done}, StreamID, StreamError).
 	stream_terminate(State, StreamID, StreamError).
 
-stream_terminate(State=#state{socket=Socket, transport=Transport,
+stream_terminate(State0=#state{socket=Socket, transport=Transport,
 		out_streamid=OutStreamID, out_state=OutState,
 		streams=Streams0, children=Children0}, StreamID, Reason) ->
 	{value, #stream{state=StreamState, version=Version}, Streams}
 		= lists:keytake(StreamID, #stream.id, Streams0),
-	_ = case OutState of
+	State = case OutState of
 		wait ->
-			%% @todo This should probably go through the stream handler info callback.
-			Transport:send(Socket, cow_http:response(204, 'HTTP/1.1', []));
+			info(State0, StreamID, {response, 204, #{}, <<>>});
 		chunked when Version =:= 'HTTP/1.1' ->
-			Transport:send(Socket, <<"0\r\n\r\n">>);
+			_ = Transport:send(Socket, <<"0\r\n\r\n">>),
+			State0;
 		_ -> %% done or Version =:= 'HTTP/1.0'
-			ok
+			State0
 	end,
 
 	stream_call_terminate(StreamID, Reason, StreamState),
diff --git a/src/cowboy_http2.erl b/src/cowboy_http2.erl
index afac968..d93a6cd 100644
--- a/src/cowboy_http2.erl
+++ b/src/cowboy_http2.erl
@@ -655,20 +655,22 @@ stream_reset(State=#state{socket=Socket, transport=Transport}, StreamID,
 	stream_terminate(State, StreamID, StreamError).
 
 stream_terminate(State=#state{socket=Socket, transport=Transport,
-		streams=Streams0, children=Children0, encode_state=EncodeState0}, StreamID, Reason) ->
+		streams=Streams0, children=Children0}, StreamID, Reason) ->
 	case lists:keytake(StreamID, #stream.id, Streams0) of
+		%% When the stream terminates normally (without sending RST_STREAM)
+		%% and no response was sent, we need to send a proper response back to the client.
 		{value, #stream{state=StreamState, local=idle}, Streams} when Reason =:= normal ->
-			Headers = #{<<":status">> => <<"204">>},
-			{HeaderBlock, EncodeState} = headers_encode(Headers, EncodeState0),
-			Transport:send(Socket, cow_http2:headers(StreamID, fin, HeaderBlock)),
+			State1 = info(State, StreamID, {response, 204, #{}, <<>>}),
 			stream_call_terminate(StreamID, Reason, StreamState),
 			Children = stream_terminate_children(Children0, StreamID, []),
-			State#state{streams=Streams, children=Children, encode_state=EncodeState};
+			State1#state{streams=Streams, children=Children};
+		%% When a response was sent but not terminated, we need to close the stream.
 		{value, #stream{state=StreamState, local=nofin}, Streams} when Reason =:= normal ->
 			Transport:send(Socket, cow_http2:data(StreamID, fin, <<>>)),
 			stream_call_terminate(StreamID, Reason, StreamState),
 			Children = stream_terminate_children(Children0, StreamID, []),
 			State#state{streams=Streams, children=Children};
+		%% Otherwise we sent an RST_STREAM and the stream is already closed.
 		{value, #stream{state=StreamState}, Streams} ->
 			stream_call_terminate(StreamID, Reason, StreamState),
 			Children = stream_terminate_children(Children0, StreamID, []),