[release/8.0][browser] BrowserWebSocket.ReceiveAsync after server initiated close (#97002)

Author: pavelsavara

src/libraries/Common/src/System/Net/WebSockets/WebSocketValidate.cs

@@ -23,10 +23,15 @@ internal static partial class WebSocketValidate
         internal const int MaxDeflateWindowBits = 15;
 
         internal const int MaxControlFramePayloadLength = 123;
+#if TARGET_BROWSER
+        private const int ValidCloseStatusCodesFrom = 3000;
+        private const int ValidCloseStatusCodesTo = 4999;
+#else
         private const int CloseStatusCodeAbort = 1006;
         private const int CloseStatusCodeFailedTLSHandshake = 1015;
         private const int InvalidCloseStatusCodesFrom = 0;
         private const int InvalidCloseStatusCodesTo = 999;
+#endif
 
         // [0x21, 0x7E] except separators "()<>@,;:\\\"/[]?={} ".
         private static readonly SearchValues<char> s_validSubprotocolChars =
@@ -84,11 +89,15 @@ internal static void ValidateCloseStatus(WebSocketCloseStatus closeStatus, strin
             }
 
             int closeStatusCode = (int)closeStatus;
-
+#if TARGET_BROWSER
+            // as defined in https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close#code
+            if (closeStatus != WebSocketCloseStatus.NormalClosure && (closeStatusCode < ValidCloseStatusCodesFrom || closeStatusCode > ValidCloseStatusCodesTo))
+#else
             if ((closeStatusCode >= InvalidCloseStatusCodesFrom &&
                 closeStatusCode <= InvalidCloseStatusCodesTo) ||
                 closeStatusCode == CloseStatusCodeAbort ||
                 closeStatusCode == CloseStatusCodeFailedTLSHandshake)
+#endif
             {
                 // CloseStatus 1006 means Aborted - this will never appear on the wire and is reflected by calling WebSocket.Abort
                 throw new ArgumentException(SR.Format(SR.net_WebSockets_InvalidCloseStatusCode,

src/libraries/Common/tests/System/Net/Prerequisites/NetCoreServer/Handlers/EchoWebSocketHandler.cs

@@ -24,11 +24,11 @@ public static async Task InvokeAsync(HttpContext context)
 
             if (context.Request.QueryString.HasValue && context.Request.QueryString.Value.Contains("delay10sec"))
             {
-                Thread.Sleep(10000);
+                await Task.Delay(10000);
             }
             else if (context.Request.QueryString.HasValue && context.Request.QueryString.Value.Contains("delay20sec"))
             {
-                Thread.Sleep(20000);
+                await Task.Delay(20000);
             }
 
             try
@@ -124,14 +124,15 @@ await socket.CloseAsync(
                 }
 
                 bool sendMessage = false;
+                string receivedMessage = null;
                 if (receiveResult.MessageType == WebSocketMessageType.Text)
                 {
-                    string receivedMessage = Encoding.UTF8.GetString(receiveBuffer, 0, offset);
+                    receivedMessage = Encoding.UTF8.GetString(receiveBuffer, 0, offset);
                     if (receivedMessage == ".close")
                     {
                         await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, receivedMessage, CancellationToken.None);
                     }
-                    if (receivedMessage == ".shutdown")
+                    else if (receivedMessage == ".shutdown")
                     {
                         await socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, receivedMessage, CancellationToken.None);
                     }
@@ -161,6 +162,14 @@ await socket.SendAsync(
                             !replyWithPartialMessages,
                             CancellationToken.None);
                 }
+                if (receivedMessage == ".closeafter")
+                {
+                    await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, receivedMessage, CancellationToken.None);
+                }
+                else if (receivedMessage == ".shutdownafter")
+                {
+                    await socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, receivedMessage, CancellationToken.None);
+                }
             }
         }
     }

src/libraries/System.Net.WebSockets.Client/src/System/Net/WebSockets/BrowserWebSockets/BrowserWebSocket.cs

