Abort request when returning a completed FlushResult (#22086)

src/Servers/Kestrel/Core/src/Internal/Http/Http1Connection.cs

@@ -6,7 +6,6 @@ using System.Buffers;
 using System.Diagnostics;
 using System.Globalization;
 using System.IO.Pipelines;
-using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Http.Features;
@@ -14,7 +13,7 @@ using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 {
-    internal partial class Http1Connection : HttpProtocol, IRequestProcessor
+    internal partial class Http1Connection : HttpProtocol, IRequestProcessor, IHttpOutputAborter
     {
         private const byte ByteAsterisk = (byte)'*';
         private const byte ByteForwardSlash = (byte)'/';
@@ -27,7 +26,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
         protected readonly long _keepAliveTicks;
         private readonly long _requestHeadersTimeoutTicks;
 
-        private int _requestAborted;
         private volatile bool _requestTimedOut;
         private uint _requestCount;
 
@@ -56,10 +54,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                 _context.Transport.Output,
                 _context.ConnectionId,
                 _context.ConnectionContext,
+                _context.MemoryPool,
                 _context.ServiceContext.Log,
                 _context.TimeoutControl,
-                this,
+                minResponseDataRateFeature: this,
-                _context.MemoryPool);
+                outputAborter: this);
 
             Input = _context.Transport.Input;
             Output = _http1Output;
@@ -94,7 +93,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
         public void OnInputOrOutputCompleted()
         {
             _http1Output.Abort(new ConnectionAbortedException(CoreStrings.ConnectionAbortedByClient));
-            AbortRequest();
+            CancelRequestAbortedToken();
         }
 
         /// <summary>
@@ -102,16 +101,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
         /// </summary>
         public void Abort(ConnectionAbortedException abortReason)
         {
-            if (Interlocked.Exchange(ref _requestAborted, 1) != 0)
-            {
-                return;
-            }
-
             _http1Output.Abort(abortReason);
-
+            CancelRequestAbortedToken();
-            AbortRequest();
+            PoisonBody(abortReason);
-
-            PoisonRequestBodyStream(abortReason);
         }
 
         protected override void ApplicationAbort()

src/Servers/Kestrel/Core/src/Internal/Http/Http1OutputProducer.cs

@@ -16,7 +16,7 @@ using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure.PipeWrite
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 {
-    internal class Http1OutputProducer : IHttpOutputProducer, IHttpOutputAborter, IDisposable
+    internal class Http1OutputProducer : IHttpOutputProducer, IDisposable
     {
         // Use C#7.3's ReadOnlySpan<byte> optimization for static data https://vcsjones.com/2019/02/01/csharp-readonly-span-bytes-static/
         // "HTTP/1.1 100 Continue\r\n\r\n"
@@ -33,10 +33,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
         private readonly string _connectionId;
         private readonly ConnectionContext _connectionContext;
+        private readonly MemoryPool<byte> _memoryPool;
         private readonly IKestrelTrace _log;
         private readonly IHttpMinResponseDataRateFeature _minResponseDataRateFeature;
+        private readonly IHttpOutputAborter _outputAborter;
         private readonly TimingPipeFlusher _flusher;
-        private readonly MemoryPool<byte> _memoryPool;
 
         // This locks access to all of the below fields
         private readonly object _contextLock = new object();
@@ -77,19 +78,22 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
             PipeWriter pipeWriter,
             string connectionId,
             ConnectionContext connectionContext,
+            MemoryPool<byte> memoryPool,
             IKestrelTrace log,
             ITimeoutControl timeoutControl,
             IHttpMinResponseDataRateFeature minResponseDataRateFeature,
-            MemoryPool<byte> memoryPool)
+            IHttpOutputAborter outputAborter)
         {
             // Allow appending more data to the PipeWriter when a flush is pending.
             _pipeWriter = new ConcurrentPipeWriter(pipeWriter, memoryPool, _contextLock);
             _connectionId = connectionId;
             _connectionContext = connectionContext;
+            _memoryPool = memoryPool;
             _log = log;
             _minResponseDataRateFeature = minResponseDataRateFeature;
+            _outputAborter = outputAborter;
+
             _flusher = new TimingPipeFlusher(_pipeWriter, timeoutControl, log);
-            _memoryPool = memoryPool;
         }
 
         public Task WriteDataAsync(ReadOnlySpan<byte> buffer, CancellationToken cancellationToken = default)
@@ -168,7 +172,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                 var bytesWritten = _unflushedBytes;
                 _unflushedBytes = 0;
 
-                return _flusher.FlushAsync(_minResponseDataRateFeature.MinDataRate, bytesWritten, this, cancellationToken);
+                return _flusher.FlushAsync(_minResponseDataRateFeature.MinDataRate, bytesWritten, _outputAborter, cancellationToken);
             }
 
             static ValueTask<FlushResult> FlushAsyncChunked(Http1OutputProducer producer, CancellationToken token)
