From ebd1d77ccfe1316258f349094369ed662b0abbff Mon Sep 17 00:00:00 2001
From: RogueException <rogueexception@gmail.com>
Date: Thu, 31 Dec 2015 15:51:28 -0400
Subject: [PATCH] Minor WS disconnect logic changes

---
 .../Net/WebSockets/WS4NetEngine.cs            | 37 ++++++++++++-------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/src/Discord.Net/Net/WebSockets/WS4NetEngine.cs b/src/Discord.Net/Net/WebSockets/WS4NetEngine.cs
index c419af0b4..4c1813501 100644
--- a/src/Discord.Net/Net/WebSockets/WS4NetEngine.cs
+++ b/src/Discord.Net/Net/WebSockets/WS4NetEngine.cs
@@ -36,21 +36,30 @@ namespace Discord.Net.WebSockets
 
         public Task Connect(string host, CancellationToken cancelToken)
         {
-            _webSocket = new WebSocketClient(host);
-            _webSocket.EnableAutoSendPing = false;
-            _webSocket.NoDelay = true;
-            _webSocket.Proxy = null;
+            try
+            {
+                _webSocket = new WebSocketClient(host);
+                _webSocket.EnableAutoSendPing = false;
+                _webSocket.NoDelay = true;
+                _webSocket.Proxy = null;
 
-            _webSocket.DataReceived += OnWebSocketBinary;
-            _webSocket.MessageReceived += OnWebSocketText;
-            _webSocket.Error += OnWebSocketError;
-            _webSocket.Closed += OnWebSocketClosed;
-            _webSocket.Opened += OnWebSocketOpened;
+                _webSocket.DataReceived += OnWebSocketBinary;
+                _webSocket.MessageReceived += OnWebSocketText;
+                _webSocket.Error += OnWebSocketError;
+                _webSocket.Closed += OnWebSocketClosed;
+                _webSocket.Opened += OnWebSocketOpened;
 
-            _waitUntilConnect.Reset();
-            _webSocket.Open();
-            _waitUntilConnect.Wait(cancelToken);
-            _taskManager.ThrowException(); //In case our connection failed
+                _waitUntilConnect.Reset();
+                _waitUntilDisconnect.Reset();
+                _webSocket.Open();
+                _waitUntilConnect.Wait(cancelToken);
+                _taskManager.ThrowException(); //In case our connection failed
+            }
+            catch
+            {
+                _waitUntilDisconnect.Set();
+                throw;
+            }
             return TaskHelper.CompletedTask;
         }
 
@@ -96,8 +105,8 @@ namespace Discord.Net.WebSockets
         }
         private void OnWebSocketOpened(object sender, EventArgs e)
         {
-            _waitUntilDisconnect.Reset();
             _waitUntilConnect.Set();
+            _waitUntilDisconnect.Reset();
         }
         private void OnWebSocketText(object sender, MessageReceivedEventArgs e)
             => OnTextMessage(e.Message);