Refactor non-body response handling.

- Add functional tests
- Remove BadHttpResponse, ResponseRejectionReasons and TestFrameProtectedMembers
- Chunk non-keepalive HTTP/1.1 responses
- Set _canHaveBody to true when StatusCode == 101
- Add test to check that upgraded connections are always closed when the app
  is done
- Log writes on responses to HEAD requests only once.

src/Microsoft.AspNetCore.Server.Kestrel/BadHttpResponse.cs

@@ -1,75 +0,0 @@
-// Copyright (c) .NET Foundation. All rights reserved.
-// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
-
-using System;
-using System.Runtime.CompilerServices;
-using Microsoft.AspNetCore.Server.Kestrel.Internal.Http;
-
-namespace Microsoft.AspNetCore.Server.Kestrel
-{
-    public static class BadHttpResponse
-    {
-        internal static void ThrowException(ResponseRejectionReasons reason)
-        {
-            throw GetException(reason);
-        }
-
-        internal static void ThrowException(ResponseRejectionReasons reason, int value)
-        {
-            throw GetException(reason, value.ToString());
-        }
-
-        internal static void ThrowException(ResponseRejectionReasons reason, ResponseRejectionParameter parameter)
-        {
-            throw GetException(reason, parameter.ToString());
-        }
-
-        internal static InvalidOperationException GetException(ResponseRejectionReasons reason, int value)
-        {
-            return GetException(reason, value.ToString());
-        }
-
-        [MethodImpl(MethodImplOptions.NoInlining)]
-        internal static InvalidOperationException GetException(ResponseRejectionReasons reason)
-        {
-            InvalidOperationException ex;
-            switch (reason)
-            {
-                case ResponseRejectionReasons.HeadersReadonlyResponseStarted:
-                    ex = new InvalidOperationException("Headers are read-only, response has already started.");
-                    break;
-                case ResponseRejectionReasons.OnStartingCannotBeSetResponseStarted:
-                    ex = new InvalidOperationException("OnStarting cannot be set, response has already started.");
-                    break;
-                default:
-                    ex = new InvalidOperationException("Bad response.");
-                    break;
-            }
-
-            return ex;
-        }
-
-        [MethodImpl(MethodImplOptions.NoInlining)]
-        private static InvalidOperationException GetException(ResponseRejectionReasons reason, string value)
-        {
-            InvalidOperationException ex;
-            switch (reason)
-            {
-                case ResponseRejectionReasons.ValueCannotBeSetResponseStarted:
-                    ex = new InvalidOperationException(value + " cannot be set, response had already started.");
-                    break;
-                case ResponseRejectionReasons.TransferEncodingSetOnNonBodyResponse:
-                    ex = new InvalidOperationException($"Transfer-Encoding set on a {value} non-body request.");
-                    break;
-                case ResponseRejectionReasons.WriteToNonBodyResponse:
-                    ex = new InvalidOperationException($"Write to non-body {value} response.");
-                    break;
-                default:
-                    ex = new InvalidOperationException("Bad response.");
-                    break;
-            }
-
-            return ex;
-        }
-    }
-}

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

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Net;
@@ -74,6 +75,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
         protected readonly long _keepAliveMilliseconds;
         private readonly long _requestHeadersTimeoutMilliseconds;
 
