[browser] WS cancelation fix (#99522)

Co-authored-by: Ilona Tomkowicz <32700855+ilonatommy@users.noreply.github.com>
dotnet · Mar 12, 2024 · 19d7768 · 19d7768
1 parent 738d238
commit 19d7768
Showing 1 changed file with 22 additions and 12 deletions.
diff --git a/...tem.Net.WebSockets.Client/src/System/Net/WebSockets/BrowserWebSockets/BrowserWebSocket.cs b/...tem.Net.WebSockets.Client/src/System/Net/WebSockets/BrowserWebSockets/BrowserWebSocket.cs
@@ -135,9 +135,10 @@ public override string? SubProtocol
 
         internal Task ConnectAsync(Uri uri, List<string>? requestedSubProtocols, CancellationToken cancellationToken)
         {
+            AbortIfCancelationRequested(cancellationToken);
+
             lock (_lockObject)
             {
-                cancellationToken.ThrowIfCancellationRequested();
                 ThrowIfDisposed();
 
                 if (FastState != WebSocketState.None)
@@ -184,7 +185,6 @@ public override Task CloseOutputAsync(WebSocketCloseStatus closeStatus, string?
             // this validation should be synchronous
             WebSocketValidate.ValidateCloseStatus(closeStatus, statusDescription);
 
-            cancellationToken.ThrowIfCancellationRequested();
             ThrowIfDisposed();
 
             return CloseAsyncCore(closeStatus, statusDescription, false, cancellationToken);
@@ -195,7 +195,6 @@ public override Task CloseAsync(WebSocketCloseStatus closeStatus, string? status
             // this validation should be synchronous
             WebSocketValidate.ValidateCloseStatus(closeStatus, statusDescription);
 
-            cancellationToken.ThrowIfCancellationRequested();
             ThrowIfDisposed();
 
             return CloseAsyncCore(closeStatus, statusDescription, true, cancellationToken);
@@ -284,6 +283,17 @@ private void ThrowIfDisposed()
             } // lock
         }
 
+        private void AbortIfCancelationRequested(CancellationToken cancellationToken)
+        {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                lock (_lockObject)
+                {
+                    Abort();
+                } // lock
+                cancellationToken.ThrowIfCancellationRequested();
+            }
+        }
 
         private void CreateCore(Uri uri, List<string>? requestedSubProtocols)
         {
@@ -366,14 +376,14 @@ private async Task SendAsyncCore(ArraySegment<byte> buffer, WebSocketMessageType
             {
                 lock (_lockObject)
                 {
-                    cancellationToken.ThrowIfCancellationRequested();
                     ThrowIfDisposed();
 
                     previousState = FastState;
                     if (previousState != WebSocketState.Open && previousState != WebSocketState.CloseReceived)
                     {
                         throw new InvalidOperationException(SR.net_WebSockets_NotConnected);
                     }
+                    AbortIfCancelationRequested(cancellationToken);
 
                     if (buffer.Count == 0)
                     {
@@ -416,14 +426,14 @@ private async Task<WebSocketReceiveResult> ReceiveAsyncCore(ArraySegment<byte> b
             {
                 lock (_lockObject)
                 {
-                    cancellationToken.ThrowIfCancellationRequested();
                     ThrowIfDisposed();
 
                     previousState = FastState;
                     if (previousState != WebSocketState.Open && previousState != WebSocketState.CloseSent)
                     {
                         throw new WebSocketException(WebSocketError.InvalidState, SR.Format(SR.net_WebSockets_InvalidState, previousState, "Open, CloseSent"));
                     }
+                    AbortIfCancelationRequested(cancellationToken);
 
                     Memory<byte> bufferMemory = buffer.AsMemory();
                     pinBuffer = bufferMemory.Pin();
@@ -502,22 +512,22 @@ private async Task CloseAsyncCore(WebSocketCloseStatus closeStatus, string? stat
             WebSocketState previousState;
             lock (_lockObject)
             {
-                cancellationToken.ThrowIfCancellationRequested();
-
                 previousState = FastState;
                 if (_aborted)
                 {
                     return;
                 }
+                if (previousState == WebSocketState.None || previousState == WebSocketState.Closed)
+                {
+                    throw new WebSocketException(WebSocketError.InvalidState, SR.Format(SR.net_WebSockets_InvalidState, previousState, "Connecting, Open, CloseSent, CloseReceived, Aborted"));
+                }
+                AbortIfCancelationRequested(cancellationToken);
+
                 if (!_closeReceived)
                 {
                     _closeStatus = closeStatus;
                     _closeStatusDescription = statusDescription;
                 }
-                if (previousState == WebSocketState.None || previousState == WebSocketState.Closed)
-                {
-                    throw new WebSocketException(WebSocketError.InvalidState, SR.Format(SR.net_WebSockets_InvalidState, previousState, "Connecting, Open, CloseSent, Aborted"));
-                }
 
                 _closeSent = true;
 
@@ -544,7 +554,6 @@ private async Task CancellationHelper(Task promise, CancellationToken cancellati
         {
             try
             {
-                cancellationToken.ThrowIfCancellationRequested();
                 if (promise.IsCompletedSuccessfully)
                 {
                     disposable?.Dispose();
@@ -556,6 +565,7 @@ private async Task CancellationHelper(Task promise, CancellationToken cancellati
                     await promise.ConfigureAwait(false);
                     return;
                 }
+                AbortIfCancelationRequested(cancellationToken);
 
                 using (var receiveRegistration = cancellationToken.Register(static s =>
                 {