Add timeouts for websocket initial connections.

In DiscordWebSocket and DiscordVoiceWebsocket the from_client factory
methods can hang indefintely on all websocket ops. We set static
timeouts and attempt to reconnect if we time out.

Additionally, do not return from DiscordVoiceWebSocket.from_client
until we set up the keep alive thread, which is created after we
receive the READY payload.

1 files changed, 44 insertions, 5 deletions

diff --git a/discord/gateway.py b/discord/gateway.py
index ed9e06f8..da02d119 100644
--- a/discord/gateway.py
+++ b/discord/gateway.py
@@ -55,6 +55,16 @@ class ResumeWebSocket(Exception):
 
 EventListener = namedtuple('EventListener', 'predicate event result future')
 
[email protected]
+def _ensure_coroutine_connect(gateway, *, loop, klass):
+    # In 3.5+ websockets.connect does not return a coroutine, but an awaitable.
+    # The problem is that in 3.5.0 and in some cases 3.5.1, asyncio.ensure_future and
+    # by proxy, asyncio.wait_for, do not accept awaitables, but rather futures or coroutines.
+    # By wrapping it up into this function we ensure that it's in a coroutine and not an awaitable
+    # even for 3.5.0 users.
+    ws = yield from websockets.connect(gateway, loop=loop, klass=klass)
+    return ws
+
 class KeepAliveHandler(threading.Thread):
     def __init__(self, *args, **kwargs):
         ws = kwargs.pop('ws', None)
@@ -71,7 +81,7 @@ class KeepAliveHandler(threading.Thread):
         while not self._stop_ev.wait(self.interval):
             if self._last_ack + 2 * self.interval < time.time():
                 log.warn("We have stopped responding to the gateway.")
-                coro = self.ws.close(1006)
+                coro = self.ws.close(1001)
                 f = compat.run_coroutine_threadsafe(coro, loop=self.ws.loop)
 
                 try:
@@ -191,7 +201,13 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
         This is for internal use only.
         """
         gateway = yield from client.http.get_gateway()
-        ws = yield from websockets.connect(gateway, loop=client.loop, klass=cls)
+        try:
+            ws = yield from asyncio.wait_for(
+                    _ensure_coroutine_connect(gateway, loop=client.loop, klass=cls),
+                    timeout=60, loop=client.loop)
+        except asyncio.TimeoutError:
+            log.warn('timed out waiting for client connect')
+            return (yield from cls.from_client(client, resume=resume))
 
         # dynamically add attributes needed
         ws.token = client.http.token
@@ -206,7 +222,12 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
         log.info('Created websocket connected to {}'.format(gateway))
 
         # poll event for OP Hello
-        yield from ws.poll_event()
+        try:
+            yield from asyncio.wait_for(ws.poll_event(), timeout=60, loop=client.loop)
+        except asyncio.TimeoutError:
+            log.warn("timed out waiting for client HELLO")
+            yield from ws.close(1001)
+            return (yield from cls.from_client(client, resume=resume))
 
         if not resume:
             yield from ws.identify()
@@ -219,7 +240,7 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
             yield from ws.ensure_open()
         except websockets.exceptions.ConnectionClosed:
             # ws got closed so let's just do a regular IDENTIFY connect.
-            log.info('RESUME failure.')
+            log.warn('RESUME failure.')
             return (yield from cls.from_client(client))
         else:
             return ws
@@ -525,6 +546,7 @@ class DiscordVoiceWebSocket(websockets.client.WebSocketClientProtocol):
     HEARTBEAT           = 3
     SESSION_DESCRIPTION = 4
     SPEAKING            = 5
+    HELLO               = 8
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -540,7 +562,14 @@ class DiscordVoiceWebSocket(websockets.client.WebSocketClientProtocol):
     def from_client(cls, client):
         """Creates a voice websocket for the :class:`VoiceClient`."""
         gateway = 'wss://' + client.endpoint
-        ws = yield from websockets.connect(gateway, loop=client.loop, klass=cls)
+        try:
+            ws = yield from asyncio.wait_for(
+                    _ensure_coroutine_connect(gateway, loop=client.loop, klass=cls),
+                    timeout=60, loop=client.loop)
+        except asyncio.TimeoutError:
+            log.warn("timed out waiting for voice client connect")
+            return (yield from cls.from_client(client))
+
         ws.gateway = gateway
         ws._connection = client
 
@@ -555,6 +584,16 @@ class DiscordVoiceWebSocket(websockets.client.WebSocketClientProtocol):
         }
 
         yield from ws.send_as_json(identify)
+
+        try:
+            # Wait until we have processed READY and keep alive is running
+            while not ws._keep_alive:
+                yield from asyncio.wait_for(ws.poll_event(), timeout=60, loop=client.loop)
+        except asyncio.TimeoutError:
+            log.warn("timed out waiting for voice client READY")
+            yield from ws.close(1001)
+            return (yield from cls.from_client(client))
+
         return ws
 
     @asyncio.coroutine