Fix WebSocket close single-consumer violation (#127183)

Author: cittaz

src/libraries/System.Net.WebSockets.Client/tests/CloseTest.cs

@@ -422,9 +422,13 @@ protected async Task RunClient_CloseAsync_DuringConcurrentReceiveAsync_ExpectedS
 
                 await cws.CloseAsync(WebSocketCloseStatus.NormalClosure, "", CancellationToken.None);
 
-                // There is a race condition in the above.  If the ReceiveAsync receives the sent close message from the server,
-                // then it will complete successfully and the socket will close successfully.  If the CloseAsync receive the sent
-                // close message from the server, then the receive async will end up getting aborted along with the socket.
+                // The outcome depends on timing. Two outcomes are acceptable:
+                //   1. The pending ReceiveAsync picks up the server's close frame cleanly: it completes with a Close
+                //      message and the socket transitions to Closed.
+                //   2. The close handshake triggers an internal Abort (e.g. WaitForServerToCloseConnectionAsync
+                //      times out because the server doesn't close TCP within its 1s window): state becomes Aborted,
+                //      the parked receive's stream read fails, and ReceiveAsyncPrivate translates the exception to
+                //      OperationCanceledException because of the Aborted state.
                 try
                 {
                     await t;

src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs

@@ -408,17 +408,23 @@ public override Task CloseOutputAsync(WebSocketCloseStatus closeStatus, string?
             if (NetEventSource.Log.IsEnabled()) NetEventSource.Trace(this);
 
             WebSocketValidate.ValidateCloseStatus(closeStatus, statusDescription);
-            return CloseOutputAsyncCore(closeStatus, statusDescription, cancellationToken);
+            return CloseOutputAsyncCore(closeStatus, statusDescription, enterReceiveMutex: true, cancellationToken: cancellationToken);
         }
 
-        private async Task CloseOutputAsyncCore(WebSocketCloseStatus closeStatus, string? statusDescription, CancellationToken cancellationToken)
+        private async Task CloseOutputAsyncCore(WebSocketCloseStatus closeStatus, string? statusDescription, bool enterReceiveMutex, CancellationToken cancellationToken)
         {
             if (NetEventSource.Log.IsEnabled()) NetEventSource.Trace(this);
 
             ThrowIfInvalidState(WebSocketStateHelper.ValidCloseOutputStates);
 
             await SendCloseFrameAsync(closeStatus, statusDescription, cancellationToken).ConfigureAwait(false);
 
+            // Polite EOF wait under the receive mutex; avoids racing the receive loop on the stream.
+            if (!_isServer && _receivedCloseFrame)
+            {
+                await WaitForServerToCloseConnectionAsync(enterReceiveMutex, cancellationToken).ConfigureAwait(false);
+            }
+
             // If we already received a close frame, since we've now also sent one, we're now closed.
             lock (StateUpdateLock)
             {
@@ -1119,41 +1125,55 @@ private async ValueTask HandleReceivedCloseAsync(MessageHeader header, Cancellat
 
             if (!_isServer && _sentCloseFrame)
             {
-                await WaitForServerToCloseConnectionAsync(cancellationToken).ConfigureAwait(false);
+                await WaitForServerToCloseConnectionAsync(enterMutex: false, cancellationToken).ConfigureAwait(false);
             }
         }
 
-        /// <summary>Issues a read on the stream to wait for EOF.</summary>
-        private async ValueTask WaitForServerToCloseConnectionAsync(CancellationToken cancellationToken)
+        /// <summary>Issues a read on the stream to wait for EOF, optionally acquiring <see cref="_receiveMutex"/> first.</summary>
+        private async ValueTask WaitForServerToCloseConnectionAsync(bool enterMutex, CancellationToken cancellationToken)
         {
-            if (NetEventSource.Log.IsEnabled()) NetEventSource.Trace(this);
+            bool mutexEntered = false;
+            Task? task = null;
+            try
+            {
+                if (enterMutex)
+                {
+                    await _receiveMutex.EnterAsync(cancellationToken).ConfigureAwait(false);
+                    mutexEntered = true;
+                    if (NetEventSource.Log.IsEnabled()) NetEventSource.MutexEntered(_receiveMutex);
+                }
 
-            // Per RFC 6455 7.1.1, try to let the server close the connection.  We give it up to a second.
-            // We simply issue a read and don't care what we get back; we could validate that we don't get
-            // additional data, but at this point we're about to close the connection and we're just stalling
-            // to try to get the server to close first.
-            ValueTask<int> finalReadTask = _stream.ReadAsync(_receiveBuffer, cancellationToken);
+                Debug.Assert(_receiveMutex.IsHeld, $"Expected {nameof(_receiveMutex)} to be held");
 
-            if (finalReadTask.IsCompletedSuccessfully)
-            {
-                finalReadTask.GetAwaiter().GetResult();
-            }
-            else
-            {
-                const int WaitForCloseTimeoutMs = 1_000; // arbitrary amount of time to give the server (same duration as .NET Framework)
-                Task task = finalReadTask.AsTask();
+                if (NetEventSource.Log.IsEnabled()) NetEventSource.Trace(this);
+
+                // Per RFC 6455 7.1.1, try to let the server close the connection.  We give it up to a second.
+                // We simply issue a read and don't care what we get back; we could validate that we don't get
+                // additional data, but at this point we're about to close the connection and we're just stalling
+                // to try to get the server to close first.
+                ValueTask<int> finalReadTask = _stream.ReadAsync(_receiveBuffer, cancellationToken);
+
+                const int WaitForCloseTimeoutMs = 1_000; // arbitrary amount of time to give the server
+                task = finalReadTask.AsTask();
 
-                try
-                {
 #pragma warning disable CA2016 // Token was already provided to the ReadAsync
-                    await task.WaitAsync(TimeSpan.FromMilliseconds(WaitForCloseTimeoutMs)).ConfigureAwait(false);
+                await task.WaitAsync(TimeSpan.FromMilliseconds(WaitForCloseTimeoutMs)).ConfigureAwait(false);
 #pragma warning restore CA2016
-                }
-                catch
+            }
+            catch
+            {
+                if (task is not null)
                 {
-                    // Eat any resulting exceptions. We were going to close the connection, anyway.
                     LogExceptions(task);
-                    Abort();
+                }
+                Abort();
+            }
+            finally
+            {
+                if (mutexEntered)
+                {
+                    _receiveMutex.Exit();
+                    if (NetEventSource.Log.IsEnabled()) NetEventSource.MutexExited(_receiveMutex);
                 }
             }
         }
@@ -1267,12 +1287,14 @@ private static bool IsValidCloseStatus(WebSocketCloseStatus closeStatus)
         private async ValueTask CloseWithReceiveErrorAndThrowAsync(
             WebSocketCloseStatus closeStatus, WebSocketError error, string? errorMessage = null, Exception? innerException = null)
         {
+            Debug.Assert(_receiveMutex.IsHeld, $"Caller should hold the {nameof(_receiveMutex)}");
+
             if (NetEventSource.Log.IsEnabled()) NetEventSource.Trace(this, errorMessage);
 
-            // Close the connection if it hasn't already been closed
+            // Caller holds _receiveMutex; don't re-enter it for the EOF wait.
             if (!_sentCloseFrame)
             {
-                await CloseOutputAsync(closeStatus, string.Empty, default).ConfigureAwait(false);
+                await CloseOutputAsyncCore(closeStatus, string.Empty, enterReceiveMutex: false, cancellationToken: default).ConfigureAwait(false);
             }
 
             // Dump our receive buffer; we're in a bad state to do any further processing
@@ -1491,6 +1513,12 @@ private async Task CloseAsyncPrivate(WebSocketCloseStatus closeStatus, string? s
                     }
                 }
 
+                // Polite EOF wait under the receive mutex; avoids racing the receive loop on the stream.
+                if (!_isServer && _receivedCloseFrame)
+                {
+                    await WaitForServerToCloseConnectionAsync(enterMutex: true, cancellationToken).ConfigureAwait(false);
+                }
+
                 // We're closed.  Close the connection and update the status.
                 lock (StateUpdateLock)
                 {
@@ -1560,11 +1588,6 @@ private async ValueTask SendCloseFrameAsync(WebSocketCloseStatus closeStatus, st
 
                 if (NetEventSource.Log.IsEnabled()) NetEventSource.Trace(this, $"State transition from {state} to {_state}");
             }
-
-            if (!_isServer && _receivedCloseFrame)
-            {
-                await WaitForServerToCloseConnectionAsync(cancellationToken).ConfigureAwait(false);
-            }
         }
 
         private void ConsumeFromBuffer(int count)