Fix NRE when aborting connection or client disconects

src/Microsoft.AspNetCore.Server.Kestrel/Http/SocketInput.cs

@@ -27,9 +27,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Http
         private MemoryPoolBlock _tail;
         private MemoryPoolBlock _pinned;
 
-        private int _consumingState;
         private object _sync = new object();
 
+        private bool _consuming;
+        private bool _disposed;
+
         public SocketInput(MemoryPool memory, IThreadPool threadPool)
         {
             _memory = memory;
@@ -163,13 +165,16 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Http
 
         public MemoryPoolIterator ConsumingStart()
         {
-            if (Interlocked.CompareExchange(ref _consumingState, 1, 0) != 0)
+            lock (_sync)
+            {
+                if (_consuming)
                 {
                     throw new InvalidOperationException("Already consuming input.");
                 }
-
+                _consuming = true;
                 return new MemoryPoolIterator(_head);
             }
+        }
 
         public void ConsumingComplete(
             MemoryPoolIterator consumed,
@@ -179,6 +184,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Http
             MemoryPoolBlock returnEnd = null;
 
             lock (_sync)
+            {
+                if (!_disposed)
                 {
                     if (!consumed.IsDefault)
                     {
@@ -201,18 +208,22 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Http
                             _awaitableIsCompleted);
                     }
                 }
-
-            while (returnStart != returnEnd)
+                else
                 {
-                var returnBlock = returnStart;
-                returnStart = returnStart.Next;
-                returnBlock.Pool.Return(returnBlock);
+                    returnStart = _head;
+                    returnEnd = null;
+                    _head = null;
+                    _tail = null;
                 }
 
-            if (Interlocked.CompareExchange(ref _consumingState, 0, 1) != 1)
+                ReturnBlocks(returnStart, returnEnd);
+
+                if (!_consuming)
                 {
                     throw new InvalidOperationException("No ongoing consuming operation to complete.");
                 }
+                _consuming = false;
+            }
         }
 
         public void CompleteAwaiting()
@@ -285,21 +296,30 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Http
         }
 
         public void Dispose()
+        {
+            lock (_sync)
             {
                 AbortAwaiting();
 
-            // Return all blocks
-            var block = _head;
-            while (block != null)
+                if (!_consuming)
+                {
+                    ReturnBlocks(_head, null);
+                    _head = null;
+                    _tail = null;
+                }
+                _disposed = true;
+            }
+        }
+
+        private static void ReturnBlocks(MemoryPoolBlock block, MemoryPoolBlock end)
+        {
+            while (block != end)
             {
                 var returnBlock = block;
                 block = block.Next;
 
                 returnBlock.Pool.Return(returnBlock);
             }
-
-            _head = null;
-            _tail = null;
         }
     }
 }

src/Microsoft.AspNetCore.Server.Kestrel/Infrastructure/TaskUtilities.cs

@@ -29,6 +29,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Infrastructure
         public static Task<int> GetCancelledZeroTask(CancellationToken cancellationToken = default(CancellationToken))
         {
 #if NETSTANDARD1_3
+            // Make sure cancellationToken is cancelled before passing to Task.FromCanceled
+            if (!cancellationToken.IsCancellationRequested)
+            {
+                cancellationToken = new CancellationToken(true);
+            }
             return Task.FromCanceled<int>(cancellationToken);
 #else
             var tcs = new TaskCompletionSource<int>();

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

@@ -104,5 +104,15 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             var stream = new FrameRequestStream();
             await stream.FlushAsync();
         }
+
+        [Fact]
+        public void AbortCausesReadToCancel()
+        {
+            var stream = new FrameRequestStream();
+            stream.StartAcceptingReads(null);
+            stream.Abort();
+            var task = stream.ReadAsync(new byte[1], 0, 1);
+            Assert.True(task.IsCanceled);
+        }
     }
 }