Don't write body for non-body responses

Author: benaadams

src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/Frame.cs

@@ -58,6 +58,7 @@ public abstract partial class Frame : ConnectionContext, IFrameControl
 
         protected RequestProcessingStatus _requestProcessingStatus;
         protected bool _keepAlive;
+        private bool _canHaveBody;
         private bool _autoChunk;
         protected Exception _applicationException;
 
@@ -465,6 +466,9 @@ public void Write(ArraySegment<byte> data)
         {
             ProduceStartAndFireOnStarting().GetAwaiter().GetResult();
 
+            // Don't write to body for HEAD requests or 101, 204, 205, 304 responses.
+            if (!_canHaveBody) return;
+
             if (_autoChunk)
             {
                 if (data.Count == 0)
@@ -486,6 +490,9 @@ public Task WriteAsync(ArraySegment<byte> data, CancellationToken cancellationTo
                 return WriteAsyncAwaited(data, cancellationToken);
             }
 
+            // Don't write to body for HEAD requests or 101, 204, 205, 304 responses.
+            if (!_canHaveBody) return TaskUtilities.CompletedTask;
+
             if (_autoChunk)
             {
                 if (data.Count == 0)
@@ -710,10 +717,12 @@ private void CreateResponseHeader(
             bool appCompleted)
         {
             var responseHeaders = FrameResponseHeaders;
-            responseHeaders.SetReadOnly();
 
             var hasConnection = responseHeaders.HasConnection;
 
+            // Set whether response can have body
+            _canHaveBody = StatusCanHaveBody(StatusCode) && Method != "HEAD";
+
             var end = SocketOutput.ProducingStart();
             if (_keepAlive && hasConnection)
             {
@@ -727,11 +736,9 @@ private void CreateResponseHeader(
                 }
             }
 
-            if (_keepAlive && !responseHeaders.HasTransferEncoding && !responseHeaders.HasContentLength)
+            if (_canHaveBody)
             {
-                // Don't set the Content-Length or Transfer-Encoding headers
-                // automatically for HEAD requests or 101, 204, 205, 304 responses.
-                if (Method != "HEAD" && StatusCanHaveBody(StatusCode))
+                if (_keepAlive && !responseHeaders.HasTransferEncoding && !responseHeaders.HasContentLength)
                 {
                     if (appCompleted)
                     {
@@ -753,15 +760,26 @@ private void CreateResponseHeader(
                             _autoChunk = true;
                             responseHeaders.SetRawTransferEncoding("chunked", _bytesTransferEncodingChunked);
                         }
+                        else
+                        {
+                            _keepAlive = false;
+                        }
                     }
                 }
-
-                if (!appCompleted && _httpVersion != HttpVersionType.Http11)
+            }
+            else
+            {
+                // Don't set the Content-Length or Transfer-Encoding headers
+                // automatically for HEAD requests or 101, 204, 205, 304 responses.
+                if (responseHeaders.HasTransferEncoding)
                 {
-                    _keepAlive = false;
+                    // Remove Transfer-Encoding encoding header if it exists
+                    responseHeaders.Remove("Transfer-Encoding");
                 }
             }
 
+            responseHeaders.SetReadOnly();
+
             if (!_keepAlive && !hasConnection && _httpVersion != HttpVersionType.Http10)
             {
                 responseHeaders.SetRawConnection("close", _bytesConnectionClose);

src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/FrameHeaders.cs

@@ -196,7 +196,7 @@ bool ICollection<KeyValuePair<string, StringValues>>.Remove(KeyValuePair<string,
                 RemoveFast(item.Key);
         }
 
-        bool IDictionary<string, StringValues>.Remove(string key)
+        public bool Remove(string key)
         {
             if (_isReadOnly)
             {