Skip to content

Commit c15cd69

Browse files
authored
Prevent WebSockets from throwing during graceful shutdown (#27123)
- Fix canceled ReadResult handling in Http1UpgradeMessageBody - Add UpgradeTests.DoesNotThrowGivenCanceledReadResult
1 parent 074069f commit c15cd69

File tree

2 files changed

+92
-4
lines changed

2 files changed

+92
-4
lines changed

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

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
1414
/// </summary>
1515
internal sealed class Http1UpgradeMessageBody : Http1MessageBody
1616
{
17+
private int _userCanceled;
18+
1719
public Http1UpgradeMessageBody(Http1Connection context, bool keepAlive)
1820
: base(context, keepAlive)
1921
{
@@ -26,13 +28,13 @@ public Http1UpgradeMessageBody(Http1Connection context, bool keepAlive)
2628
public override ValueTask<ReadResult> ReadAsync(CancellationToken cancellationToken = default)
2729
{
2830
ThrowIfCompleted();
29-
return _context.Input.ReadAsync(cancellationToken);
31+
return ReadAsyncInternal(cancellationToken);
3032
}
3133

3234
public override bool TryRead(out ReadResult result)
3335
{
3436
ThrowIfCompleted();
35-
return _context.Input.TryRead(out result);
37+
return TryReadInternal(out result);
3638
}
3739

3840
public override void AdvanceTo(SequencePosition consumed)
@@ -54,6 +56,7 @@ public override void Complete(Exception exception)
5456

5557
public override void CancelPendingRead()
5658
{
59+
Interlocked.Exchange(ref _userCanceled, 1);
5760
_context.Input.CancelPendingRead();
5861
}
5962

@@ -69,12 +72,49 @@ public override Task StopAsync()
6972

7073
public override bool TryReadInternal(out ReadResult readResult)
7174
{
72-
return _context.Input.TryRead(out readResult);
75+
// Ignore the canceled readResult unless it was canceled by the user.
76+
do
77+
{
78+
if (!_context.Input.TryRead(out readResult))
79+
{
80+
return false;
81+
}
82+
} while (readResult.IsCanceled && Interlocked.Exchange(ref _userCanceled, 0) == 0);
83+
84+
return true;
7385
}
7486

7587
public override ValueTask<ReadResult> ReadAsyncInternal(CancellationToken cancellationToken = default)
7688
{
77-
return _context.Input.ReadAsync(cancellationToken);
89+
ReadResult readResult;
90+
91+
// Ignore the canceled readResult unless it was canceled by the user.
92+
do
93+
{
94+
var readTask = _context.Input.ReadAsync(cancellationToken);
95+
96+
if (!readTask.IsCompletedSuccessfully)
97+
{
98+
return ReadAsyncInternalAwaited(readTask, cancellationToken);
99+
}
100+
101+
readResult = readTask.GetAwaiter().GetResult();
102+
} while (readResult.IsCanceled && Interlocked.Exchange(ref _userCanceled, 0) == 0);
103+
104+
return new ValueTask<ReadResult>(readResult);
105+
}
106+
107+
private async ValueTask<ReadResult> ReadAsyncInternalAwaited(ValueTask<ReadResult> readTask, CancellationToken cancellationToken = default)
108+
{
109+
var readResult = await readTask;
110+
111+
// Ignore the canceled readResult unless it was canceled by the user.
112+
while (readResult.IsCanceled && Interlocked.Exchange(ref _userCanceled, 0) == 0)
113+
{
114+
readResult = await _context.Input.ReadAsync(cancellationToken);
115+
}
116+
117+
return readResult;
78118
}
79119
}
80120
}

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

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,10 @@
33

44
using System;
55
using System.IO;
6+
using System.IO.Pipelines;
67
using System.Threading.Tasks;
8+
using Microsoft.AspNetCore.Connections;
9+
using Microsoft.AspNetCore.Connections.Features;
710
using Microsoft.AspNetCore.Http.Features;
811
using Microsoft.AspNetCore.Server.Kestrel.Core;
912
using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
@@ -375,5 +378,50 @@ await connection.Receive("HTTP/1.1 101 Switching Protocols",
375378
}
376379
}
377380
}
381+
382+
[Fact]
383+
public async Task DoesNotThrowGivenCanceledReadResult()
384+
{
385+
var appCompletedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
386+
387+
await using var server = new TestServer(async context =>
388+
{
389+
try
390+
{
391+
var upgradeFeature = context.Features.Get<IHttpUpgradeFeature>();
392+
var duplexStream = await upgradeFeature.UpgradeAsync();
393+
394+
// Kestrel will call Transport.Input.CancelPendingRead() during shutdown so idle connections
395+
// can wake up and shutdown gracefully. We manually call CancelPendingRead() to simulate this and
396+
// ensure the Stream returned by UpgradeAsync doesn't throw in this case.
397+
// https://github.com/dotnet/aspnetcore/issues/26482
398+
var connectionTransportFeature = context.Features.Get<IConnectionTransportFeature>();
399+
connectionTransportFeature.Transport.Input.CancelPendingRead();
400+
401+
// Use ReadAsync() instead of CopyToAsync() for this test since IsCanceled is only checked in
402+
// HttpRequestStream.ReadAsync() and not HttpRequestStream.CopyToAsync()
403+
Assert.Equal(0, await duplexStream.ReadAsync(new byte[1]));
404+
appCompletedTcs.SetResult(null);
405+
}
406+
catch (Exception ex)
407+
{
408+
appCompletedTcs.SetException(ex);
409+
throw;
410+
}
411+
},
412+
new TestServiceContext(LoggerFactory));
413+
414+
using (var connection = server.CreateConnection())
415+
{
416+
await connection.SendEmptyGetWithUpgrade();
417+
await connection.Receive("HTTP/1.1 101 Switching Protocols",
418+
"Connection: Upgrade",
419+
$"Date: {server.Context.DateHeaderValue}",
420+
"",
421+
"");
422+
}
423+
424+
await appCompletedTcs.Task.DefaultTimeout();
425+
}
378426
}
379427
}

0 commit comments

Comments
 (0)