+        private int _responseBytesWritten;
+
         public Frame(ConnectionContext context)
         {
             ConnectionContext = context;
@@ -172,7 +175,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             {
                 if (HasResponseStarted)
                 {
-                    BadHttpResponse.ThrowException(ResponseRejectionReasons.ValueCannotBeSetResponseStarted, ResponseRejectionParameter.StatusCode);
+                    ThrowResponseAlreadyStartedException(nameof(StatusCode));
                 }
 
                 _statusCode = value;
@@ -190,7 +193,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             {
                 if (HasResponseStarted)
                 {
-                    BadHttpResponse.ThrowException(ResponseRejectionReasons.ValueCannotBeSetResponseStarted, ResponseRejectionParameter.ReasonPhrase);
+                    ThrowResponseAlreadyStartedException(nameof(ReasonPhrase));
                 }
 
                 _reasonPhrase = value;
@@ -346,6 +349,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
 
             _remainingRequestHeadersBytesAllowed = ServerOptions.Limits.MaxRequestHeadersTotalSize;
             _requestHeadersParsed = 0;
+
+            _responseBytesWritten = 0;
         }
 
         /// <summary>
@@ -426,7 +431,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             {
                 if (HasResponseStarted)
                 {
-                    BadHttpResponse.ThrowException(ResponseRejectionReasons.OnStartingCannotBeSetResponseStarted, ResponseRejectionParameter.OnStarting);
+                    ThrowResponseAlreadyStartedException(nameof(OnStarting));
                 }
 
                 if (_onStarting == null)
@@ -512,6 +517,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
         public void Write(ArraySegment<byte> data)
         {
             ProduceStartAndFireOnStarting().GetAwaiter().GetResult();
+            _responseBytesWritten += data.Count;
 
             if (_canHaveBody)
             {
@@ -530,7 +536,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             }
             else
             {
-                HandleNonBodyResponseWrite(data.Count);
+                HandleNonBodyResponseWrite();
             }
         }
 
@@ -541,6 +547,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
                 return WriteAsyncAwaited(data, cancellationToken);
             }
 
+            _responseBytesWritten += data.Count;
+
             if (_canHaveBody)
             {
                 if (_autoChunk)
@@ -558,7 +566,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             }
             else
             {
-                HandleNonBodyResponseWrite(data.Count);
+                HandleNonBodyResponseWrite();
                 return TaskCache.CompletedTask;
             }
         }
@@ -566,6 +574,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
         public async Task WriteAsyncAwaited(ArraySegment<byte> data, CancellationToken cancellationToken)
         {
             await ProduceStartAndFireOnStarting();
+            _responseBytesWritten += data.Count;
 
             if (_canHaveBody)
             {
@@ -584,10 +593,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             }
             else
             {
-                HandleNonBodyResponseWrite(data.Count);
+                HandleNonBodyResponseWrite();
                 return;
             }
-
         }
 
         private void WriteChunked(ArraySegment<byte> data)
@@ -700,11 +708,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
                     return TaskCache.CompletedTask;
                 }
 
-                // If the request was rejected, StatusCode has already been set by SetBadRequestState
+                // If the request was rejected, the error state has already been set by SetBadRequestState
                 if (!_requestRejected)
                 {
                     // 500 Internal Server Error
-                    ErrorResetHeadersToDefaults(statusCode: 500);
+                    SetErrorResponseHeaders(statusCode: 500);
                 }
             }
 
@@ -740,6 +748,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
                 ConnectionControl.End(ProduceEndType.ConnectionKeepAlive);
             }
 
+            if (HttpMethods.IsHead(Method) && _responseBytesWritten > 0)
+            {
+                Log.ConnectionHeadResponseBodyWrite(ConnectionId, _responseBytesWritten);
+            }
+
             return TaskCache.CompletedTask;
         }
 