@@ -188,7 +192,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                 // If there is an empty write, we still need to update the current chunk
                 producer._currentChunkMemoryUpdated = false;
 
-                return producer._flusher.FlushAsync(producer._minResponseDataRateFeature.MinDataRate, bytesWritten, producer, token);
+                return producer._flusher.FlushAsync(producer._minResponseDataRateFeature.MinDataRate, bytesWritten, producer._outputAborter, token);
             }
         }
 
@@ -585,7 +589,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
             return _flusher.FlushAsync(
                 _minResponseDataRateFeature.MinDataRate,
                 bytesWritten,
-                this,
+                _outputAborter,
                 cancellationToken);
         }
 

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

@@ -445,7 +445,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
         protected abstract bool TryParseRequest(ReadResult result, out bool endConnection);
 
-        private void CancelRequestAbortedToken()
+        private void CancelRequestAbortedTokenCallback()
         {
             try
             {
@@ -470,7 +470,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
             }
         }
 
-        protected void AbortRequest()
+        protected void CancelRequestAbortedToken()
         {
             var shouldScheduleCancellation = false;
 
@@ -488,11 +488,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
             if (shouldScheduleCancellation)
             {
                 // Potentially calling user code. CancelRequestAbortedToken logs any exceptions.
-                ServiceContext.Scheduler.Schedule(state => ((HttpProtocol)state).CancelRequestAbortedToken(), this);
+                ServiceContext.Scheduler.Schedule(state => ((HttpProtocol)state).CancelRequestAbortedTokenCallback(), this);
             }
         }
 
-        protected void PoisonRequestBodyStream(Exception abortReason)
+        protected void PoisonBody(Exception abortReason)
         {
             _bodyControl?.Abort(abortReason);
         }

src/Servers/Kestrel/Core/src/Internal/Http/HttpRequestPipeReader.cs

@@ -92,7 +92,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
             if (_state != HttpStreamState.Closed)
             {
                 _state = HttpStreamState.Aborted;
-                if (error != null)
+
+                if (error is object && _error is null)
                 {
                     _error = ExceptionDispatchInfo.Capture(error);
                 }
@@ -113,7 +114,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
             }
             else
             {
-                if (_error != null)
+                if (_error is object)
                 {
                     _error.Throw();
                 }

src/Servers/Kestrel/Core/src/Internal/Http/IHttpOutputAborter.cs

@@ -8,5 +8,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
     internal interface IHttpOutputAborter
     {
         void Abort(ConnectionAbortedException abortReason);
+        void OnInputOrOutputCompleted();
     }
 }

src/Servers/Kestrel/Core/src/Internal/Http2/Http2OutputProducer.cs

@@ -129,6 +129,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
             _stream.ResetAndAbort(abortReason, Http2ErrorCode.INTERNAL_ERROR);
         }
 
+        void IHttpOutputAborter.OnInputOrOutputCompleted()
+        {
+            _stream.ResetAndAbort(new ConnectionAbortedException($"{nameof(Http2OutputProducer)}.{nameof(ProcessDataWrites)} has completed."), Http2ErrorCode.INTERNAL_ERROR);
+        }
+
         public ValueTask<FlushResult> FlushAsync(CancellationToken cancellationToken)
         {
             if (cancellationToken.IsCancellationRequested)

src/Servers/Kestrel/Core/src/Internal/Http2/Http2Stream.FeatureCollection.cs

@@ -58,7 +58,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
         void IHttpResetFeature.Reset(int errorCode)
         {
             var abortReason = new ConnectionAbortedException(CoreStrings.FormatHttp2StreamResetByApplication((Http2ErrorCode)errorCode));
-            ResetAndAbort(abortReason, (Http2ErrorCode)errorCode);
+            ApplicationAbort(abortReason, (Http2ErrorCode)errorCode);
         }
     }
 }

