Allow OperationCanceledExceptions to propagate out of ClientWebSocket.ConnectAsync

They're currently being wrapped in WebSocketExceptions, but they should be allowed to escape unwrapped.

Author: stephentoub

src/System.Net.WebSockets.Client/src/System/Net/WebSockets/WebSocketHandle.Managed.cs

@@ -223,7 +223,7 @@ public async Task ConnectAsyncCore(Uri uri, CancellationToken cancellationToken,
                 Abort();
                 response?.Dispose();
 
-                if (exc is WebSocketException)
+                if (exc is WebSocketException || exc is OperationCanceledException)
                 {
                     throw;
                 }

src/System.Net.WebSockets.Client/tests/ConnectTest.cs

@@ -267,5 +267,43 @@ public async Task ConnectAndCloseAsync_UseProxyServer_ExpectedClosedState(Uri se
                 Assert.Equal(expectedCloseStatusDescription, cws.CloseStatusDescription);
             }
         }
+
+        [ConditionalFact(nameof(WebSocketsSupported))]
+        public async Task ConnectAsync_CancellationRequestedBeforeConnect_ThrowsOperationCanceledException()
+        {
+            using (var clientSocket = new ClientWebSocket())
+            {
+                var cts = new CancellationTokenSource();
+                cts.Cancel();
+                Task t = clientSocket.ConnectAsync(new Uri("ws://" + Guid.NewGuid().ToString("N")), cts.Token);
+                await Assert.ThrowsAnyAsync<OperationCanceledException>(() => t);
+            }
+        }
+
+        [ConditionalFact(nameof(WebSocketsSupported))]
+        public async Task ConnectAsync_CancellationRequestedAfterConnect_ThrowsOperationCanceledException()
+        {
+            var releaseServer = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            await LoopbackServer.CreateClientAndServerAsync(async uri =>
+            {
+                var clientSocket = new ClientWebSocket();
+                try
+                {
+                    var cts = new CancellationTokenSource();
+                    Task t = clientSocket.ConnectAsync(uri, cts.Token);
+                    Assert.False(t.IsCompleted);
+                    cts.Cancel();
+                    await Assert.ThrowsAnyAsync<OperationCanceledException>(() => t);
+                }
+                finally
+                {
+                    releaseServer.SetResult(true);
+                    clientSocket.Dispose();
+                }
+            }, server => server.AcceptConnectionAsync(async connection =>
+            {
+                await releaseServer.Task;
+            }), new LoopbackServer.Options { WebSocketEndpoint = true });
+        }
     }
 }