@@ -758,24 +771,26 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             bool appCompleted)
         {
             var responseHeaders = FrameResponseHeaders;
-
             var hasConnection = responseHeaders.HasConnection;
 
-            // Set whether response can have body
-            _canHaveBody = StatusCanHaveBody(StatusCode) && Method != "HEAD";
-
             var end = SocketOutput.ProducingStart();
+
             if (_keepAlive && hasConnection)
             {
                 var connectionValue = responseHeaders.HeaderConnection.ToString();
                 _keepAlive = connectionValue.Equals("keep-alive", StringComparison.OrdinalIgnoreCase);
             }
-            
+
+            // Set whether response can have body
+            _canHaveBody = StatusCanHaveBody(StatusCode) && Method != "HEAD";
+
+            // Don't set the Content-Length or Transfer-Encoding headers
+            // automatically for HEAD requests or 204, 205, 304 responses.
             if (_canHaveBody)
             {
                 if (!responseHeaders.HasTransferEncoding && !responseHeaders.HasContentLength)
                 {
-                    if (appCompleted)
+                    if (appCompleted && StatusCode != 101)
                     {
                         // Since the app has completed and we are only now generating
                         // the headers we can safely set the Content-Length to 0.
@@ -790,7 +805,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
                         //
                         // A server MUST NOT send a response containing Transfer-Encoding unless the corresponding
                         // request indicates HTTP/1.1 (or later).
-                        if (_httpVersion == Http.HttpVersion.Http11)
+                        if (_httpVersion == Http.HttpVersion.Http11 && StatusCode != 101)
                         {
                             _autoChunk = true;
                             responseHeaders.SetRawTransferEncoding("chunked", _bytesTransferEncodingChunked);
@@ -804,8 +819,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             }
             else
             {
-                // Don't set the Content-Length or Transfer-Encoding headers
-                // automatically for HEAD requests or 101, 204, 205, 304 responses.
                 if (responseHeaders.HasTransferEncoding)
                 {
                     RejectNonBodyTransferEncodingResponse(appCompleted);
@@ -1271,15 +1284,19 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
         public bool StatusCanHaveBody(int statusCode)
         {
             // List of status codes taken from Microsoft.Net.Http.Server.Response
-            return statusCode != 101 &&
+            return statusCode != 204 &&
-                   statusCode != 204 &&
                    statusCode != 205 &&
                    statusCode != 304;
         }
 
+        private void ThrowResponseAlreadyStartedException(string value)
+        {
+            throw new InvalidOperationException($"{value} cannot be set, response has already started.");
+        }
+
         private void RejectNonBodyTransferEncodingResponse(bool appCompleted)
         {
-            var ex = BadHttpResponse.GetException(ResponseRejectionReasons.TransferEncodingSetOnNonBodyResponse, StatusCode);
+            var ex = new InvalidOperationException($"Transfer-Encoding set on a {StatusCode} non-body request.");
             if (!appCompleted)
             {
                 // Back out of header creation surface exeception in user code
@@ -1289,18 +1306,22 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             else
             {
                 ReportApplicationError(ex);
+
                 // 500 Internal Server Error
-                ErrorResetHeadersToDefaults(statusCode: 500);
+                SetErrorResponseHeaders(statusCode: 500);
             }
         }
 
-        private void ErrorResetHeadersToDefaults(int statusCode)
+        private void SetErrorResponseHeaders(int statusCode)
         {
-            // Setting status code will throw if response has already started
+            Debug.Assert(!HasResponseStarted, $"{nameof(SetErrorResponseHeaders)} called after response had already started.");
-            if (!HasResponseStarted)
+
+            StatusCode = statusCode;
+            ReasonPhrase = null;
+
+            if (FrameResponseHeaders == null)
             {
-                StatusCode = statusCode;
+                InitializeHeaders();
-                ReasonPhrase = null;
             }
 
             var responseHeaders = FrameResponseHeaders;
@@ -1316,17 +1337,13 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             }
         }
 
-        public void HandleNonBodyResponseWrite(int count)
+        public void HandleNonBodyResponseWrite()
         {
-            if (Method == "HEAD")
+            // Writes to HEAD response are ignored and logged at the end of the request
+            if (Method != "HEAD")
             {
-                // Don't write to body for HEAD requests.
+                // Throw Exception for 204, 205, 304 responses.
-                Log.ConnectionHeadResponseBodyWrite(ConnectionId, count);
+                throw new InvalidOperationException($"Write to non-body {StatusCode} response.");
-            }
-            else
-            {
-                // Throw Exception for 101, 204, 205, 304 responses.
-                BadHttpResponse.ThrowException(ResponseRejectionReasons.WriteToNonBodyResponse, StatusCode);
             }
         }
 
@@ -1365,7 +1382,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
 
         public void SetBadRequestState(BadHttpRequestException ex)
         {
-            ErrorResetHeadersToDefaults(statusCode: ex.StatusCode);
+            // Setting status code will throw if response has already started
+            if (!HasResponseStarted)
+            {
+                SetErrorResponseHeaders(statusCode: ex.StatusCode);
+            }
 
             _keepAlive = false;
             _requestProcessingStopping = true;

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

@@ -29,7 +29,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             {
                 if (_isReadOnly)
                 {
-                    BadHttpResponse.ThrowException(ResponseRejectionReasons.HeadersReadonlyResponseStarted);
+                    ThrowHeadersReadOnlyException();
                 }
                 SetValueFast(key, value);
             }
@@ -48,6 +48,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             }
         }
 
+        protected void ThrowHeadersReadOnlyException()
+        {
+            throw new InvalidOperationException("Headers are read-only, response has already started.");
+        }
+
         protected void ThrowArgumentException()
         {
             throw new ArgumentException();
@@ -139,7 +144,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
         {
             if (_isReadOnly)
             {
-                BadHttpResponse.ThrowException(ResponseRejectionReasons.HeadersReadonlyResponseStarted);
+                ThrowHeadersReadOnlyException();
             }
             AddValueFast(key, value);
         }
@@ -148,7 +153,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
         {
             if (_isReadOnly)
             {
-                BadHttpResponse.ThrowException(ResponseRejectionReasons.HeadersReadonlyResponseStarted);
+                ThrowHeadersReadOnlyException();
             }
             ClearFast();
         }
@@ -195,7 +200,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
         {
             if (_isReadOnly)
             {
-                BadHttpResponse.ThrowException(ResponseRejectionReasons.HeadersReadonlyResponseStarted);
+                ThrowHeadersReadOnlyException();
             }
             return RemoveFast(key);
         }

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

@@ -1,22 +0,0 @@
-// Copyright (c) .NET Foundation. All rights reserved.
-// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
-
-namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
-{
-    public enum ResponseRejectionReasons
-    {
-        HeadersReadonlyResponseStarted,
-        ValueCannotBeSetResponseStarted,
-        TransferEncodingSetOnNonBodyResponse,
-        WriteToNonBodyResponse,
-        OnStartingCannotBeSetResponseStarted
-    }
-
-    public enum ResponseRejectionParameter
-    {
-        StatusCode,
-        ReasonPhrase,
-        OnStarting
-    }
-    
-}

test/Microsoft.AspNetCore.Server.Kestrel.FunctionalTests/ResponseTests.cs

@@ -5,13 +5,17 @@ using System;
 using System.Linq;
 using System.Net;
 using System.Net.Http;
+using System.Text;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Server.Kestrel.Internal.Http;
+using Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure;
 using Microsoft.AspNetCore.Testing;
 using Microsoft.Extensions.Internal;
 using Microsoft.Extensions.Primitives;
+using Moq;
 using Xunit;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
@@ -213,6 +217,116 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             }
         }
 
+        [Fact]
+        public async Task TransferEncodingChunkedSetOnUnknownLengthHttp11Response()
+        {
+            using (var server = new TestServer(async httpContext =>
+            {
+                await httpContext.Response.WriteAsync("hello, ");
+                await httpContext.Response.WriteAsync("world");
+            }, new TestServiceContext()))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.Send(
+                        "GET / HTTP/1.1",
+                        "",
+                        "");
+                    await connection.Receive(
+                        "HTTP/1.1 200 OK",
+                        $"Date: {server.Context.DateHeaderValue}",
+                        "Transfer-Encoding: chunked",
+                        "",
+                        "7",
+                        "hello, ",
+                        "5",
+                        "world",
+                        "0",
+                        "",
+                        "");
+                }
+            }
+        }
+
+        [Theory]
+        [InlineData(204)]
+        [InlineData(205)]
+        [InlineData(304)]
+        public async Task TransferEncodingChunkedNotSetOnNonBodyResponse(int statusCode)
+        {
+            using (var server = new TestServer(httpContext =>
+            {
+                httpContext.Response.StatusCode = statusCode;
+                return TaskCache.CompletedTask;
+            }, new TestServiceContext()))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.Send(
+                        "GET / HTTP/1.1",
+                        "",
+                        "");
+                    await connection.Receive(
+                        $"HTTP/1.1 {Encoding.ASCII.GetString(ReasonPhrases.ToStatusBytes(statusCode))}",
+                        $"Date: {server.Context.DateHeaderValue}",
+                        "",
+                        "");
+                }
+            }
+        }
+
+        [Fact]
+        public async Task TransferEncodingNotSetOnHeadResponse()
+        {
+            using (var server = new TestServer(httpContext =>
+            {
+                return TaskCache.CompletedTask;
+            }, new TestServiceContext()))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.Send(
+                        "HEAD / HTTP/1.1",
+                        "",
+                        "");
+                    await connection.Receive(
+                        $"HTTP/1.1 200 OK",
+                        $"Date: {server.Context.DateHeaderValue}",
+                        "",
+                        "");
+                }
+            }
+        }
+
+        [Fact]
+        public async Task ResponseBodyNotWrittenOnHeadResponse()
+        {
+            var mockKestrelTrace = new Mock<IKestrelTrace>();
+
+            using (var server = new TestServer(async httpContext =>
+            {
+                await httpContext.Response.WriteAsync("hello, world");
+                await httpContext.Response.Body.FlushAsync();
+            }, new TestServiceContext { Log = mockKestrelTrace.Object }))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.Send(
+                        "HEAD / HTTP/1.1",
+                        "",
+                        "");
+                    await connection.Receive(
+                        $"HTTP/1.1 200 OK",
+                        $"Date: {server.Context.DateHeaderValue}",
+                        "",
+                        "");
+                }
+            }
+
+            mockKestrelTrace.Verify(kestrelTrace =>
+                kestrelTrace.ConnectionHeadResponseBodyWrite(It.IsAny<string>(), "hello, world".Length));
+        }
+
         public static TheoryData<string, StringValues, string> NullHeaderData
         {
             get

test/Microsoft.AspNetCore.Server.Kestrel.FunctionalTests/project.json

@@ -7,6 +7,7 @@
     "Microsoft.AspNetCore.Server.Kestrel.Https": "1.1.0-*",
     "Microsoft.AspNetCore.Testing": "1.1.0-*",
     "Microsoft.Extensions.Logging.Testing": "1.1.0-*",
+    "Moq": "4.6.36-*",
     "Newtonsoft.Json": "9.0.1",
     "xunit": "2.2.0-*"
   },

