Fix ClientWebSocket closing handshake logic (#36975)

An HttpListener websocket test was failing after the change from PR #36928. That PR made changes
to tighten up the transition to the Closed state after the closing handshake had completed.

That PR missed some additional changes needed especially in cases where a server (such as loopback)
sends the close frame almost concurrently with the client sending the close frame.

This PR fixes the close frame handshake logic and also makes some optimizations in cases where the
websocket has already transitioned to the Aborted state during the closing handshake. In that case,
the websocket should not wait for a close frame to be received but should simply proceed to closing
the connection.

Fixes #36963

Author: davidsh

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

@@ -1063,50 +1063,56 @@ private async Task CloseAsyncPrivate(WebSocketCloseStatus closeStatus, string st
                 await SendCloseFrameAsync(closeStatus, statusDescription, cancellationToken).ConfigureAwait(false);
             }
 
-            // We should now either be in a CloseSent case (because we just sent one), or in a CloseReceived state, in case
+            // We should now either be in a CloseSent case (because we just sent one), or in a Closed state, in case
             // there was a concurrent receive that ended up handling an immediate close frame response from the server.
             // Of course it could also be Aborted if something happened concurrently to cause things to blow up.
             Debug.Assert(
                 State == WebSocketState.CloseSent ||
-                State == WebSocketState.CloseReceived ||
+                State == WebSocketState.Closed ||
                 State == WebSocketState.Aborted,
                 $"Unexpected state {State}.");
 
-            // Wait until we've received a close response
-            byte[] closeBuffer = ArrayPool<byte>.Shared.Rent(MaxMessageHeaderLength + MaxControlPayloadLength);
-            try
+            // We only need to wait for a received close frame if we are in the CloseSent State. If we are in the Closed
+            // State then it means we already received a close frame. If we are in the Aborted State, then we should not
+            // wait for a close frame as per RFC 6455 Section 7.1.7 "Fail the WebSocket Connection".
+            if (State == WebSocketState.CloseSent)
             {
-                while (!_receivedCloseFrame)
+                // Wait until we've received a close response
+                byte[] closeBuffer = ArrayPool<byte>.Shared.Rent(MaxMessageHeaderLength + MaxControlPayloadLength);
+                try
                 {
-                    Debug.Assert(!Monitor.IsEntered(StateUpdateLock), $"{nameof(StateUpdateLock)} must never be held when acquiring {nameof(ReceiveAsyncLock)}");
-                    Task receiveTask;
-                    lock (ReceiveAsyncLock)
+                    while (!_receivedCloseFrame)
                     {
-                        // Now that we're holding the ReceiveAsyncLock, double-check that we've not yet received the close frame.
-                        // It could have been received between our check above and now due to a concurrent receive completing.
-                        if (_receivedCloseFrame)
+                        Debug.Assert(!Monitor.IsEntered(StateUpdateLock), $"{nameof(StateUpdateLock)} must never be held when acquiring {nameof(ReceiveAsyncLock)}");
+                        Task receiveTask;
+                        lock (ReceiveAsyncLock)
                         {
-                            break;
+                            // Now that we're holding the ReceiveAsyncLock, double-check that we've not yet received the close frame.
+                            // It could have been received between our check above and now due to a concurrent receive completing.
+                            if (_receivedCloseFrame)
+                            {
+                                break;
+                            }
+
+                            // We've not yet processed a received close frame, which means we need to wait for a received close to complete.
+                            // There may already be one in flight, in which case we want to just wait for that one rather than kicking off
+                            // another (we don't support concurrent receive operations).  We need to kick off a new receive if either we've
+                            // never issued a receive or if the last issued receive completed for reasons other than a close frame.  There is
+                            // a race condition here, e.g. if there's a in-flight receive that completes after we check, but that's fine: worst
+                            // case is we then await it, find that it's not what we need, and try again.
+                            receiveTask = _lastReceiveAsync;
+                            _lastReceiveAsync = receiveTask = ValidateAndReceiveAsync(receiveTask, closeBuffer, cancellationToken);
                         }
 
-                        // We've not yet processed a received close frame, which means we need to wait for a received close to complete.
-                        // There may already be one in flight, in which case we want to just wait for that one rather than kicking off
-                        // another (we don't support concurrent receive operations).  We need to kick off a new receive if either we've
-                        // never issued a receive or if the last issued receive completed for reasons other than a close frame.  There is
-                        // a race condition here, e.g. if there's a in-flight receive that completes after we check, but that's fine: worst
-                        // case is we then await it, find that it's not what we need, and try again.
-                        receiveTask = _lastReceiveAsync;
-                        _lastReceiveAsync = receiveTask = ValidateAndReceiveAsync(receiveTask, closeBuffer, cancellationToken);
+                        // Wait for whatever receive task we have.  We'll then loop around again to re-check our state.
+                        Debug.Assert(receiveTask != null);
+                        await receiveTask.ConfigureAwait(false);
                     }
-
-                    // Wait for whatever receive task we have.  We'll then loop around again to re-check our state.
-                    Debug.Assert(receiveTask != null);
-                    await receiveTask.ConfigureAwait(false);
                 }
-            }
-            finally
-            {
-                ArrayPool<byte>.Shared.Return(closeBuffer);
+                finally
+                {
+                    ArrayPool<byte>.Shared.Return(closeBuffer);
+                }
             }
 
             // We're closed.  Close the connection and update the status.

src/System.Net.HttpListener/tests/HttpListenerWebSocketTests.cs

@@ -172,7 +172,6 @@ public static IEnumerable<object[]> CloseStatus_Valid_TestData()
             yield return new object[] { WebSocketCloseStatus.MandatoryExtension, "StatusDescription", WebSocketCloseStatus.MandatoryExtension };
         }
 
-        [ActiveIssue(36963)]
         [ConditionalTheory(nameof(IsNotWindows7AndIsWindowsImplementation))] // [ActiveIssue(20396, TestPlatforms.AnyUnix)]
         [MemberData(nameof(CloseStatus_Valid_TestData))]
         public async Task CloseOutputAsync_HandshakeStartedFromClient_Success(WebSocketCloseStatus status, string statusDescription, WebSocketCloseStatus expectedCloseStatus)