From 1b587b408fdbc6bb6323e858979f172cb7840099 Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Thu, 28 Jan 2021 22:30:35 -0500 Subject: [PATCH] Dispose of ClientWebSocket when it fails to connect Somehow when it was ported to .NET Core, the transition from connecting to connected was switched to being done before the connect rather than after. As a result, if the connection fails and the websocket is never initialized, subsequent use (misuse) ends up null ref'ing. This restores the .NET Framework behavior, which was to dispose of the instance if the connect fails; subsequent operations then throw ObjectDisposedException. --- .../System/Net/WebSockets/ClientWebSocket.cs | 18 +++++++++++++----- .../tests/ConnectTest.cs | 6 ++++++ 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/src/libraries/System.Net.WebSockets.Client/src/System/Net/WebSockets/ClientWebSocket.cs b/src/libraries/System.Net.WebSockets.Client/src/System/Net/WebSockets/ClientWebSocket.cs index 4d71e31863ce03..7159283212d601 100644 --- a/src/libraries/System.Net.WebSockets.Client/src/System/Net/WebSockets/ClientWebSocket.cs +++ b/src/libraries/System.Net.WebSockets.Client/src/System/Net/WebSockets/ClientWebSocket.cs @@ -82,17 +82,25 @@ public Task ConnectAsync(Uri uri, CancellationToken cancellationToken) return ConnectAsyncCore(uri, cancellationToken); } - private Task ConnectAsyncCore(Uri uri, CancellationToken cancellationToken) + private async Task ConnectAsyncCore(Uri uri, CancellationToken cancellationToken) { _innerWebSocket = new WebSocketHandle(); - // Change internal state to 'connected' to enable the other methods - if ((InternalState)Interlocked.CompareExchange(ref _state, (int)InternalState.Connected, (int)InternalState.Connecting) != InternalState.Connecting) + try + { + await _innerWebSocket.ConnectAsync(uri, cancellationToken, Options).ConfigureAwait(false); + } + catch { - return Task.FromException(new ObjectDisposedException(nameof(ClientWebSocket))); // Aborted/Disposed during connect. + Dispose(); + throw; } - return _innerWebSocket.ConnectAsync(uri, cancellationToken, Options); + if ((InternalState)Interlocked.CompareExchange(ref _state, (int)InternalState.Connected, (int)InternalState.Connecting) != InternalState.Connecting) + { + Debug.Assert(_state == (int)InternalState.Disposed); + throw new ObjectDisposedException(GetType().FullName); + } } public override Task SendAsync(ArraySegment buffer, WebSocketMessageType messageType, bool endOfMessage, CancellationToken cancellationToken) => diff --git a/src/libraries/System.Net.WebSockets.Client/tests/ConnectTest.cs b/src/libraries/System.Net.WebSockets.Client/tests/ConnectTest.cs index 68faae7d7408cf..d6f9512b15ce4c 100644 --- a/src/libraries/System.Net.WebSockets.Client/tests/ConnectTest.cs +++ b/src/libraries/System.Net.WebSockets.Client/tests/ConnectTest.cs @@ -33,6 +33,12 @@ public async Task ConnectAsync_NotWebSocketServer_ThrowsWebSocketExceptionWithMe } Assert.Equal(WebSocketState.Closed, cws.State); Assert.Equal(exceptionMessage, ex.Message); + + // Other operations throw after failed connect + await Assert.ThrowsAsync(() => cws.ReceiveAsync(new byte[1], default)); + await Assert.ThrowsAsync(() => cws.SendAsync(new byte[1], WebSocketMessageType.Binary, true, default)); + await Assert.ThrowsAsync(() => cws.CloseAsync(WebSocketCloseStatus.NormalClosure, null, default)); + await Assert.ThrowsAsync(() => cws.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, null, default)); } }