test/Microsoft.AspNetCore.Server.KestrelTests/ChunkedResponseTests.cs

@@ -5,6 +5,7 @@ using System;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Testing;
 using Xunit;
 
@@ -63,13 +64,12 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
 
         [Theory]
         [MemberData(nameof(ConnectionFilterData))]
-        public async Task ResponsesAreNotChunkedAutomaticallyForHttp10RequestsAndHttp11NonKeepAliveRequests(TestServiceContext testContext)
+        public async Task ResponsesAreNotChunkedAutomaticallyForHttp10Requests(TestServiceContext testContext)
         {
             using (var server = new TestServer(async httpContext =>
             {
-                var response = httpContext.Response;
+                await httpContext.Response.WriteAsync("Hello ");
-                await response.Body.WriteAsync(Encoding.ASCII.GetBytes("Hello "), 0, 6);
+                await httpContext.Response.WriteAsync("World!");
-                await response.Body.WriteAsync(Encoding.ASCII.GetBytes("World!"), 0, 6);
             }, testContext))
             {
                 using (var connection = server.CreateConnection())
@@ -86,7 +86,19 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                         "",
                         "Hello World!");
                 }
+            }
+        }
 
+        [Theory]
+        [MemberData(nameof(ConnectionFilterData))]
+        public async Task ResponsesAreChunkedAutomaticallyForHttp11NonKeepAliveRequests(TestServiceContext testContext)
+        {
+            using (var server = new TestServer(async httpContext =>
+            {
+                await httpContext.Response.WriteAsync("Hello ");
+                await httpContext.Response.WriteAsync("World!");
+            }, testContext))
+            {
                 using (var connection = server.CreateConnection())
                 {
                     await connection.SendEnd(
@@ -98,23 +110,28 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                         "HTTP/1.1 200 OK",
                         "Connection: close",
                         $"Date: {testContext.DateHeaderValue}",
+                        "Transfer-Encoding: chunked",
                         "",
-                        "Hello World!");
+                        "6",
+                        "Hello ",
+                        "6",
+                        "World!",
+                        "0",
+                        "",
+                        "");
                 }
             }
         }
 