@@ -24,6 +24,8 @@ internal sealed class BrowserWebSocket : WebSocket
         private WebSocketState _state;
         private bool _disposed;
         private bool _aborted;
+        private bool _closeReceived;
+        private bool _closeSent;
         private int[] responseStatus = new int[3];
         private MemoryHandle? responseStatusHandle;
 
@@ -37,7 +39,7 @@ public override WebSocketState State
                 lock (_thisLock)
                 {
 #endif
-                    if (_innerWebSocket == null || _disposed || (_state != WebSocketState.Connecting && _state != WebSocketState.Open && _state != WebSocketState.CloseSent))
+                    if (_innerWebSocket == null || _disposed || _state == WebSocketState.Aborted || _state == WebSocketState.Closed)
                     {
                         return _state;
                     }
@@ -46,15 +48,9 @@ public override WebSocketState State
 #endif
 
 #if FEATURE_WASM_THREADS
-                return FastState = _innerWebSocket!.SynchronizationContext.Send(static (BrowserWebSocket self) =>
-                {
-                    lock (self._thisLock)
-                    {
-                        return GetReadyState(self._innerWebSocket!);
-                    } //lock
-                }, this);
+                return _innerWebSocket!.SynchronizationContext.Send(GetReadyState, this);
 #else
-                return FastState = GetReadyState(_innerWebSocket!);
+                return GetReadyState(this);
 #endif
             }
         }