src/Servers/Kestrel/Core/src/Internal/Http2/Http2Stream.cs

@@ -517,10 +517,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
             ResetAndAbort(abortReason, Http2ErrorCode.INTERNAL_ERROR);
         }
 
-        protected override void ApplicationAbort()
+        protected override void ApplicationAbort() => ApplicationAbort(new ConnectionAbortedException(CoreStrings.ConnectionAbortedByApplication), Http2ErrorCode.INTERNAL_ERROR);
+
+        private void ApplicationAbort(ConnectionAbortedException abortReason, Http2ErrorCode error)
         {
-            var abortReason = new ConnectionAbortedException(CoreStrings.ConnectionAbortedByApplication);
+            ResetAndAbort(abortReason, error);
-            ResetAndAbort(abortReason, Http2ErrorCode.INTERNAL_ERROR);
         }
 
         internal void ResetAndAbort(ConnectionAbortedException abortReason, Http2ErrorCode error)
@@ -548,10 +549,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
             // ensure that an app that completes early due to the abort doesn't result in header frames being sent.
             _http2Output.Stop();
 
-            AbortRequest();
+            CancelRequestAbortedToken();
 
             // Unblock the request body.
-            PoisonRequestBodyStream(abortReason);
+            PoisonBody(abortReason);
             RequestBodyPipe.Writer.Complete(abortReason);
 
             _inputFlowControl.Abort();

src/Servers/Kestrel/Core/src/Internal/Http3/Http3OutputProducer.cs

@@ -80,6 +80,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http3
             _stream.Abort(abortReason, Http3ErrorCode.InternalError);
         }
 
+        void IHttpOutputAborter.OnInputOrOutputCompleted()
+        {
+            _stream.Abort(new ConnectionAbortedException($"{nameof(Http3OutputProducer)}.{nameof(ProcessDataWrites)} has completed."), Http3ErrorCode.InternalError);
+        }
+
         public void Advance(int bytes)
         {
             lock (_dataWriterLock)

src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/TimingPipeFlusher.cs

@@ -1,6 +1,8 @@
 // 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.
 
+#nullable enable
+
 using System;
 using System.IO.Pipelines;
 using System.Threading;
@@ -36,45 +38,57 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure.PipeW
             return FlushAsync(outputAborter: null, cancellationToken: default);
         }
 
-        public ValueTask<FlushResult> FlushAsync(IHttpOutputAborter outputAborter, CancellationToken cancellationToken)
+        public ValueTask<FlushResult> FlushAsync(IHttpOutputAborter? outputAborter, CancellationToken cancellationToken)
         {
             return FlushAsync(minRate: null, count: 0, outputAborter: outputAborter, cancellationToken: cancellationToken);
         }
 
-        public ValueTask<FlushResult> FlushAsync(MinDataRate minRate, long count)
+        public ValueTask<FlushResult> FlushAsync(MinDataRate? minRate, long count)
         {
             return FlushAsync(minRate, count, outputAborter: null, cancellationToken: default);
         }
 
-        public ValueTask<FlushResult> FlushAsync(MinDataRate minRate, long count, IHttpOutputAborter outputAborter, CancellationToken cancellationToken)
+        public ValueTask<FlushResult> FlushAsync(MinDataRate? minRate, long count, IHttpOutputAborter? outputAborter, CancellationToken cancellationToken)
         {
             var pipeFlushTask = _writer.FlushAsync(cancellationToken);
 
-            if (minRate != null)
+            if (minRate is object)
             {
                 _timeoutControl.BytesWrittenToBuffer(minRate, count);
             }
 
             if (pipeFlushTask.IsCompletedSuccessfully)
             {
-                return new ValueTask<FlushResult>(pipeFlushTask.Result);
+                var flushResult = pipeFlushTask.Result;
+
+                if (flushResult.IsCompleted && outputAborter is object)
+                {
+                    outputAborter.OnInputOrOutputCompleted();
+                }
+
+                return new ValueTask<FlushResult>(flushResult);
             }
 
             return TimeFlushAsyncAwaited(pipeFlushTask, minRate, outputAborter, cancellationToken);
         }
 
