Calling abort inside TestServer handler should set RequestAborted (#9975)

src/Hosting/TestHost/src/HttpContextBuilder.cs

@@ -19,7 +19,7 @@ namespace Microsoft.AspNetCore.TestHost
         private TaskCompletionSource<HttpContext> _responseTcs = new TaskCompletionSource<HttpContext>(TaskCreationOptions.RunContinuationsAsynchronously);
         private ResponseStream _responseStream;
         private ResponseFeature _responseFeature = new ResponseFeature();
-        private CancellationTokenSource _requestAbortedSource = new CancellationTokenSource();
+        private RequestLifetimeFeature _requestLifetimeFeature = new RequestLifetimeFeature();
         private bool _pipelineFinished;
         private Context _testContext;
 
@@ -35,9 +35,7 @@ namespace Microsoft.AspNetCore.TestHost
 
             _httpContext.Features.Set<IHttpBodyControlFeature>(this);
             _httpContext.Features.Set<IHttpResponseFeature>(_responseFeature);
-            var requestLifetimeFeature = new HttpRequestLifetimeFeature();
+            _httpContext.Features.Set<IHttpRequestLifetimeFeature>(_requestLifetimeFeature);
-            requestLifetimeFeature.RequestAborted = _requestAbortedSource.Token;
-            _httpContext.Features.Set<IHttpRequestLifetimeFeature>(requestLifetimeFeature);
             
             _responseStream = new ResponseStream(ReturnResponseMessageAsync, AbortRequest, () => AllowSynchronousIO);
             _responseFeature.Body = _responseStream;
@@ -92,7 +90,7 @@ namespace Microsoft.AspNetCore.TestHost
         {
             if (!_pipelineFinished)
             {
-                _requestAbortedSource.Cancel();
+                _requestLifetimeFeature.Abort();
             }
             _responseStream.CompleteWrites();
         }

src/Hosting/TestHost/src/RequestLifetimeFeature.cs

@@ -0,0 +1,22 @@
+// 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.Threading;
+using Microsoft.AspNetCore.Http.Features;
+
+namespace Microsoft.AspNetCore.TestHost
+{
+    internal class RequestLifetimeFeature : IHttpRequestLifetimeFeature
+    {
+        private readonly CancellationTokenSource _cancellationTokenSource = new CancellationTokenSource();
+
+        public RequestLifetimeFeature()
+        {
+            RequestAborted = _cancellationTokenSource.Token;
+        }
+
+        public CancellationToken RequestAborted { get; set; }
+
+        public void Abort() => _cancellationTokenSource.Cancel();
+    }
+}

src/Hosting/TestHost/test/HttpContextBuilderTests.cs

@@ -310,6 +310,24 @@ namespace Microsoft.AspNetCore.TestHost
             var ctx = await server.SendAsync(c => { });
         }
 
+        [Fact]
+        public async Task CallingAbortInsideHandlerShouldSetRequestAborted()
+        {
+            var builder = new WebHostBuilder()
+                .Configure(app =>
+                {
+                    app.Run(context =>
+                    {
+                        context.Abort();
+                        return Task.CompletedTask;
+                    });
+                });
+            var server = new TestServer(builder);
+
+            var ctx = await server.SendAsync(c => { });
+            Assert.True(ctx.RequestAborted.IsCancellationRequested);
+        }
+
         private class VerifierLogger : ILogger<IWebHost>
         {
             public IDisposable BeginScope<TState>(TState state) => new NoopDispoasble();