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..48c62ac94cfefd 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)); + + 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); + 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); + + // 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) =>