From 790e85f39112fb7fe5c4eb81de5d0ea03c3d5b7f Mon Sep 17 00:00:00 2001 From: Steve Pfister Date: Thu, 3 Sep 2026 15:38:42 -0400 Subject: [PATCH 1/2] Fix ManagedWebSocket send cancellation leaving state Open (#132031) SendFrameFallbackAsync awaited the send-mutex acquisition task before entering its try block and before registering the cancellation -> Abort() callback. If cancellation raced with acquiring _sendMutex (e.g. because a keep-alive ping or another in-flight send already held it), the OperationCanceledException from the canceled mutex-wait task propagated out without ever calling Abort(), leaving WebSocketState stuck at Open instead of transitioning to Aborted. This caused the flaky System.Net.WebSockets.Client.Tests.CancelTest_* .SendAsync_Cancel_Success failures tracked by #132031. Move the cancellationToken.Register(... => Abort()) registration to wrap the entire method body, including the mutex-acquisition wait, matching the existing pattern already used by the receive path (ReceiveAsyncPrivate registers cancellation before entering _receiveMutex). Also trace a canceled mutex wait via NetEventSource.TraceException before rethrowing, for diagnostics parity with the write/flush cancellation path. Add two deterministic regression tests to WebSocketTests.cs: - SendAsync_AlreadyCanceledToken_AbortsConnectionAndThrowsOperationCanceledException, using an already-canceled token so AsyncMutex.EnterAsync returns Task.FromCanceled immediately, hitting the pre-lock cancellation path with no contention or timing. - SendAsync_CancelWhileWaitingForSendMutex_AbortsConnectionAndThrowsOperationCanceledException, using a new GatedWriteStream test helper whose WriteAsync blocks on a TaskCompletionSource to deterministically create genuine send-mutex contention (no sleeps), then cancels a second, waiting send. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../System/Net/WebSockets/ManagedWebSocket.cs | 63 ++++++++++------- .../tests/WebSocketTests.cs | 68 +++++++++++++++++++ 2 files changed, 106 insertions(+), 25 deletions(-) diff --git a/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs b/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs index aa07aef340474a..888e93ff3d9d8c 100644 --- a/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs +++ b/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs @@ -599,40 +599,53 @@ private async ValueTask WaitForWriteTaskAsync(ValueTask writeTask, bool shouldFl private async ValueTask SendFrameFallbackAsync(MessageOpcode opcode, bool endOfMessage, bool disableCompression, ReadOnlyMemory payloadBuffer, Task lockTask, CancellationToken cancellationToken) { - await lockTask.ConfigureAwait(false); - if (NetEventSource.Log.IsEnabled()) NetEventSource.MutexEntered(_sendMutex); - - try + // Register for cancellation before waiting on the lock so that if cancellation races with + // acquiring the mutex, we still abort the connection, just as we would if cancellation raced + // with the write itself. Without this, a cancellation that fires while we're waiting to enter + // the mutex would propagate out without transitioning the WebSocket to the Aborted state. + using (cancellationToken.Register(static s => ((ManagedWebSocket)s!).Abort(), this)) { - int sendBytes = WriteFrameToSendBuffer(opcode, endOfMessage, disableCompression, payloadBuffer.Span); - using (cancellationToken.Register(static s => ((ManagedWebSocket)s!).Abort(), this)) + try { - await _stream.WriteAsync(new ReadOnlyMemory(_sendBuffer, 0, sendBytes), cancellationToken).ConfigureAwait(false); - await _stream.FlushAsync(cancellationToken).ConfigureAwait(false); + await lockTask.ConfigureAwait(false); } - } - catch (Exception exc) - { - if (NetEventSource.Log.IsEnabled()) NetEventSource.TraceException(this, exc); - - if (exc is OperationCanceledException) + catch (Exception exc) { + if (NetEventSource.Log.IsEnabled()) NetEventSource.TraceException(this, exc); throw; } - throw _state == WebSocketState.Aborted ? - CreateOperationCanceledException(exc, cancellationToken) : - new WebSocketException(WebSocketError.ConnectionClosedPrematurely, exc); - } - finally - { - ReleaseSendBuffer(); - _sendMutex.Exit(); + if (NetEventSource.Log.IsEnabled()) NetEventSource.MutexEntered(_sendMutex); - if (NetEventSource.Log.IsEnabled()) + try { - NetEventSource.MutexExited(_sendMutex); - NetEventSource.SendFrameAsyncCompleted(this); + int sendBytes = WriteFrameToSendBuffer(opcode, endOfMessage, disableCompression, payloadBuffer.Span); + await _stream.WriteAsync(new ReadOnlyMemory(_sendBuffer, 0, sendBytes), cancellationToken).ConfigureAwait(false); + await _stream.FlushAsync(cancellationToken).ConfigureAwait(false); + } + catch (Exception exc) + { + if (NetEventSource.Log.IsEnabled()) NetEventSource.TraceException(this, exc); + + if (exc is OperationCanceledException) + { + throw; + } + + throw _state == WebSocketState.Aborted ? + CreateOperationCanceledException(exc, cancellationToken) : + new WebSocketException(WebSocketError.ConnectionClosedPrematurely, exc); + } + finally + { + ReleaseSendBuffer(); + _sendMutex.Exit(); + + if (NetEventSource.Log.IsEnabled()) + { + NetEventSource.MutexExited(_sendMutex); + NetEventSource.SendFrameAsyncCompleted(this); + } } } } diff --git a/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs b/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs index cf01437fc3d97d..ce662345542081 100644 --- a/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs +++ b/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs @@ -252,6 +252,74 @@ public async Task ReceiveAsync_AfterCancellationDoReceiveAsync_ThrowsWebSocketEx ex.Message); } + [Fact] + public async Task SendAsync_AlreadyCanceledToken_AbortsConnectionAndThrowsOperationCanceledException() + { + using var stream = new WebSocketTestStream(); + using var websocket = WebSocket.CreateFromStream(stream, new WebSocketCreationOptions()); + + var cts = new CancellationTokenSource(); + cts.Cancel(); + + await Assert.ThrowsAnyAsync(() => + websocket.SendAsync(new byte[] { 1, 2, 3 }, WebSocketMessageType.Binary, endOfMessage: true, cts.Token).AsTask()); + + Assert.Equal(WebSocketState.Aborted, websocket.State); + } + + [Fact] + public async Task SendAsync_CancelWhileWaitingForSendMutex_AbortsConnectionAndThrowsOperationCanceledException() + { + using var stream = new GatedWriteStream(); + using var websocket = WebSocket.CreateFromStream(stream, new WebSocketCreationOptions()); + + // Start a send with a non-cancelable token; it acquires the send mutex synchronously + // and then blocks inside the (gated) write, simulating another in-flight send -- e.g. a + // keep-alive ping -- holding the mutex. + Task firstSend = websocket.SendAsync(new byte[] { 1 }, WebSocketMessageType.Binary, endOfMessage: true, CancellationToken.None).AsTask(); + await stream.WriteStarted; + + // Issue a second, cancelable send. Since the mutex is held, it must wait to acquire it. + using var cts = new CancellationTokenSource(); + Task secondSend = websocket.SendAsync(new byte[] { 2 }, WebSocketMessageType.Binary, endOfMessage: true, cts.Token).AsTask(); + + // Cancel while the second send is still waiting on the send mutex. + cts.Cancel(); + await Assert.ThrowsAnyAsync(() => secondSend); + + Assert.Equal(WebSocketState.Aborted, websocket.State); + + // Unblock the first send so it can finish (successfully or not -- that's not what this + // test is about) and the test can complete deterministically. + stream.ReleaseWrite(); + try + { + await firstSend; + } + catch + { + } + } + + /// A test stream whose first WriteAsync blocks until released, allowing tests to + /// deterministically create contention on the WebSocket's internal send mutex. + private sealed class GatedWriteStream : WebSocketTestStream + { + private readonly TaskCompletionSource _writeStarted = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _releaseWrite = new(TaskCreationOptions.RunContinuationsAsynchronously); + + public Task WriteStarted => _writeStarted.Task; + + public void ReleaseWrite() => _releaseWrite.TrySetResult(); + + public override async ValueTask WriteAsync(ReadOnlyMemory buffer, CancellationToken cancellationToken) + { + _writeStarted.TrySetResult(); + await _releaseWrite.Task.ConfigureAwait(false); + await base.WriteAsync(buffer, cancellationToken).ConfigureAwait(false); + } + } + public abstract class ExposeProtectedWebSocket : WebSocket { public static new bool IsStateTerminal(WebSocketState state) => From a39cd0bb09e8b0c9928acd96cdcc1f179fe4ecf4 Mon Sep 17 00:00:00 2001 From: Steve Pfister Date: Fri, 11 Sep 2026 18:12:25 -0400 Subject: [PATCH 2/2] Fix WebSocket send cancellation tests Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs b/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs index ce662345542081..48c62ac94cfefd 100644 --- a/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs +++ b/src/libraries/System.Net.WebSockets/tests/WebSocketTests.cs @@ -262,7 +262,7 @@ public async Task SendAsync_AlreadyCanceledToken_AbortsConnectionAndThrowsOperat cts.Cancel(); await Assert.ThrowsAnyAsync(() => - websocket.SendAsync(new byte[] { 1, 2, 3 }, WebSocketMessageType.Binary, endOfMessage: true, cts.Token).AsTask()); + websocket.SendAsync(new byte[] { 1, 2, 3 }, WebSocketMessageType.Binary, endOfMessage: true, cts.Token)); Assert.Equal(WebSocketState.Aborted, websocket.State); } @@ -276,12 +276,12 @@ public async Task SendAsync_CancelWhileWaitingForSendMutex_AbortsConnectionAndTh // Start a send with a non-cancelable token; it acquires the send mutex synchronously // and then blocks inside the (gated) write, simulating another in-flight send -- e.g. a // keep-alive ping -- holding the mutex. - Task firstSend = websocket.SendAsync(new byte[] { 1 }, WebSocketMessageType.Binary, endOfMessage: true, CancellationToken.None).AsTask(); + Task firstSend = websocket.SendAsync(new byte[] { 1 }, WebSocketMessageType.Binary, endOfMessage: true, CancellationToken.None); await stream.WriteStarted; // Issue a second, cancelable send. Since the mutex is held, it must wait to acquire it. using var cts = new CancellationTokenSource(); - Task secondSend = websocket.SendAsync(new byte[] { 2 }, WebSocketMessageType.Binary, endOfMessage: true, cts.Token).AsTask(); + Task secondSend = websocket.SendAsync(new byte[] { 2 }, WebSocketMessageType.Binary, endOfMessage: true, cts.Token); // Cancel while the second send is still waiting on the send mutex. cts.Cancel();