From 0622513058eff9a4374a0edca726d2c62959ea1e Mon Sep 17 00:00:00 2001
From: Stephen Halter <halter73@gmail.com>
Date: Tue, 15 Jan 2019 14:54:10 -0800
Subject: [PATCH] Wait to dispose RequestAborted CTS (#4447)

---
 eng/PatchConfig.props                         |  1 +
 .../Core/src/Internal/Http/HttpProtocol.cs    | 36 +++++++++----------
 .../Kestrel/Core/test/Http1ConnectionTests.cs | 25 ++++++++++++-
 3 files changed, 41 insertions(+), 21 deletions(-)

diff --git a/eng/PatchConfig.props b/eng/PatchConfig.props
index cb12734457..d9ee2e8f4e 100644
--- a/eng/PatchConfig.props
+++ b/eng/PatchConfig.props
@@ -33,6 +33,7 @@ Later on, this will be checked using this condition:
       Microsoft.AspNetCore.Http;
       Microsoft.AspNetCore.Mvc.Core;
       Microsoft.AspNetCore.Server.IIS;
+      Microsoft.AspNetCore.Server.Kestrel.Core;
       java:signalr;
     </PackagesInPatch>
   </PropertyGroup>
diff --git a/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs b/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs
index 3309701e0a..cf7ca09e9a 100644
--- a/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs
+++ b/src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs
@@ -348,11 +348,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
             // Lock to prevent CancelRequestAbortedToken from attempting to cancel an disposed CTS.
             lock (_abortLock)
             {
-                if (!_requestAborted)
-                {
-                    _abortedCts?.Dispose();
-                    _abortedCts = null;
-                }
+                _abortedCts?.Dispose();
+                _abortedCts = null;
             }
 
             _requestHeadersParsed = 0;
@@ -394,15 +391,16 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
         private void CancelRequestAbortedToken()
         {
-            try
+            lock (_abortLock)
             {
-                _abortedCts.Cancel();
-                _abortedCts.Dispose();
-                _abortedCts = null;
-            }
-            catch (Exception ex)
-            {
-                Log.ApplicationError(ConnectionId, TraceIdentifier, ex);
+                try
+                {
+                    _abortedCts?.Cancel();
+                }
+                catch (Exception ex)
+                {
+                    Log.ApplicationError(ConnectionId, TraceIdentifier, ex);
+                }
             }
         }
 
@@ -416,12 +414,12 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                 }
 
                 _requestAborted = true;
-            }
 
-            if (_abortedCts != null)
-            {
-                // Potentially calling user code. CancelRequestAbortedToken logs any exceptions.
-                ServiceContext.Scheduler.Schedule(state => ((HttpProtocol)state).CancelRequestAbortedToken(), this);
+                if (_abortedCts != null && !_preventRequestAbortedCancellation)
+                {
+                    // Potentially calling user code. CancelRequestAbortedToken logs any exceptions.
+                    ServiceContext.Scheduler.Schedule(state => ((HttpProtocol)state).CancelRequestAbortedToken(), this);
+                }
             }
         }
 
@@ -441,8 +439,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                 }
 
                 _preventRequestAbortedCancellation = true;
-                _abortedCts?.Dispose();
-                _abortedCts = null;
             }
         }
 
diff --git a/src/Servers/Kestrel/Core/test/Http1ConnectionTests.cs b/src/Servers/Kestrel/Core/test/Http1ConnectionTests.cs
index d60303d625..70301631ed 100644
--- a/src/Servers/Kestrel/Core/test/Http1ConnectionTests.cs
+++ b/src/Servers/Kestrel/Core/test/Http1ConnectionTests.cs
@@ -53,7 +53,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             var connectionFeatures = new FeatureCollection();
             connectionFeatures.Set(Mock.Of<IConnectionLifetimeFeature>());
 
-            _serviceContext = new TestServiceContext();
+            _serviceContext = new TestServiceContext()
+            {
+                Scheduler = PipeScheduler.Inline
+            };
+
             _timeoutControl = new Mock<ITimeoutControl>();
             _http1ConnectionContext = new HttpConnectionContext
             {
@@ -724,6 +728,25 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             Assert.False(_http1Connection.RequestAborted.IsCancellationRequested);
         }
 
+        [Fact]
+        public void RequestAbortedTokenIsFullyUsableAfterCancellation()
+        {
+            var originalToken = _http1Connection.RequestAborted;
+            var originalRegistration = originalToken.Register(() => { });
+
+            _http1Connection.Abort(new ConnectionAbortedException());
+
+            Assert.True(originalToken.WaitHandle.WaitOne(TestConstants.DefaultTimeout));
+            Assert.True(_http1Connection.RequestAborted.WaitHandle.WaitOne(TestConstants.DefaultTimeout));
+
+#if NETCOREAPP2_2
+            Assert.Equal(originalToken, originalRegistration.Token);
+#elif NET461
+#else
+#error Target framework needs to be updated
+#endif
+        }
+
         [Fact]
         public async Task ExceptionDetailNotIncludedWhenLogLevelInformationNotEnabled()
         {