From 7da2048d1aa0fe25d63a993f973494d7be25a4eb Mon Sep 17 00:00:00 2001
From: Rapptz <rapptz@gmail.com>
Date: Fri, 5 Aug 2022 08:17:52 -0400
Subject: [PATCH] Add support for resume_gateway_url
---
discord/client.py | 17 +++++++++++++++--
discord/gateway.py | 18 ++++++++++++++++--
discord/shard.py | 1 +
discord/types/gateway.py | 1 +
4 files changed, 33 insertions(+), 4 deletions(-)
diff --git a/discord/client.py b/discord/client.py
index b8adcf525..517e4d97c 100644
--- a/discord/client.py
+++ b/discord/client.py
@@ -682,6 +682,8 @@ class Client:
_log.info('Got a request to %s the websocket.', e.op)
self.dispatch('disconnect')
ws_params.update(sequence=self.ws.sequence, resume=e.resume, session=self.ws.session_id)
+ if e.resume:
+ ws_params['gateway'] = self.ws.gateway
continue
except (
OSError,
@@ -705,7 +707,13 @@ class Client:
# If we get connection reset by peer then try to RESUME
if isinstance(exc, OSError) and exc.errno in (54, 10054):
- ws_params.update(sequence=self.ws.sequence, initial=False, resume=True, session=self.ws.session_id)
+ ws_params.update(
+ sequence=self.ws.sequence,
+ gateway=self.ws.gateway,
+ initial=False,
+ resume=True,
+ session=self.ws.session_id,
+ )
continue
# We should only get this when an unhandled close code happens,
@@ -725,7 +733,12 @@ class Client:
# Always try to RESUME the connection
# If the connection is not RESUME-able then the gateway will invalidate the session.
# This is apparently what the official Discord client does.
- ws_params.update(sequence=self.ws.sequence, resume=True, session=self.ws.session_id)
+ ws_params.update(
+ sequence=self.ws.sequence,
+ gateway=self.ws.gateway,
+ resume=True,
+ session=self.ws.session_id,
+ )
async def close(self) -> None:
"""|coro|
diff --git a/discord/gateway.py b/discord/gateway.py
index 3bb58af36..7e6e75537 100644
--- a/discord/gateway.py
+++ b/discord/gateway.py
@@ -291,6 +291,7 @@ class DiscordWebSocket:
_max_heartbeat_timeout: float
# fmt: off
+ DEFAULT_GATEWAY = 'wss://gateway.discord.gg/'
DISPATCH = 0
HEARTBEAT = 1
IDENTIFY = 2
@@ -350,13 +351,24 @@ class DiscordWebSocket:
session: Optional[str] = None,
sequence: Optional[int] = None,
resume: bool = False,
+ encoding: str = 'json',
+ zlib: bool = True,
) -> Self:
"""Creates a main websocket for Discord from a :class:`Client`.
This is for internal use only.
"""
- gateway = gateway or await client.http.get_gateway()
- socket = await client.http.ws_connect(gateway)
+ # Circular import
+ from .http import INTERNAL_API_VERSION
+
+ gateway = gateway or cls.DEFAULT_GATEWAY
+
+ if zlib:
+ url = f'{gateway}?v={INTERNAL_API_VERSION}&encoding={encoding}&compress=zlib-stream'
+ else:
+ url = f'{gateway}?v={INTERNAL_API_VERSION}&encoding={encoding}'
+
+ socket = await client.http.ws_connect(url)
ws = cls(socket, loop=client.loop)
# dynamically add attributes needed
@@ -533,6 +545,7 @@ class DiscordWebSocket:
self.sequence = None
self.session_id = None
+ self.gateway = self.DEFAULT_GATEWAY
_log.info('Shard ID %s session has been invalidated.', self.shard_id)
await self.close(code=1000)
raise ReconnectWebSocket(self.shard_id, resume=False)
@@ -543,6 +556,7 @@ class DiscordWebSocket:
if event == 'READY':
self.sequence = msg['s']
self.session_id = data['session_id']
+ self.gateway = data['resume_gateway_url']
_log.info('Shard ID %s has connected to Gateway (Session ID: %s).', self.shard_id, self.session_id)
elif event == 'RESUMED':
diff --git a/discord/shard.py b/discord/shard.py
index bd431c5f8..dfff2baa7 100644
--- a/discord/shard.py
+++ b/discord/shard.py
@@ -183,6 +183,7 @@ class Shard:
coro = DiscordWebSocket.from_client(
self._client,
resume=exc.resume,
+ gateway=None if not exc.resume else self.ws.gateway,
shard_id=self.id,
session=self.ws.session_id,
sequence=self.ws.sequence,
diff --git a/discord/types/gateway.py b/discord/types/gateway.py
index 4321a2a6a..f5b9fbdfc 100644
--- a/discord/types/gateway.py
+++ b/discord/types/gateway.py
@@ -66,6 +66,7 @@ class ReadyEvent(TypedDict):
user: User
guilds: List[UnavailableGuild]
session_id: str
+ resume_gateway_url: str
shard: List[int] # shard_id, num_shards
application: GatewayAppInfo