-
         [Theory]
         [MemberData(nameof(ConnectionFilterData))]
-        public async Task SettingConnectionCloseHeaderInAppDisablesChunking(TestServiceContext testContext)
+        public async Task SettingConnectionCloseHeaderInAppDoesNotDisableChunking(TestServiceContext testContext)
         {
             using (var server = new TestServer(async httpContext =>
             {
-                var response = httpContext.Response;
+                httpContext.Response.Headers["Connection"] = "close";
-                response.Headers["Connection"] = "close";
+                await httpContext.Response.WriteAsync("Hello ");
-                await response.Body.WriteAsync(Encoding.ASCII.GetBytes("Hello "), 0, 6);
+                await httpContext.Response.WriteAsync("World!");
-                await response.Body.WriteAsync(Encoding.ASCII.GetBytes("World!"), 0, 6);
             }, testContext))
             {
                 using (var connection = server.CreateConnection())
@@ -127,8 +144,15 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                         "HTTP/1.1 200 OK",
                         "Connection: close",
                         $"Date: {testContext.DateHeaderValue}",
+                        "Transfer-Encoding: chunked",
                         "",
-                        "Hello World!");
+                        "6",
+                        "Hello ",
+                        "6",
+                        "World!",
+                        "0",
+                        "",
+                        "");
                 }
             }
         }

test/Microsoft.AspNetCore.Server.KestrelTests/EngineTests.cs

@@ -548,9 +548,6 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                         "POST / HTTP/1.1",
                         "Content-Length: 3",
                         "",
-                        "101POST / HTTP/1.1",
-                        "Content-Length: 3",
-                        "",
                         "204POST / HTTP/1.1",
                         "Content-Length: 3",
                         "",
