Fix ManagedWebSocket send cancellation leaving state Open

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<byte> 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<byte>(_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<byte>(_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);
+                    }
                 }
             }
         }

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<OperationCanceledException>(() =>
+                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<OperationCanceledException>(() => 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
+            {
+            }
+        }
+
+        /// <summary>A test stream whose first WriteAsync blocks until released, allowing tests to
+        /// deterministically create contention on the WebSocket's internal send mutex.</summary>
+        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<byte> 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) =>