Don't send body in HEAD response when using PipeWriter.Advance before headers flushed

src/Servers/Kestrel/Core/src/Internal/Http/Http1OutputProducer.cs

@@ -333,7 +333,7 @@ private void CommitChunkInternal(ref BufferWriter<PipeWriter> writer, ReadOnlySp
         writer.Commit();
     }
 
-    public void WriteResponseHeaders(int statusCode, string? reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk, bool appComplete)
+    public void WriteResponseHeaders(int statusCode, string? reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk, bool appComplete, bool canWriteBody)
     {
         lock (_contextLock)
         {
@@ -346,11 +346,11 @@ public void WriteResponseHeaders(int statusCode, string? reasonPhrase, HttpRespo
 
             var buffer = _pipeWriter;
             var writer = new BufferWriter<PipeWriter>(buffer);
-            WriteResponseHeadersInternal(ref writer, statusCode, reasonPhrase, responseHeaders, autoChunk);
+            WriteResponseHeadersInternal(ref writer, statusCode, reasonPhrase, responseHeaders, autoChunk, canWriteBody);
         }
     }
 
-    private void WriteResponseHeadersInternal(ref BufferWriter<PipeWriter> writer, int statusCode, string? reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk)
+    private void WriteResponseHeadersInternal(ref BufferWriter<PipeWriter> writer, int statusCode, string? reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk, bool canWriteBody)
     {
         writer.Write(HttpVersion11Bytes);
         var statusBytes = ReasonPhrases.ToStatusBytes(statusCode, reasonPhrase);
@@ -361,26 +361,29 @@ private void WriteResponseHeadersInternal(ref BufferWriter<PipeWriter> writer, i
         writer.Commit();
 
         _autoChunk = autoChunk;
-        WriteDataWrittenBeforeHeaders(ref writer);
+        WriteDataWrittenBeforeHeaders(ref writer, canWriteBody);
         _unflushedBytes += writer.BytesCommitted;
 
         _startCalled = true;
     }
 
-    private void WriteDataWrittenBeforeHeaders(ref BufferWriter<PipeWriter> writer)
+    private void WriteDataWrittenBeforeHeaders(ref BufferWriter<PipeWriter> writer, bool canWriteBody)
     {
         if (_completedSegments != null)
         {
             foreach (var segment in _completedSegments)
             {
-                if (_autoChunk)
+                if (canWriteBody)
                 {
-                    CommitChunkInternal(ref writer, segment.Span);
-                }
-                else
-                {
-                    writer.Write(segment.Span);
-                    writer.Commit();
+                    if (_autoChunk)
+                    {
+                        CommitChunkInternal(ref writer, segment.Span);
+                    }
+                    else
+                    {
+                        writer.Write(segment.Span);
+                        writer.Commit();
+                    }
                 }
                 segment.Return();
             }
@@ -391,16 +394,19 @@ private void WriteDataWrittenBeforeHeaders(ref BufferWriter<PipeWriter> writer)
 
         if (!_currentSegment.IsEmpty)
         {
-            var segment = _currentSegment.Slice(0, _position);
-
-            if (_autoChunk)
+            if (canWriteBody)
             {
-                CommitChunkInternal(ref writer, segment.Span);
-            }
-            else
-            {
-                writer.Write(segment.Span);
-                writer.Commit();
+                var segment = _currentSegment.Slice(0, _position);
+
+                if (_autoChunk)
+                {
+                    CommitChunkInternal(ref writer, segment.Span);
+                }
+                else
+                {
+                    writer.Write(segment.Span);
+                    writer.Commit();
+                }
             }
 
             _position = 0;
@@ -505,7 +511,8 @@ public ValueTask<FlushResult> FirstWriteAsync(int statusCode, string? reasonPhra
             // Uses same BufferWriter to write response headers and response
             var writer = new BufferWriter<PipeWriter>(_pipeWriter);
 
-            WriteResponseHeadersInternal(ref writer, statusCode, reasonPhrase, responseHeaders, autoChunk);
+            // canWriteBody hardcoded to true as we already check if a body is allowed before calling this method
+            WriteResponseHeadersInternal(ref writer, statusCode, reasonPhrase, responseHeaders, autoChunk, canWriteBody: true);
 
             return WriteAsyncInternal(ref writer, buffer, cancellationToken);
         }
@@ -525,7 +532,8 @@ public ValueTask<FlushResult> FirstWriteChunkedAsync(int statusCode, string? rea
             // Uses same BufferWriter to write response headers and chunk
             var writer = new BufferWriter<PipeWriter>(_pipeWriter);
 
-            WriteResponseHeadersInternal(ref writer, statusCode, reasonPhrase, responseHeaders, autoChunk);
+            // canWriteBody hardcoded to true as we already check if a body is allowed before calling this method
+            WriteResponseHeadersInternal(ref writer, statusCode, reasonPhrase, responseHeaders, autoChunk, canWriteBody: true);
 
             CommitChunkInternal(ref writer, buffer);
 

src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs

@@ -999,7 +999,7 @@ private void ProduceStart(bool appCompleted)
 
         var responseHeaders = CreateResponseHeaders(appCompleted);
 
-        Output.WriteResponseHeaders(StatusCode, ReasonPhrase, responseHeaders, _autoChunk, appCompleted);
+        Output.WriteResponseHeaders(StatusCode, ReasonPhrase, responseHeaders, _autoChunk, appCompleted, canWriteBody: _canWriteResponseBody);
     }
 
     private void VerifyInitializeState(int firstWriteByteCount)
@@ -1645,7 +1645,7 @@ private ValueTask<FlushResult> FirstWriteAsyncInternal(ReadOnlyMemory<byte> data
             {
                 if (data.Length == 0)
                 {
-                    Output.WriteResponseHeaders(StatusCode, ReasonPhrase, responseHeaders, _autoChunk, appCompleted: false);
+                    Output.WriteResponseHeaders(StatusCode, ReasonPhrase, responseHeaders, _autoChunk, appCompleted: false, canWriteBody: _canWriteResponseBody);
                     return Output.FlushAsync(cancellationToken);
                 }
 
@@ -1659,7 +1659,7 @@ private ValueTask<FlushResult> FirstWriteAsyncInternal(ReadOnlyMemory<byte> data
         }
         else
         {
-            Output.WriteResponseHeaders(StatusCode, ReasonPhrase, responseHeaders, _autoChunk, appCompleted: false);
+            Output.WriteResponseHeaders(StatusCode, ReasonPhrase, responseHeaders, _autoChunk, appCompleted: false, canWriteBody: _canWriteResponseBody);
             HandleNonBodyResponseWrite();
             return Output.FlushAsync(cancellationToken);
         }

src/Servers/Kestrel/Core/src/Internal/Http/IHttpOutputProducer.cs

@@ -13,7 +13,7 @@ internal interface IHttpOutputProducer
     ValueTask<FlushResult> WriteChunkAsync(ReadOnlySpan<byte> data, CancellationToken cancellationToken);
     ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken);
     ValueTask<FlushResult> Write100ContinueAsync();
-    void WriteResponseHeaders(int statusCode, string? reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk, bool appCompleted);
+    void WriteResponseHeaders(int statusCode, string? reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk, bool appCompleted, bool canWriteBody);
     // This takes ReadOnlySpan instead of ReadOnlyMemory because it always synchronously copies data before flushing.
     ValueTask<FlushResult> WriteDataToPipeAsync(ReadOnlySpan<byte> data, CancellationToken cancellationToken);
     // Test hook

src/Servers/Kestrel/Core/src/Internal/Http2/Http2OutputProducer.cs

@@ -358,7 +358,8 @@ public ValueTask<FlushResult> Write100ContinueAsync()
         }
     }
 
-    public void WriteResponseHeaders(int statusCode, string? reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk, bool appCompleted)
+    // canWriteBody is ignored as we don't have chunked bodies in HTTP/2 and so writing headers is not affected by canWriteBody
+    public void WriteResponseHeaders(int statusCode, string? reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk, bool appCompleted, bool canWriteBody)
     {
         lock (_dataWriterLock)
         {
@@ -556,7 +557,8 @@ public ValueTask<FlushResult> FirstWriteAsync(int statusCode, string? reasonPhra
     {
         lock (_dataWriterLock)
         {
-            WriteResponseHeaders(statusCode, reasonPhrase, responseHeaders, autoChunk, appCompleted: false);
+            // canWriteBody is hardcoded to true as we don't have chunked bodies in HTTP/2 and so writing headers is not affected by canWriteBody
+            WriteResponseHeaders(statusCode, reasonPhrase, responseHeaders, autoChunk, appCompleted: false, canWriteBody: true);
 
             return WriteDataToPipeAsync(data, cancellationToken);
         }

src/Servers/Kestrel/Core/src/Internal/Http3/Http3OutputProducer.cs

@@ -153,7 +153,8 @@ public ValueTask<FlushResult> FirstWriteAsync(int statusCode, string? reasonPhra
     {
         lock (_dataWriterLock)
         {
-            WriteResponseHeaders(statusCode, reasonPhrase, responseHeaders, autoChunk, appCompleted: false);
+            // canWriteBody is hardcoded to true as we don't have chunked bodies in HTTP/3 and so writing headers is not affected by canWriteBody
+            WriteResponseHeaders(statusCode, reasonPhrase, responseHeaders, autoChunk, appCompleted: false, canWriteBody: true);
 
             return WriteDataToPipeAsync(data, cancellationToken);
         }
@@ -375,7 +376,8 @@ public ValueTask<FlushResult> WriteDataToPipeAsync(ReadOnlySpan<byte> data, Canc
         }
     }
 
-    public void WriteResponseHeaders(int statusCode, string? reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk, bool appCompleted)
+    // canWriteBody is ignored as we don't have chunked bodies in HTTP/3 and so writing headers is not affected by canWriteBody
+    public void WriteResponseHeaders(int statusCode, string? reasonPhrase, HttpResponseHeaders responseHeaders, bool autoChunk, bool appCompleted, bool canWriteBody)
     {
         // appCompleted flag is not used here. The write FIN is sent via the transport and not via the frame.
         // Headers are written to buffer and flushed with a FIN when Http3FrameWriter.CompleteAsync is called

src/Servers/Kestrel/test/InMemory.FunctionalTests/ResponseTests.cs

@@ -1494,6 +1494,44 @@ await connection.Receive(
         }
     }
 
+    [Fact]
+    public async Task HeadResponseBodyNotWrittenWithAdvanceBeforeFlush()
+    {
+        var flushed = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        var serviceContext = new TestServiceContext(LoggerFactory) { ServerOptions = { AllowSynchronousIO = true } };
+
+        await using (var server = new TestServer(async httpContext =>
+        {
+            var span = httpContext.Response.BodyWriter.GetSpan(5);
+            for (var i = 0; i < span.Length; i++)
+            {
+                span[i] = (byte)'h';
+            }
+            httpContext.Response.BodyWriter.Advance(span.Length);
+            await httpContext.Response.BodyWriter.FlushAsync();
+            await flushed.Task;
+        }, serviceContext))
+        {
+            using (var connection = server.CreateConnection())
+            {
+                await connection.Send(
+                    "HEAD / HTTP/1.1",
+                    "Host:",
+                    "",
+                    "");
+                await connection.Receive(
+                    "HTTP/1.1 200 OK",
+                    $"Date: {server.Context.DateHeaderValue}",
+                    "Transfer-Encoding: chunked",
+                    "",
+                    "");
+
+                flushed.SetResult();
+            }
+        }
+    }
+
     [Fact]
     public async Task ZeroLengthWritesFlushHeaders()
     {