@@ -562,9 +559,6 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                         "",
                         "200");
                     await connection.ReceiveEnd(
-                        "HTTP/1.1 101 Switching Protocols",
-                        $"Date: {testContext.DateHeaderValue}",
-                        "",
                         "HTTP/1.1 204 No Content",
                         $"Date: {testContext.DateHeaderValue}",
                         "",
@@ -583,6 +577,78 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             }
         }
 
+        [Theory]
+        [MemberData(nameof(ConnectionFilterData))]
+        public async Task ConnectionClosedAfter101Response(TestServiceContext testContext)
+        {
+            using (var server = new TestServer(async httpContext =>
+            {
+                var request = httpContext.Request;
+                var stream = await httpContext.Features.Get<IHttpUpgradeFeature>().UpgradeAsync();
+                var response = Encoding.ASCII.GetBytes("hello, world");
+                await stream.WriteAsync(response, 0, response.Length);
+            }, testContext))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.Send(
+                        "GET / HTTP/1.1",
+                        "",
+                        "");
+                    await connection.ReceiveEnd(
+                        "HTTP/1.1 101 Switching Protocols",
+                        "Connection: Upgrade",
+                        $"Date: {testContext.DateHeaderValue}",
+                        "",
+                        "hello, world");
+                }
+
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.Send(
+                        "GET / HTTP/1.0",
+                        "Connection: keep-alive",
+                        "",
+                        "");
+                    await connection.ReceiveEnd(
+                        "HTTP/1.1 101 Switching Protocols",
+                        "Connection: Upgrade",
+                        $"Date: {testContext.DateHeaderValue}",
+                        "",
+                        "hello, world");
+                }
+            }
+        }
+
+        [Theory]
+        [MemberData(nameof(ConnectionFilterData))]
+        public async Task WriteOnHeadResponseLoggedOnlyOnce(TestServiceContext testContext)
+        {
+            using (var server = new TestServer(async httpContext =>
+            {
+                await httpContext.Response.WriteAsync("hello, ");
+                await httpContext.Response.WriteAsync("world");
+                await httpContext.Response.WriteAsync("!");
+            }, testContext))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.SendEnd(
+                        "HEAD / HTTP/1.1",
+                        "",
+                        "");
+                    await connection.ReceiveEnd(
+                        "HTTP/1.1 200 OK",
+                        $"Date: {testContext.DateHeaderValue}",
+                        "",
+                        "");
+                }
+
+                Assert.Equal(1, ((TestKestrelTrace)testContext.Log).HeadResponseWrites);
+                Assert.Equal(13, ((TestKestrelTrace)testContext.Log).HeadResponseWriteByteCount);
+            }
+        }
+
         [Theory]
         [MemberData(nameof(ConnectionFilterData))]
         public async Task ThrowingResultsIn500Response(TestServiceContext testContext)

test/Microsoft.AspNetCore.Server.KestrelTests/FrameTests.cs

@@ -5,6 +5,7 @@ using System;
 using System.IO;
 using System.Text;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel;
 using Microsoft.AspNetCore.Server.Kestrel.Internal;
@@ -1263,7 +1264,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
         }
 
         [Fact]
