Merge pull request #26994 from dotnet-maestro-bot/merge/release/3.1-to-master

[automated] Merge branch 'release/3.1' => 'master'

Author: John Luo

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

@@ -192,7 +192,7 @@ public static MessageBody For(
         }
 
         [StackTraceHidden]
-        private void ThrowIfReaderCompleted()
+        protected void ThrowIfReaderCompleted()
         {
             if (_readerCompleted)
             {

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

@@ -14,6 +14,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
     /// </summary>
     internal sealed class Http1UpgradeMessageBody : Http1MessageBody
     {
+        private int _userCanceled;
+
         public Http1UpgradeMessageBody(Http1Connection context, bool keepAlive)
             : base(context, keepAlive)
         {
@@ -23,13 +25,26 @@ public Http1UpgradeMessageBody(Http1Connection context, bool keepAlive)
         // This returns IsEmpty so we can avoid draining the body (since it's basically an endless stream)
         public override bool IsEmpty => true;
 
+        public override ValueTask<ReadResult> ReadAsync(CancellationToken cancellationToken = default)
+        {
+            ThrowIfReaderCompleted();
+            return ReadAsyncInternal(cancellationToken);
+        }
+
+        public override bool TryRead(out ReadResult result)
+        {
+            ThrowIfReaderCompleted();
+            return TryReadInternal(out result);
+        }
+
         public override void AdvanceTo(SequencePosition consumed, SequencePosition examined)
         {
             _context.Input.AdvanceTo(consumed, examined);
         }
 
         public override void CancelPendingRead()
         {
+            Interlocked.Exchange(ref _userCanceled, 1);
             _context.Input.CancelPendingRead();
         }
 
@@ -45,12 +60,49 @@ public override ValueTask StopAsync()
 
         public override bool TryReadInternal(out ReadResult readResult)
         {
-            return _context.Input.TryRead(out readResult);
+            // Ignore the canceled readResult unless it was canceled by the user.
+            do
+            {
+                if (!_context.Input.TryRead(out readResult))
+                {
+                    return false;
+                }
+            } while (readResult.IsCanceled && Interlocked.Exchange(ref _userCanceled, 0) == 0);
+
+            return true;
         }
 
         public override ValueTask<ReadResult> ReadAsyncInternal(CancellationToken cancellationToken = default)
         {
-            return _context.Input.ReadAsync(cancellationToken);
+            ReadResult readResult;
+
+            // Ignore the canceled readResult unless it was canceled by the user.
+            do
+            {
+                var readTask = _context.Input.ReadAsync(cancellationToken);
+
+                if (!readTask.IsCompletedSuccessfully)
+                {
+                    return ReadAsyncInternalAwaited(readTask, cancellationToken);
+                }
+
+                readResult = readTask.GetAwaiter().GetResult();
+            } while (readResult.IsCanceled && Interlocked.Exchange(ref _userCanceled, 0) == 0);
+
+            return new ValueTask<ReadResult>(readResult);
+        }
+
+        private async ValueTask<ReadResult> ReadAsyncInternalAwaited(ValueTask<ReadResult> readTask, CancellationToken cancellationToken = default)
+        {
+            var readResult = await readTask;
+
+            // Ignore the canceled readResult unless it was canceled by the user.
+            while (readResult.IsCanceled && Interlocked.Exchange(ref _userCanceled, 0) == 0)
+            {
+                readResult = await _context.Input.ReadAsync(cancellationToken);
+            }
+
+            return readResult;
         }
     }
 }

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

@@ -3,7 +3,10 @@
 
 using System;
 using System.IO;
+using System.IO.Pipelines;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Connections;
+using Microsoft.AspNetCore.Connections.Features;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
@@ -343,6 +346,51 @@ await connection.Receive("HTTP/1.1 101 Switching Protocols",
             }
         }
 
+        [Fact]
+        public async Task DoesNotThrowGivenCanceledReadResult()
+        {
+            var appCompletedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            await using var server = new TestServer(async context =>
+            {
+                try
+                {
+                    var upgradeFeature = context.Features.Get<IHttpUpgradeFeature>();
+                    var duplexStream = await upgradeFeature.UpgradeAsync();
+
+                    // Kestrel will call Transport.Input.CancelPendingRead() during shutdown so idle connections
+                    // can wake up and shutdown gracefully. We manually call CancelPendingRead() to simulate this and
+                    // ensure the Stream returned by UpgradeAsync doesn't throw in this case.
+                    // https://github.com/dotnet/aspnetcore/issues/26482
+                    var connectionTransportFeature = context.Features.Get<IConnectionTransportFeature>();
+                    connectionTransportFeature.Transport.Input.CancelPendingRead();
+
+                    // Use ReadAsync() instead of CopyToAsync() for this test since IsCanceled is only checked in
+                    // HttpRequestStream.ReadAsync() and not HttpRequestStream.CopyToAsync() 
+                    Assert.Equal(0, await duplexStream.ReadAsync(new byte[1]));
+                    appCompletedTcs.SetResult(null);
+                }
+                catch (Exception ex)
+                {
+                    appCompletedTcs.SetException(ex);
+                    throw;
+                }
+            },
+            new TestServiceContext(LoggerFactory));
+
+            using (var connection = server.CreateConnection())
+            {
+                await connection.SendEmptyGetWithUpgrade();
+                await connection.Receive("HTTP/1.1 101 Switching Protocols",
+                    "Connection: Upgrade",
+                    $"Date: {server.Context.DateHeaderValue}",
+                    "",
+                    "");
+            }
+
+            await appCompletedTcs.Task.DefaultTimeout();
+        }
+
         [Fact]
         public async Task DoesNotCloseConnectionWithout101Response()
         {