-        private async ValueTask<FlushResult> TimeFlushAsyncAwaited(ValueTask<FlushResult> pipeFlushTask, MinDataRate minRate, IHttpOutputAborter outputAborter, CancellationToken cancellationToken)
+        private async ValueTask<FlushResult> TimeFlushAsyncAwaited(ValueTask<FlushResult> pipeFlushTask, MinDataRate? minRate, IHttpOutputAborter? outputAborter, CancellationToken cancellationToken)
         {
-            if (minRate != null)
+            if (minRate is object)
             {
                 _timeoutControl.StartTimingWrite();
             }
 
             try
             {
-                return await pipeFlushTask;
+                var flushResult = await pipeFlushTask;
+
+                if (flushResult.IsCompleted && outputAborter is object)
+                {
+                    outputAborter.OnInputOrOutputCompleted();
+                }
             }
-            catch (OperationCanceledException ex) when (outputAborter != null)
+            catch (OperationCanceledException ex) when (outputAborter is object)
             {
                 outputAborter.Abort(new ConnectionAbortedException(CoreStrings.ConnectionOrStreamAbortedByCancellationToken, ex));
             }
@@ -85,7 +99,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure.PipeW
             }
             finally
             {
-                if (minRate != null)
+                if (minRate is object)
                 {
                     _timeoutControl.StopTimingWrite();
                 }

src/Servers/Kestrel/Core/test/Http1ConnectionTests.cs

@@ -806,6 +806,20 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             Assert.Equal(originalToken, originalRegistration.Token);
         }
 
+        [Fact]
+        public async Task RequestAbortedTokenIsFiredAfterTransportReturnsCompletedFlushResult()
+        {
+            var originalToken = _http1Connection.RequestAborted;
+
+            // Ensure the next call to _transport.Output.FlushAsync() returns a completed FlushResult.
+            _application.Input.Complete();
+
+            await _http1Connection.WritePipeAsync(ReadOnlyMemory<byte>.Empty, default).DefaultTimeout();
+
+            Assert.True(originalToken.WaitHandle.WaitOne(TestConstants.DefaultTimeout));
+            Assert.True(_http1Connection.RequestAborted.WaitHandle.WaitOne(TestConstants.DefaultTimeout));
+        }
+
         [Fact]
         public async Task ExceptionDetailNotIncludedWhenLogLevelInformationNotEnabled()
         {

src/Servers/Kestrel/Core/test/OutputProducerTests.cs

@@ -87,17 +87,19 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
                 pipe,
                 "0",
                 connectionContext,
+                _memoryPool,
                 serviceContext.Log,
                 Mock.Of<ITimeoutControl>(),
                 Mock.Of<IHttpMinResponseDataRateFeature>(),
-                _memoryPool);
+                Mock.Of<IHttpOutputAborter>());
 
             return socketOutput;
         }
 
         private class TestHttpOutputProducer : Http1OutputProducer
         {
-            public TestHttpOutputProducer(Pipe pipe, string connectionId, ConnectionContext connectionContext, IKestrelTrace log, ITimeoutControl timeoutControl, IHttpMinResponseDataRateFeature minResponseDataRateFeature, MemoryPool<byte> memoryPool) : base(pipe.Writer, connectionId, connectionContext, log, timeoutControl, minResponseDataRateFeature, memoryPool)
+            public TestHttpOutputProducer(Pipe pipe, string connectionId, ConnectionContext connectionContext, MemoryPool<byte> memoryPool, IKestrelTrace log, ITimeoutControl timeoutControl, IHttpMinResponseDataRateFeature minResponseDataRateFeature, IHttpOutputAborter outputAborter)
+                : base(pipe.Writer, connectionId, connectionContext, memoryPool, log, timeoutControl, minResponseDataRateFeature, outputAborter)
             {
                 Pipe = pipe;
             }

src/Servers/Kestrel/test/InMemory.FunctionalTests/ChunkedRequestTests.cs

@@ -1047,6 +1047,13 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
                        "0",
                        "",
                        "");
+
+                    await connection.Receive(
+                        "HTTP/1.1 200 OK",
+                        $"Date: {testContext.DateHeaderValue}",
+                        "Content-Length: 0",
+                        "",
+                        "");
                 }
             }