-        public void FlushSetsTransferEncodingSetForUnknownLengthBodyResponse()
+        public void WriteThrowsForNonBodyResponse()
         {
             // Arrange
             var serviceContext = new ServiceContext
@@ -1280,114 +1281,40 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             {
                 SocketOutput = new MockSocketOuptut(),
             };
-            var frame = new TestFrameProtectedMembers<object>(application: null, context: connectionContext);
+            var frame = new Frame<object>(application: null, context: connectionContext);
             frame.InitializeHeaders();
-            frame.KeepAlive = true;
-            frame.HttpVersion = "HTTP/1.1";
-
-            // Act
-            frame.Flush();
-
-            // Assert
-            Assert.True(frame.HasResponseStarted);
-            Assert.True(frame.ResponseHeaders.ContainsKey("Transfer-Encoding"));
-        }
-
-        [Fact]
-        public void FlushDoesNotSetTransferEncodingSetForNoBodyResponse()
-        {
-            // Arrange
-            var serviceContext = new ServiceContext
-            {
-                DateHeaderValueManager = new DateHeaderValueManager(),
-                ServerOptions = new KestrelServerOptions(),
-                Log = new TestKestrelTrace()
-            };
-            var listenerContext = new ListenerContext(serviceContext)
-            {
-                ServerAddress = ServerAddress.FromUrl("http://localhost:5000")
-            };
-            var connectionContext = new ConnectionContext(listenerContext)
-            {
-                SocketOutput = new MockSocketOuptut(),
-            };
-            var frame = new TestFrameProtectedMembers<object>(application: null, context: connectionContext);
-            frame.InitializeHeaders();
-            frame.KeepAlive = true;
             frame.HttpVersion = "HTTP/1.1";
             ((IHttpResponseFeature)frame).StatusCode = 304;
 
-            // Act
+            // Act/Assert
-            frame.Flush();
-
-            // Assert
-            Assert.True(frame.HasResponseStarted);
-            Assert.False(frame.ResponseHeaders.ContainsKey("Transfer-Encoding"));
-        }
-
-        [Fact]
-        public void FlushDoesNotSetTransferEncodingSetForHeadResponse()
-        {
-            // Arrange
-            var serviceContext = new ServiceContext
-            {
-                DateHeaderValueManager = new DateHeaderValueManager(),
-                ServerOptions = new KestrelServerOptions(),
-                Log = new TestKestrelTrace()
-            };
-            var listenerContext = new ListenerContext(serviceContext)
-            {
-                ServerAddress = ServerAddress.FromUrl("http://localhost:5000")
-            };
-            var connectionContext = new ConnectionContext(listenerContext)
-            {
-                SocketOutput = new MockSocketOuptut(),
-            };
-            var frame = new TestFrameProtectedMembers<object>(application: null, context: connectionContext);
-            frame.InitializeHeaders();
-            frame.KeepAlive = true;
-            frame.HttpVersion = "HTTP/1.1";
-            ((IHttpRequestFeature)frame).Method = "HEAD";
-
-            // Act
-            frame.Flush();
-
-            // Assert
-            Assert.True(frame.HasResponseStarted);
-            Assert.False(frame.ResponseHeaders.ContainsKey("Transfer-Encoding"));
-        }
-
-        [Fact]
-        public void WriteThrowsForNoBodyResponse()
-        {
-            // Arrange
-            var serviceContext = new ServiceContext
-            {
-                DateHeaderValueManager = new DateHeaderValueManager(),
-                ServerOptions = new KestrelServerOptions(),
-                Log = new TestKestrelTrace()
-            };
-            var listenerContext = new ListenerContext(serviceContext)
-            {
-                ServerAddress = ServerAddress.FromUrl("http://localhost:5000")
-            };
-            var connectionContext = new ConnectionContext(listenerContext)
-            {
-                SocketOutput = new MockSocketOuptut(),
-            };
-            var frame = new TestFrameProtectedMembers<object>(application: null, context: connectionContext);
-            frame.InitializeHeaders();
-            frame.KeepAlive = true;
-            frame.HttpVersion = "HTTP/1.1";
-            ((IHttpResponseFeature)frame).StatusCode = 304;
-
-            // Assert
-            frame.Flush(); // Does not throw
-
             Assert.Throws<InvalidOperationException>(() => frame.Write(new ArraySegment<byte>(new byte[1])));
-            Assert.ThrowsAsync<InvalidOperationException>(() => frame.WriteAsync(new ArraySegment<byte>(new byte[1]), default(CancellationToken)));
+        }
 
-            frame.Flush(); // Does not throw
+        [Fact]
+        public async Task WriteAsyncThrowsForNonBodyResponse()
+        {
+            // Arrange
+            var serviceContext = new ServiceContext
+            {
+                DateHeaderValueManager = new DateHeaderValueManager(),
+                ServerOptions = new KestrelServerOptions(),
+                Log = new TestKestrelTrace()
+            };
+            var listenerContext = new ListenerContext(serviceContext)
+            {
+                ServerAddress = ServerAddress.FromUrl("http://localhost:5000")
+            };
+            var connectionContext = new ConnectionContext(listenerContext)
+            {
+                SocketOutput = new MockSocketOuptut(),
+            };
+            var frame = new Frame<object>(application: null, context: connectionContext);
+            frame.InitializeHeaders();
+            frame.HttpVersion = "HTTP/1.1";
+            ((IHttpResponseFeature)frame).StatusCode = 304;
+
+            // Act/Assert
+            await Assert.ThrowsAsync<InvalidOperationException>(() => frame.WriteAsync(new ArraySegment<byte>(new byte[1]), default(CancellationToken)));
         }
 
         [Fact]
@@ -1408,20 +1335,41 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             {
                 SocketOutput = new MockSocketOuptut(),
             };
-            var frame = new TestFrameProtectedMembers<object>(application: null, context: connectionContext);
+            var frame = new Frame<object>(application: null, context: connectionContext);
             frame.InitializeHeaders();
