Fix some connection woes

Author: dolfies

discord/gateway.py

@@ -220,6 +220,9 @@ def ack(self) -> None:
 
 
 class VoiceKeepAliveHandler(KeepAliveHandler):
+    if TYPE_CHECKING:
+        ws: DiscordVoiceWebSocket
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.recent_ack_latencies: deque[float] = deque(maxlen=20)
@@ -232,7 +235,10 @@ def __init__(self, *args, **kwargs):
     def get_payload(self) -> Dict[str, Any]:
         return {
             'op': self.ws.HEARTBEAT,
-            'd': int(time.time() * 1000),
+            'd': {
+                't': int(time.time() * 1000),
+                'seq_ack': self.ws.seq_ack,
+            },
         }
 
     def ack(self) -> None:
@@ -657,6 +663,9 @@ async def poll_event(self) -> None:
             if isinstance(e, CurlError):
                 reason = str(e)
 
+            if not socket.closed:
+                await socket.close(code or 4000, (reason or 'Unknown error').encode('utf-8'))
+
             _log.info(f'Gateway received close code {code} and reason {reason!r}.')
 
             if self._can_handle_close(code or None):
@@ -926,6 +935,7 @@ def __init__(
         self._keep_alive: Optional[VoiceKeepAliveHandler] = None
         self._close_code: Optional[int] = None
         self.secret_key: Optional[List[int]] = None
+        self.seq_ack: int = -1
         if hook:
             self._hook = hook  # type: ignore
 
@@ -956,6 +966,7 @@ async def resume(self) -> None:
                 'token': state.token,
                 'server_id': str(state.server_id),
                 'session_id': state.session_id,
+                'seq_ack': self.seq_ack,
             },
         }
         await self.send_as_json(payload)
@@ -980,14 +991,16 @@ async def from_connection_state(
         *,
         resume: bool = False,
         hook: Optional[Callable[..., Coroutine[Any, Any, Any]]] = None,
+        seq_ack: int = -1
     ) -> Self:
         """Creates a voice websocket for the :class:`VoiceClient`."""
-        gateway = f'wss://{state.endpoint}/?v=4'
+        gateway = f'wss://{state.endpoint}/?v=8'
         client = state.voice_client
         http = client._state.http
         socket = await http.ws_connect(gateway)
         ws = cls(socket, loop=client.loop, hook=hook)
         ws.gateway = gateway
+        ws.seq_ack = seq_ack
         ws._connection = state
         ws._max_heartbeat_timeout = 60.0
         ws.thread_id = threading.get_ident()
@@ -1040,6 +1053,7 @@ async def received_message(self, msg: Dict[str, Any]) -> None:
         _log.debug('Voice socket event: %s.', msg)
         op = msg['op']
         data = msg['d']  # According to Discord this key is always given
+        self.seq_ack = msg.get('seq', self.seq_ack)
 
         if op == self.READY:
             await self.initial_connection(data)

discord/voice_state.py

@@ -323,7 +323,7 @@ async def voice_server_update(self, data: VoiceServerUpdatePayload) -> None:
             )
             return
 
-        self.endpoint, _, _ = endpoint.rpartition(':')
+        self.endpoint = endpoint
         if self.endpoint.startswith('wss://'):
             # Just in case, strip it off since we're going to add it later
             self.endpoint = self.endpoint[6:]
@@ -603,7 +603,10 @@ async def _voice_disconnect(self) -> None:
         self._disconnected.clear()
 
     async def _connect_websocket(self, resume: bool) -> DiscordVoiceWebSocket:
-        ws = await DiscordVoiceWebSocket.from_connection_state(self, resume=resume, hook=self.hook)
+        seq_ack = -1
+        if self.ws is not MISSING:
+            seq_ack = self.ws.seq_ack
+        ws = await DiscordVoiceWebSocket.from_connection_state(self, resume=resume, hook=self.hook, seq_ack=seq_ack)
         self.state = ConnectionFlowState.websocket_connected
         return ws
 
@@ -632,18 +635,20 @@ async def _poll_voice_ws(self, reconnect: bool) -> None:
                     # The following close codes are undocumented so I will document them here.
                     # 1000 - normal closure (obviously)
                     # 4014 - we were externally disconnected (voice channel deleted, we were moved, etc)
-                    # 4015 - voice server has crashed
+                    # 4015 - voice server has crashed, we should resume
+                    # 4021 - rate limited, we should not reconnect
+                    # 4022 - call terminated, similar to 4014
 
                     code = getattr(exc, 'code', self.ws._close_code)
 
-                    if code in (1000, 4015):
+                    if code == 1000:
                         # Don't call disconnect a second time if the websocket closed from a disconnect call
                         if not self._expecting_disconnect:
                             _log.info('Disconnecting from voice normally, close code %d.', code)
                             await self.disconnect()
                         break
 
-                    if code == 4014:
+                    if code in (4014, 4022):
                         # We were disconnected by discord
                         # This condition is a race between the main ws event and the voice ws closing
                         if self._disconnected.is_set():
@@ -663,6 +668,32 @@ async def _poll_voice_ws(self, reconnect: bool) -> None:
                         else:
                             continue
 
+                    if code == 4021:
+                        _log.warning('We are being rate limited while trying to connect to voice. Disconnecting...')
+                        if self.state is not ConnectionFlowState.disconnected:
+                            await self.disconnect()
+                        break
+
+                    if exc.code == 4015:
+                        _log.info('Disconnected from voice, attempting a resume...')
+                        try:
+                            await self._connect(
+                                reconnect=reconnect,
+                                timeout=self.timeout,
+                                self_deaf=(self.self_voice_state or self).self_deaf,
+                                self_mute=(self.self_voice_state or self).self_mute,
+                                self_video=(self.self_voice_state or self).self_video,
+                                resume=True,
+                            )
+                        except asyncio.TimeoutError:
+                            _log.info('Could not resume the voice connection. Disconnecting...')
+                            if self.state is not ConnectionFlowState.disconnected:
+                                await self.disconnect()
+                            break
+                        else:
+                            _log.info('Successfully resumed voice connection.')
+                            continue
+
                     _log.debug(
                         'Not handling voice socket close code %s (reason: %s).',
                         code,