fb

Author: BrennanConroy

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

@@ -360,6 +360,7 @@ private void WriteResponseHeadersInternal(ref BufferWriter<PipeWriter> writer, i
 
         writer.Commit();
 
+        Debug.Assert(responseBodyMode != ResponseBodyMode.Uninitialized);
         _responseBodyMode = responseBodyMode;
         WriteDataWrittenBeforeHeaders(ref writer);
         _unflushedBytes += writer.BytesCommitted;
@@ -544,7 +545,7 @@ public void Reset()
         Debug.Assert(_completedSegments == null || _completedSegments.Count == 0);
         // Cleared in sequential address ascending order
         _currentMemoryPrefixBytes = 0;
-        _responseBodyMode = ResponseBodyMode.ContentLength;
+        _responseBodyMode = ResponseBodyMode.Uninitialized;
         _writeStreamSuffixCalled = false;
         _currentChunkMemoryUpdated = false;
         _startCalled = false;

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

@@ -56,7 +56,7 @@ internal abstract partial class HttpProtocol : IHttpResponseControl
     protected volatile bool _keepAlive = true;
     // _responseBodyMode is set in CreateResponseHeaders.
     // If we are writing with GetMemory/Advance before calling StartAsync, assume we can write and throw away contents if we can't.
-    private ResponseBodyMode _responseBodyMode = ResponseBodyMode.ContentLength;
+    private ResponseBodyMode _responseBodyMode = ResponseBodyMode.Uninitialized;
     private bool _hasAdvanced;
     private bool _isLeasedMemoryInvalid = true;
     protected Exception? _applicationException;
@@ -346,7 +346,7 @@ public void Reset()
         _routeValues?.Clear();
 
         _requestProcessingStatus = RequestProcessingStatus.RequestPending;
-        _responseBodyMode = ResponseBodyMode.ContentLength;
+        _responseBodyMode = ResponseBodyMode.Uninitialized;
         _applicationException = null;
         _requestRejectedException = null;
 
@@ -1592,28 +1592,28 @@ public ValueTask<FlushResult> WritePipeAsync(ReadOnlyMemory<byte> data, Cancella
             VerifyAndUpdateWrite(data.Length);
         }
 
-        if (_responseBodyMode != ResponseBodyMode.Disabled)
+        switch (_responseBodyMode)
         {
-            if (_responseBodyMode == ResponseBodyMode.Chunked)
-            {
+            case ResponseBodyMode.Disabled:
+                HandleNonBodyResponseWrite();
+                return default;
+            case ResponseBodyMode.Chunked:
                 if (data.Length == 0)
                 {
                     return default;
                 }
 
                 return Output.WriteChunkAsync(data.Span, cancellationToken);
-            }
-            else
-            {
+            case ResponseBodyMode.ContentLength:
                 CheckLastWrite();
                 return Output.WriteDataToPipeAsync(data.Span, cancellationToken: cancellationToken);
-            }
-        }
-        else
-        {
-            HandleNonBodyResponseWrite();
-            return default;
+            case ResponseBodyMode.Uninitialized:
+                ThrowInvalidOperation();
+                break;
         }
+
+        Debug.Assert(false, "Should not reach here, all cases in above switch statement should return");
+        return default;
     }
 
     private ValueTask<FlushResult> FirstWriteAsync(ReadOnlyMemory<byte> data, CancellationToken cancellationToken)
@@ -1640,30 +1640,30 @@ private ValueTask<FlushResult> FirstWriteAsyncInternal(ReadOnlyMemory<byte> data
     {
         var responseHeaders = InitializeResponseFirstWrite(data.Length);
 
-        if (_responseBodyMode != ResponseBodyMode.Disabled)
+        switch (_responseBodyMode)
         {
-            if (_responseBodyMode == ResponseBodyMode.Chunked)
-            {
+            case ResponseBodyMode.Disabled:
+                Output.WriteResponseHeaders(StatusCode, ReasonPhrase, responseHeaders, _responseBodyMode, appCompleted: false);
+                HandleNonBodyResponseWrite();
+                return Output.FlushAsync(cancellationToken);
+            case ResponseBodyMode.Chunked:
                 if (data.Length == 0)
                 {
                     Output.WriteResponseHeaders(StatusCode, ReasonPhrase, responseHeaders, _responseBodyMode, appCompleted: false);
                     return Output.FlushAsync(cancellationToken);
                 }
 
                 return Output.FirstWriteChunkedAsync(StatusCode, ReasonPhrase, responseHeaders, _responseBodyMode, data.Span, cancellationToken);
-            }
-            else
-            {
+            case ResponseBodyMode.ContentLength:
                 CheckLastWrite();
                 return Output.FirstWriteAsync(StatusCode, ReasonPhrase, responseHeaders, _responseBodyMode, data.Span, cancellationToken);
-            }
-        }
-        else
-        {
-            Output.WriteResponseHeaders(StatusCode, ReasonPhrase, responseHeaders, _responseBodyMode, appCompleted: false);
-            HandleNonBodyResponseWrite();
-            return Output.FlushAsync(cancellationToken);
+            case ResponseBodyMode.Uninitialized:
+                ThrowInvalidOperation();
+                break;
         }
+
+        Debug.Assert(false, "Should not reach here, all cases in above switch statement should return");
+        return default;
     }
 
     public Task FlushAsync(CancellationToken cancellationToken = default)
@@ -1689,27 +1689,34 @@ public async ValueTask<FlushResult> WriteAsyncAwaited(Task initializeTask, ReadO
 
         // WriteAsyncAwaited is only called for the first write to the body.
         // Ensure headers are flushed if Write(Chunked)Async isn't called.
-        if (_responseBodyMode != ResponseBodyMode.Disabled)
+        switch (_responseBodyMode)
         {
-            if (_responseBodyMode == ResponseBodyMode.Chunked)
-            {
+            case ResponseBodyMode.Disabled:
+                HandleNonBodyResponseWrite();
+                return await Output.FlushAsync(cancellationToken);
+            case ResponseBodyMode.Chunked:
                 if (data.Length == 0)
                 {
                     return await Output.FlushAsync(cancellationToken);
                 }
 
                 return await Output.WriteChunkAsync(data.Span, cancellationToken);
-            }
-            else
-            {
+            case ResponseBodyMode.ContentLength:
                 CheckLastWrite();
                 return await Output.WriteDataToPipeAsync(data.Span, cancellationToken: cancellationToken);
-            }
-        }
-        else
-        {
-            HandleNonBodyResponseWrite();
-            return await Output.FlushAsync(cancellationToken);
+            case ResponseBodyMode.Uninitialized:
+                ThrowInvalidOperation();
+                break;
         }
+
+        Debug.Assert(false, "Should not reach here, all cases in above switch statement should return");
+        return default;
+    }
+
+    [DoesNotReturn]
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static void ThrowInvalidOperation()
+    {
+        throw new InvalidOperationException();
     }
 }

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

@@ -32,6 +32,7 @@ internal interface IHttpOutputProducer
 
 internal enum ResponseBodyMode
 {
+    Uninitialized,
     Disabled,
     Chunked,
     ContentLength