-            frame.KeepAlive = true;
             frame.HttpVersion = "HTTP/1.1";
             ((IHttpRequestFeature)frame).Method = "HEAD";
 
-            // Assert
+            // Act/Assert
-            frame.Flush(); // Does not throw
-
             frame.Write(new ArraySegment<byte>(new byte[1]));
-
-            frame.Flush(); // Does not throw
         }
 
+        [Fact]
+        public async Task WriteAsyncDoesNotThrowForHeadResponse()
+        {
+            // Arrange
+            var serviceContext = new ServiceContext
+            {
+                DateHeaderValueManager = new DateHeaderValueManager(),
+                ServerOptions = new KestrelServerOptions(),
+                Log = new TestKestrelTrace()
+            };
+            var listenerContext = new ListenerContext(serviceContext)
+            {
+                ServerAddress = ServerAddress.FromUrl("http://localhost:5000")
+            };
+            var connectionContext = new ConnectionContext(listenerContext)
+            {
+                SocketOutput = new MockSocketOuptut(),
+            };
+            var frame = new Frame<object>(application: null, context: connectionContext);
+            frame.InitializeHeaders();
+            frame.HttpVersion = "HTTP/1.1";
+            ((IHttpRequestFeature)frame).Method = "HEAD";
+
+            // Act/Assert
+            await frame.WriteAsync(new ArraySegment<byte>(new byte[1]), default(CancellationToken));
+        }
 
         [Fact]
         public void ManuallySettingTransferEncodingThrowsForHeadResponse()
@@ -1441,13 +1389,12 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             {
                 SocketOutput = new MockSocketOuptut(),
             };
-            var frame = new TestFrameProtectedMembers<object>(application: null, context: connectionContext);
+            var frame = new Frame<object>(application: null, context: connectionContext);
             frame.InitializeHeaders();
-            frame.KeepAlive = true;
             frame.HttpVersion = "HTTP/1.1";
             ((IHttpRequestFeature)frame).Method = "HEAD";
 
-            //Act
+            // Act
             frame.ResponseHeaders.Add("Transfer-Encoding", "chunked");
 
             // Assert
@@ -1472,13 +1419,12 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             {
                 SocketOutput = new MockSocketOuptut(),
             };
-            var frame = new TestFrameProtectedMembers<object>(application: null, context: connectionContext);
+            var frame = new Frame<object>(application: null, context: connectionContext);
             frame.InitializeHeaders();
-            frame.KeepAlive = true;
             frame.HttpVersion = "HTTP/1.1";
             ((IHttpResponseFeature)frame).StatusCode = 304;
 
-            //Act
+            // Act
             frame.ResponseHeaders.Add("Transfer-Encoding", "chunked");
 
             // Assert

test/Microsoft.AspNetCore.Server.KestrelTests/TestFrameProtectedMembers.cs

@@ -1,23 +0,0 @@
-// Copyright (c) .NET Foundation. All rights reserved.
-// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
-
-
-using Microsoft.AspNetCore.Hosting.Server;
-using Microsoft.AspNetCore.Server.Kestrel.Internal.Http;
-
-namespace Microsoft.AspNetCore.Server.KestrelTests
-{
-    public class TestFrameProtectedMembers<TContext> : Frame<TContext>
-    {
-        public TestFrameProtectedMembers(IHttpApplication<TContext> application, ConnectionContext context)
-            : base(application, context)
-        {
-        }
-
-        public bool KeepAlive
-        {
-            get { return _keepAlive; }
-            set { _keepAlive = value; }
-        }
-    }
-}

test/shared/TestKestrelTrace.cs

@@ -13,6 +13,10 @@ namespace Microsoft.AspNetCore.Testing
         {
         }
 
+        public int HeadResponseWrites { get; set; }
+
+        public int HeadResponseWriteByteCount { get; set; }
+
         public override void ConnectionRead(string connectionId, int count)
         {
             //_logger.LogDebug(1, @"Connection id ""{ConnectionId}"" recv {count} bytes.", connectionId, count);
@@ -27,5 +31,11 @@ namespace Microsoft.AspNetCore.Testing
         {
             //_logger.LogDebug(1, @"Connection id ""{ConnectionId}"" send finished with status {status}.", connectionId, status);
         }
+
+        public override void ConnectionHeadResponseBodyWrite(string connectionId, int count)
+        {
+            HeadResponseWrites++;
+            HeadResponseWriteByteCount = count;
+        }
     }
 }