TestServer Reset support, fix up Abort

src/Hosting/TestHost/ref/Microsoft.AspNetCore.TestHost.netcoreapp3.0.cs

@@ -14,6 +14,11 @@ public static partial class HostBuilderTestServerExtensions
         public static System.Net.Http.HttpClient GetTestClient(this Microsoft.Extensions.Hosting.IHost host) { throw null; }
         public static Microsoft.AspNetCore.TestHost.TestServer GetTestServer(this Microsoft.Extensions.Hosting.IHost host) { throw null; }
     }
+    public partial class HttpResetTestException : System.Exception
+    {
+        public HttpResetTestException(int errorCode) { }
+        public int ErrorCode { [System.Runtime.CompilerServices.CompilerGeneratedAttribute]get { throw null; } }
+    }
     public partial class RequestBuilder
     {
         public RequestBuilder(Microsoft.AspNetCore.TestHost.TestServer server, string path) { }

src/Hosting/TestHost/src/ClientHandler.cs

@@ -70,7 +70,15 @@ protected override async Task<HttpResponseMessage> SendAsync(
             {
                 var req = context.Request;
 
-                req.Protocol = "HTTP/" + request.Version.ToString(fieldCount: 2);
+                if (request.Version == HttpVersion.Version20)
+                {
+                    // https://tools.ietf.org/html/rfc7540
+                    req.Protocol = "HTTP/2";
+                }
+                else
+                {
+                    req.Protocol = "HTTP/" + request.Version.ToString(fieldCount: 2);
+                }
                 req.Method = request.Method.ToString();
 
                 req.Scheme = request.RequestUri.Scheme;

src/Hosting/TestHost/src/HttpContextBuilder.cs

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.IO;
 using System.IO.Pipelines;
 using System.Threading;
 using System.Threading.Tasks;
@@ -10,7 +11,7 @@
 
 namespace Microsoft.AspNetCore.TestHost
 {
-    internal class HttpContextBuilder : IHttpBodyControlFeature
+    internal class HttpContextBuilder : IHttpBodyControlFeature, IHttpResetFeature
     {
         private readonly ApplicationWrapper _application;
         private readonly bool _preserveExecutionContext;
@@ -20,7 +21,7 @@ internal class HttpContextBuilder : IHttpBodyControlFeature
         private readonly ResponseBodyReaderStream _responseReaderStream;
         private readonly ResponseBodyPipeWriter _responsePipeWriter;
         private readonly ResponseFeature _responseFeature;
-        private readonly RequestLifetimeFeature _requestLifetimeFeature = new RequestLifetimeFeature();
+        private readonly RequestLifetimeFeature _requestLifetimeFeature;
         private readonly ResponseTrailersFeature _responseTrailersFeature = new ResponseTrailersFeature();
         private bool _pipelineFinished;
         private bool _returningResponse;
@@ -34,13 +35,14 @@ internal HttpContextBuilder(ApplicationWrapper application, bool allowSynchronou
             _preserveExecutionContext = preserveExecutionContext;
             _httpContext = new DefaultHttpContext();
             _responseFeature = new ResponseFeature(Abort);
+            _requestLifetimeFeature = new RequestLifetimeFeature(Abort);
 
             var request = _httpContext.Request;
             request.Protocol = "HTTP/1.1";
             request.Method = HttpMethods.Get;
 
             var pipe = new Pipe();
-            _responseReaderStream = new ResponseBodyReaderStream(pipe, AbortRequest, () => _responseReadCompleteCallback?.Invoke(_httpContext));
+            _responseReaderStream = new ResponseBodyReaderStream(pipe, ClientInitiatedAbort, () => _responseReadCompleteCallback?.Invoke(_httpContext));
             _responsePipeWriter = new ResponseBodyPipeWriter(pipe, ReturnResponseMessageAsync);
             _responseFeature.Body = new ResponseBodyWriterStream(_responsePipeWriter, () => AllowSynchronousIO);
             _responseFeature.BodySnapshot = _responseFeature.Body;
@@ -77,11 +79,17 @@ internal void RegisterResponseReadCompleteCallback(Action<HttpContext> responseR
         /// <returns></returns>
         internal Task<HttpContext> SendAsync(CancellationToken cancellationToken)
         {
-            var registration = cancellationToken.Register(AbortRequest);
+            var registration = cancellationToken.Register(ClientInitiatedAbort);
 
             // Everything inside this function happens in the SERVER's execution context (unless PreserveExecutionContext is true)
             async Task RunRequestAsync()
             {
+                // HTTP/2 specific features must be added after the request has been configured.
+                if (string.Equals("HTTP/2", _httpContext.Request.Protocol, StringComparison.OrdinalIgnoreCase))
+                {
+                    _httpContext.Features.Set<IHttpResetFeature>(this);
+                }
+
                 // This will configure IHttpContextAccessor so it needs to happen INSIDE this function,
                 // since we are now inside the Server's execution context. If it happens outside this cont
                 // it will be lost when we abandon the execution context.
@@ -120,13 +128,16 @@ async Task RunRequestAsync()
             return _responseTcs.Task;
         }
 
-        internal void AbortRequest()
+        // Triggered by request CancellationToken canceling or response stream Disposal.
+        internal void ClientInitiatedAbort()
         {
             if (!_pipelineFinished)
             {
-                _requestLifetimeFeature.Abort();
+                // We don't want to trigger the token for already completed responses.
+                _requestLifetimeFeature.Cancel();
             }
-            _responsePipeWriter.Complete();
+            // Writes will still succeed, the app will only get an error if they check the CT.
+            _responseReaderStream.Abort(new IOException("The client aborted the request."));
         }
 
         internal async Task CompleteResponseAsync()
@@ -178,10 +189,15 @@ internal async Task ReturnResponseMessageAsync()
 
         internal void Abort(Exception exception)
         {
-            _pipelineFinished = true;
             _responsePipeWriter.Abort(exception);
             _responseReaderStream.Abort(exception);
+            _requestLifetimeFeature.Cancel();
             _responseTcs.TrySetException(exception);
         }
+
+        void IHttpResetFeature.Reset(int errorCode)
+        {
+            Abort(new HttpResetTestException(errorCode));
+        }
     }
 }

src/Hosting/TestHost/src/HttpResetTestException.cs

@@ -0,0 +1,29 @@
+// 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;
+using Microsoft.AspNetCore.Http.Features;
+
+namespace Microsoft.AspNetCore.TestHost
+{
+    /// <summary>
+    /// Used to surface to the test client that the application invoked <see cref="IHttpResetFeature.Reset"/>
+    /// </summary>
+    public class HttpResetTestException : Exception
+    {
+        /// <summary>
+        /// Creates a new test exception
+        /// </summary>
+        /// <param name="errorCode">The error code passed to <see cref="IHttpResetFeature.Reset"/></param>
+        public HttpResetTestException(int errorCode)
+            : base($"The application reset the request with error code {errorCode}.")
+        {
+            ErrorCode = errorCode;
+        }
+
+        /// <summary>
+        /// The error code passed to <see cref="IHttpResetFeature.Reset"/>
+        /// </summary>
+        public int ErrorCode { get; }
+    }
+}

src/Hosting/TestHost/src/RequestLifetimeFeature.cs

@@ -1,6 +1,7 @@
 // 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;
 using System.Threading;
 using Microsoft.AspNetCore.Http.Features;
 
@@ -9,14 +10,25 @@ namespace Microsoft.AspNetCore.TestHost
     internal class RequestLifetimeFeature : IHttpRequestLifetimeFeature
     {
         private readonly CancellationTokenSource _cancellationTokenSource = new CancellationTokenSource();
+        private readonly Action<Exception> _abort;
 
-        public RequestLifetimeFeature()
+        public RequestLifetimeFeature(Action<Exception> abort)
         {
             RequestAborted = _cancellationTokenSource.Token;
+            _abort = abort;
         }
 
         public CancellationToken RequestAborted { get; set; }
 
-        public void Abort() => _cancellationTokenSource.Cancel();
+        internal void Cancel()
+        {
+            _cancellationTokenSource.Cancel();
+        }
+
+        void IHttpRequestLifetimeFeature.Abort()
+        {
+            _abort(new Exception("The application aborted the request."));
+            _cancellationTokenSource.Cancel();
+        }
     }
 }

src/Hosting/TestHost/src/ResponseBodyReaderStream.cs

@@ -80,6 +80,11 @@ public async override Task<int> ReadAsync(byte[] buffer, int offset, int count,
             using var registration = cancellationToken.Register(Cancel);
             var result = await _pipe.Reader.ReadAsync(cancellationToken);
 
+            if (result.IsCanceled)
+            {
+                throw new OperationCanceledException();
+            }
+
             if (result.Buffer.IsEmpty && result.IsCompleted)
             {
                 _pipe.Reader.Complete();
@@ -114,16 +119,16 @@ private static void VerifyBuffer(byte[] buffer, int offset, int count)
 
         internal void Cancel()
         {
-            _aborted = true;
-            _abortException = new OperationCanceledException();
-            _pipe.Writer.Complete(_abortException);
+            Abort(new OperationCanceledException());
         }
 
         internal void Abort(Exception innerException)
         {
             Contract.Requires(innerException != null);
             _aborted = true;
             _abortException = innerException;
+            _pipe.Reader.CancelPendingRead();
+            _pipe.Reader.Complete();
         }
 
         private void CheckAborted()

src/Hosting/TestHost/test/ClientHandlerTests.cs

@@ -301,8 +301,7 @@ public async Task ClientDisposalCloses()
             Task<int> readTask = responseStream.ReadAsync(new byte[100], 0, 100);
             Assert.False(readTask.IsCompleted);
             responseStream.Dispose();
-            var read = await readTask.WithTimeout();
-            Assert.Equal(0, read);
+            await Assert.ThrowsAsync<OperationCanceledException>(() => readTask.WithTimeout());
             block.SetResult(0);
         }
 

src/Hosting/TestHost/test/HttpContextBuilderTests.cs

@@ -194,8 +194,7 @@ public async Task ClientDisposalCloses()
             Task<int> readTask = responseStream.ReadAsync(new byte[100], 0, 100);
             Assert.False(readTask.IsCompleted);
             responseStream.Dispose();
-            var read = await readTask.WithTimeout();
-            Assert.Equal(0, read);
+            await Assert.ThrowsAsync<OperationCanceledException>(() => readTask.WithTimeout());
             block.SetResult(0);
         }
 
@@ -313,19 +312,22 @@ public async Task ClientHandlerCreateContextWithDefaultRequestParameters()
         [Fact]
         public async Task CallingAbortInsideHandlerShouldSetRequestAborted()
         {
+            var requestAborted = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var builder = new WebHostBuilder()
                 .Configure(app =>
                 {
                     app.Run(context =>
                     {
+                        context.RequestAborted.Register(() => requestAborted.SetResult(0));
                         context.Abort();
                         return Task.CompletedTask;
                     });
                 });
             var server = new TestServer(builder);
 
-            var ctx = await server.SendAsync(c => { });
-            Assert.True(ctx.RequestAborted.IsCancellationRequested);
+            var ex = await Assert.ThrowsAsync<Exception>(() => server.SendAsync(c => { }));
+            Assert.Equal("The application aborted the request.", ex.Message);
+            await requestAborted.Task.WithTimeout();
         }
 
         private class VerifierLogger : ILogger<IWebHost>

src/Hosting/TestHost/test/RequestLifetimeTests.cs

@@ -0,0 +1,115 @@
+// 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;
+using System.Net.Http;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Hosting;
+using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Features;
+using Microsoft.Extensions.Hosting;
+using Xunit;
+
+namespace Microsoft.AspNetCore.TestHost
+{
+    public class RequestLifetimeTests
+    {
+        [Fact]
+        public async Task LifetimeFeature_Abort_TriggersRequestAbortedToken()
+        {
+            var requestAborted = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
+            using var host = await CreateHost(async httpContext =>
+            {
+                httpContext.RequestAborted.Register(() => requestAborted.SetResult(0));
+                httpContext.Abort();
+
+                await requestAborted.Task.WithTimeout();
+            });
+
+            var client = host.GetTestServer().CreateClient();
+            var ex = await Assert.ThrowsAsync<Exception>(() => client.GetAsync("/", HttpCompletionOption.ResponseHeadersRead));
+            Assert.Equal("The application aborted the request.", ex.Message);
+            await requestAborted.Task.WithTimeout();
+        }
+
+        [Fact]
+        public async Task LifetimeFeature_AbortBeforeHeadersSent_ClientThrows()
+        {
+            var abortReceived = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
+            using var host = await CreateHost(async httpContext =>
+            {
+                httpContext.Abort();
+                await abortReceived.Task.WithTimeout();
+            });
+
+            var client = host.GetTestServer().CreateClient();
+            var ex = await Assert.ThrowsAsync<Exception>(() => client.GetAsync("/", HttpCompletionOption.ResponseHeadersRead));
+            Assert.Equal("The application aborted the request.", ex.Message);
+            abortReceived.SetResult(0);
+        }
+
+        [Fact]
+        public async Task LifetimeFeature_AbortAfterHeadersSent_ClientBodyThrows()
+        {
+            var responseReceived = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var abortReceived = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
+            using var host = await CreateHost(async httpContext =>
+            {
+                await httpContext.Response.Body.FlushAsync();
+                await responseReceived.Task.WithTimeout();
+                httpContext.Abort();
+                await abortReceived.Task.WithTimeout();
+            });
+
+            var client = host.GetTestServer().CreateClient();
+            var response = await client.GetAsync("/", HttpCompletionOption.ResponseHeadersRead);
+            responseReceived.SetResult(0);
+            response.EnsureSuccessStatusCode();
+            var ex = await Assert.ThrowsAsync<HttpRequestException>(() => response.Content.ReadAsByteArrayAsync());
+            var rex = ex.GetBaseException();
+            Assert.Equal("The application aborted the request.", rex.Message);
+            abortReceived.SetResult(0);
+        }
+
+        [Fact]
+        public async Task LifetimeFeature_AbortAfterSomeDataSent_ClientBodyThrows()
+        {
+            var responseReceived = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var abortReceived = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
+            using var host = await CreateHost(async httpContext =>
+            {
+                await httpContext.Response.WriteAsync("Hello World");
+                await responseReceived.Task.WithTimeout();
+                httpContext.Abort();
+                await abortReceived.Task.WithTimeout();
+            });
+
+            var client = host.GetTestServer().CreateClient();
+            var response = await client.GetAsync("/", HttpCompletionOption.ResponseHeadersRead);
+            responseReceived.SetResult(0);
+            response.EnsureSuccessStatusCode();
+            var ex = await Assert.ThrowsAsync<HttpRequestException>(() => response.Content.ReadAsByteArrayAsync());
+            var rex = ex.GetBaseException();
+            Assert.Equal("The application aborted the request.", rex.Message);
+            abortReceived.SetResult(0);
+        }
+
+        // TODO: Abort after CompleteAsync - No-op, the request is already complete.
+
+        private Task<IHost> CreateHost(RequestDelegate appDelegate)
+        {
+            return new HostBuilder()
+                .ConfigureWebHost(webBuilder =>
+                {
+                    webBuilder
+                        .UseTestServer()
+                        .Configure(app =>
+                        {
+                            app.Run(appDelegate);
+                        });
+                })
+                .StartAsync();
+        }
+    }
+}