Don't throw on FIN from upgraded connections (#2533)

* Allow app to drain request buffer after FIN

src/Kestrel.Core/Internal/Http/Http1MessageBody.cs

@@ -89,6 +89,14 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                         }
                         else if (result.IsCompleted)
                         {
+                            // Treat any FIN from an upgraded request as expected.
+                            // It's up to higher-level consumer (i.e. WebSocket middleware) to determine 
+                            // if the end is actually expected based on higher-level framing.
+                            if (RequestUpgrade)
+                            {
+                                break;
+                            }
+
                             BadHttpRequestException.Throw(RequestRejectionReason.UnexpectedEndOfRequestContent);
                         }
 

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

@@ -8,6 +8,7 @@ using System.IO;
 using System.Net;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Features;
@@ -282,7 +283,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
         void IHttpRequestLifetimeFeature.Abort()
         {
-            Abort(error: null);
+            Abort(new ConnectionAbortedException());
         }
     }
 }

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

@@ -411,21 +411,29 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
         }
 
         /// <summary>
-        /// Immediate kill the connection and poison the request and response streams.
+        /// Immediately kill the connection and poison the request and response streams with an error if there is one.
         /// </summary>
         public void Abort(Exception error)
         {
-            if (Interlocked.Exchange(ref _requestAborted, 1) == 0)
+            if (Interlocked.Exchange(ref _requestAborted, 1) != 0)
             {
-                _keepAlive = false;
+                return;
-
-                _streams?.Abort(error);
-
-                Output.Abort(error);
-
-                // Potentially calling user code. CancelRequestAbortedToken logs any exceptions.
-                ServiceContext.Scheduler.Schedule(state => ((HttpProtocol)state).CancelRequestAbortedToken(), this);
             }
+
+            _keepAlive = false;
+
+            // If Abort() isn't called with an exception, there was a FIN. In this case, even though the connection is
+            // still closed immediately, we allow the app to drain the data in the request buffer. If the request data
+            // was truncated, MessageBody will complete the RequestBodyPipe with an error.
+            if (error != null)
+            {
+                _streams?.Abort(error);
+            }
+
+            Output.Abort(error);
+
+            // Potentially calling user code. CancelRequestAbortedToken logs any exceptions.
+            ServiceContext.Scheduler.Schedule(state => ((HttpProtocol)state).CancelRequestAbortedToken(), this);
         }
 
         public void OnHeader(Span<byte> name, Span<byte> value)

test/Kestrel.Core.Tests/MessageBodyTests.cs

@@ -786,7 +786,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
 
                 input.Fin();
 
-                await Assert.ThrowsAsync<BadHttpRequestException>(async () => await body.ReadAsync(new ArraySegment<byte>(new byte[1])));
+                Assert.Equal(0, await body.ReadAsync(new ArraySegment<byte>(new byte[1])));
 
                 mockTimeoutControl.Verify(timeoutControl => timeoutControl.StartTimingReads(), Times.Never);
                 mockTimeoutControl.Verify(timeoutControl => timeoutControl.StopTimingReads(), Times.Never);

test/Kestrel.FunctionalTests/UpgradeTests.cs

@@ -297,5 +297,42 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             var exception = await Assert.ThrowsAsync<InvalidOperationException>(async () => await upgradeTcs.Task.TimeoutAfter(TimeSpan.FromSeconds(60)));
             Assert.Equal(CoreStrings.UpgradedConnectionLimitReached, exception.Message);
         }
+
+        [Fact]
+        public async Task DoesNotThrowOnFin()
+        {
+            var appCompletedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            using (var server = new TestServer(async context =>
+            {
+                var feature = context.Features.Get<IHttpUpgradeFeature>();
+                var duplexStream = await feature.UpgradeAsync();
+
+                try
+                {
+                    await duplexStream.CopyToAsync(Stream.Null);
+                    appCompletedTcs.SetResult(null);
+                }
+                catch (Exception ex)
+                {
+                    appCompletedTcs.SetException(ex);
+                    throw;
+                }
+
+            }, new TestServiceContext(LoggerFactory)))
+            {
+                using (var connection = server.CreateConnection())
+                {
+                    await connection.SendEmptyGetWithUpgrade();
+                    await connection.Receive("HTTP/1.1 101 Switching Protocols",
+                        "Connection: Upgrade",
+                        $"Date: {server.Context.DateHeaderValue}",
+                        "",
+                        "");
+                }
+
+                await appCompletedTcs.Task.TimeoutAfter(TestConstants.DefaultTimeout);
+            }
+        }
     }
 }