[browser][ws] Fix WebSocket state propagation delay in ReceiveAsync

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

@@ -144,6 +144,18 @@ await socket.CloseAsync(
                     {
                         await Task.Delay(5000);
                     }
+                    else if (receivedMessage == ".receiveMessageAfterClose")
+                    {
+                        byte[] buffer = new byte[1024];
+                        string message = $"{receivedMessage} {DateTime.Now.ToString("HH:mm:ss")}";
+                        buffer = System.Text.Encoding.UTF8.GetBytes(message);
+                        await socket.SendAsync(
+                            new ArraySegment<byte>(buffer, 0, message.Length),
+                            WebSocketMessageType.Text,
+                            true,
+                            CancellationToken.None);
+                        await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, receivedMessage, CancellationToken.None);
+                    }
                     else if (socket.State == WebSocketState.Open)
                     {
                         sendMessage = true;

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

@@ -53,15 +53,13 @@ public static int GetReadyState(JSObject? webSocket)
                 return -1;
             }
 
-            int? readyState = webSocket.GetPropertyAsInt32("readyState");
-            if (!readyState.HasValue)
-            {
-                return -1;
-            }
-
-            return readyState.Value;
+            return BrowserInterop.WebSocketGetState(webSocket);
         }
 
+        [JSImport("INTERNAL.ws_get_state")]
+        public static partial int WebSocketGetState(
+            JSObject webSocket);
+
         [JSImport("INTERNAL.ws_wasm_create")]
         public static partial JSObject WebSocketCreate(
             string uri,

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

@@ -8,6 +8,7 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Linq;
 
 using Xunit;
 using Xunit.Abstractions;
@@ -362,6 +363,43 @@ await cws.SendAsync(
             }
         }
 
+        public static IEnumerable<object[]> EchoServersSyncState =>
+            EchoServers.SelectMany(server => new List<object[]>
+            {
+                new object[] { server[0], true },
+                new object[] { server[0], false }
+            });
+
+        [ActiveIssue("https://github.com/dotnet/runtime/issues/28957", typeof(PlatformDetection), nameof(PlatformDetection.IsNotBrowser))]
+        [ConditionalTheory(nameof(WebSocketsSupported)), MemberData(nameof(EchoServersSyncState))]
+        public async Task CloseOutputAsync_ServerInitiated_CanReceiveAfterClose(Uri server, bool syncState)
+        {
+            using (ClientWebSocket cws = await GetConnectedWebSocket(server, TimeOutMilliseconds, _output))
+            {
+                var cts = new CancellationTokenSource(TimeOutMilliseconds);
+                await cws.SendAsync(
+                    WebSocketData.GetBufferFromText(".receiveMessageAfterClose"),
+                    WebSocketMessageType.Text,
+                    true,
+                    cts.Token);
+
+                await Task.Delay(2000);
+
+                if (syncState)
+                {
+                    var state = cws.State;
+                    Assert.Equal(WebSocketState.Open, state);
+                    // should be able to receive after this sync
+                }
+
+                var recvBuffer = new ArraySegment<byte>(new byte[1024]);
+                WebSocketReceiveResult recvResult = await cws.ReceiveAsync(recvBuffer, cts.Token);
+                var message = Encoding.UTF8.GetString(recvBuffer.ToArray(), 0, recvResult.Count);
+
+                Assert.Contains(".receiveMessageAfterClose", message);
+            }
+        }
+
         [OuterLoop("Uses external servers", typeof(PlatformDetection), nameof(PlatformDetection.LocalEchoServerIsNotAvailable))]
         [ConditionalTheory(nameof(WebSocketsSupported)), MemberData(nameof(EchoServers))]
         public async Task CloseOutputAsync_CloseDescriptionIsNull_Success(Uri server)

src/mono/browser/runtime/exports-internal.ts

@@ -10,7 +10,7 @@ import { http_wasm_supports_streaming_request, http_wasm_supports_streaming_resp
 import { exportedRuntimeAPI, Module, runtimeHelpers } from "./globals";
 import { get_property, set_property, has_property, get_typeof_property, get_global_this, dynamic_import } from "./invoke-js";
 import { mono_wasm_stringify_as_error_with_stack } from "./logging";
-import { ws_wasm_create, ws_wasm_open, ws_wasm_send, ws_wasm_receive, ws_wasm_close, ws_wasm_abort } from "./web-socket";
+import { ws_wasm_create, ws_wasm_open, ws_wasm_send, ws_wasm_receive, ws_wasm_close, ws_wasm_abort, ws_get_state } from "./web-socket";
 import { mono_wasm_get_loaded_files } from "./assets";
 import { jiterpreter_dump_stats } from "./jiterpreter";
 import { interp_pgo_load_data, interp_pgo_save_data } from "./interp-pgo";
@@ -71,6 +71,7 @@ export function export_internal(): any {
         ws_wasm_receive,
         ws_wasm_close,
         ws_wasm_abort,
+        ws_get_state,
 
         // BrowserHttpHandler
         http_wasm_supports_streaming_request,

src/mono/browser/runtime/web-socket.ts

@@ -44,6 +44,17 @@ function verifyEnvironment() {
     }
 }
 
+export function ws_get_state(ws: WebSocketExtension) : number
+{
+    if (ws.readyState != WebSocket.CLOSED)
+        return ws.readyState ?? -1;
+    const receive_event_queue = ws[wasm_ws_pending_receive_event_queue];
+    const queued_events_count = receive_event_queue.getLength();
+    if (queued_events_count == 0)
+        return ws.readyState ?? -1;
+    return WebSocket.OPEN;
+}
+
 export function ws_wasm_create(uri: string, sub_protocols: string[] | null, receive_status_ptr: VoidPtr): WebSocketExtension {
     verifyEnvironment();
     assert_js_interop();