@@ -148,7 +144,7 @@ public override Task SendAsync(ArraySegment<byte> buffer, WebSocketMessageType m
             ThrowIfDisposed();
 
             // fast check of previous _state instead of GetReadyState(), the readyState would be validated on JS side
-            if (FastState != WebSocketState.Open)
+            if (FastState != WebSocketState.Open && FastState != WebSocketState.CloseReceived)
             {
                 throw new InvalidOperationException(SR.net_WebSockets_NotConnected);
             }
@@ -240,7 +236,7 @@ public override Task CloseOutputAsync(WebSocketCloseStatus closeStatus, string?
                     {
                         throw new WebSocketException(WebSocketError.InvalidState, SR.Format(SR.net_WebSockets_InvalidState, state, "Connecting, Open, CloseSent, Aborted"));
                     }
-                    if(state != WebSocketState.Open && state != WebSocketState.Connecting && state != WebSocketState.Aborted)
+                    if (state == WebSocketState.CloseSent)
                     {
                         return Task.CompletedTask;
                     }
@@ -280,10 +276,6 @@ public override Task CloseAsync(WebSocketCloseStatus closeStatus, string? status
                     {
                         throw new WebSocketException(WebSocketError.InvalidState, SR.Format(SR.net_WebSockets_InvalidState, state, "Connecting, Open, CloseSent, Aborted"));
                     }
-                    if (state != WebSocketState.Open && state != WebSocketState.Connecting && state != WebSocketState.Aborted && state != WebSocketState.CloseSent)
-                    {
-                        return Task.CompletedTask;
-                    }
 
 #if FEATURE_WASM_THREADS
                     promise = CloseAsyncCore(closeStatus, statusDescription, state != WebSocketState.Aborted, cancellationToken);
@@ -387,12 +379,13 @@ private void CreateCore(Uri uri, List<string>? requestedSubProtocols)
                 string[]? subProtocols = requestedSubProtocols?.ToArray();
                 var onClose = (int code, string reason) =>
                 {
-                    _closeStatus = (WebSocketCloseStatus)code;
-                    _closeStatusDescription = reason;
 #if FEATURE_WASM_THREADS
                     lock (_thisLock)
                     {
 #endif
+                        _closeStatus = (WebSocketCloseStatus)code;
+                        _closeStatusDescription = reason;
+                        _closeReceived = true;
                         WebSocketState state = State;
                         if (state == WebSocketState.Connecting || state == WebSocketState.Open || state == WebSocketState.CloseSent)
                         {
@@ -545,15 +538,21 @@ private static WebSocketReceiveResult ConvertResponse(BrowserWebSocket self)
             WebSocketMessageType messageType = (WebSocketMessageType)self.responseStatus[typeIndex];
             if (messageType == WebSocketMessageType.Close)
             {
+                self._closeReceived = true;
+                self.FastState = self._closeSent ? WebSocketState.Closed : WebSocketState.CloseReceived;
                 return new WebSocketReceiveResult(self.responseStatus[countIndex], messageType, self.responseStatus[endIndex] != 0, self.CloseStatus, self.CloseStatusDescription);
             }
             return new WebSocketReceiveResult(self.responseStatus[countIndex], messageType, self.responseStatus[endIndex] != 0);
         }
 
         private async Task CloseAsyncCore(WebSocketCloseStatus closeStatus, string? statusDescription, bool waitForCloseReceived, CancellationToken cancellationToken)
         {
-            _closeStatus = closeStatus;
-            _closeStatusDescription = statusDescription;
+            if (!_closeReceived)
+            {
+                _closeStatus = closeStatus;
+                _closeStatusDescription = statusDescription;
+            }
+            _closeSent = true;
 
             var closeTask = BrowserInterop.WebSocketClose(_innerWebSocket!, (int)closeStatus, statusDescription, waitForCloseReceived) ?? Task.CompletedTask;
             await CancelationHelper(closeTask, cancellationToken, FastState).ConfigureAwait(true);
@@ -562,6 +561,10 @@ private async Task CloseAsyncCore(WebSocketCloseStatus closeStatus, string? stat
             lock (_thisLock)
             {
 #endif
+                if (waitForCloseReceived)
+                {
+                    _closeReceived = true;
+                }
                 var state = State;
                 if (state == WebSocketState.Open || state == WebSocketState.Connecting || state == WebSocketState.CloseSent)
                 {
@@ -614,18 +617,42 @@ private async Task CancelationHelper(Task jsTask, CancellationToken cancellation
             }
         }
 
-        private static WebSocketState GetReadyState(JSObject innerWebSocket)
+        private static WebSocketState GetReadyState(BrowserWebSocket self)
         {
-            var readyState = BrowserInterop.GetReadyState(innerWebSocket);
-            // https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/readyState
-            return readyState switch
-            {
-                0 => WebSocketState.Connecting, // 0 (CONNECTING)
-                1 => WebSocketState.Open, // 1 (OPEN)
-                2 => WebSocketState.CloseSent, // 2 (CLOSING)
-                3 => WebSocketState.Closed, // 3 (CLOSED)
-                _ => WebSocketState.None
-            };
+#if FEATURE_WASM_THREADS
+            lock (self._thisLock)
+            {
+#endif
+                var readyState = BrowserInterop.GetReadyState(self._innerWebSocket);
+                // https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/readyState
+                var st = readyState switch
+                {
+                    0 => WebSocketState.Connecting, // 0 (CONNECTING)
+                    1 => WebSocketState.Open, // 1 (OPEN)
+                    2 => WebSocketState.CloseSent, // 2 (CLOSING)
+                    3 => WebSocketState.Closed, // 3 (CLOSED)
+                    _ => WebSocketState.None
+                };
+                if (st == WebSocketState.Closed || st == WebSocketState.CloseSent)
+                {
+                    if (self._closeReceived && self._closeSent)
+                    {
+                        st = WebSocketState.Closed;
+                    }
+                    else if (self._closeReceived && !self._closeSent)
+                    {
+                        st = WebSocketState.CloseReceived;
+                    }
+                    else if (!self._closeReceived && self._closeSent)
+                    {
+                        st = WebSocketState.CloseSent;
+                    }
+                }
+                self.FastState = st;
+                return st;
+#if FEATURE_WASM_THREADS
+            } //lock
+#endif
         }
 
         #endregion