Skip to content

Commit c68f50c

Browse files
authored
fix: Fix constant gateway resume/invalid session loop. (#1070)
* fix: Added zlib + compression, disabled msg size limit * fix: Removed `$` on identify call for future-proofing.
1 parent 7397c7c commit c68f50c

File tree

1 file changed

+29
-12
lines changed

1 file changed

+29
-12
lines changed

interactions/api/gateway/client.py

Lines changed: 29 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from sys import platform, version_info
2020
from time import perf_counter
2121
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
22+
from zlib import decompressobj
2223

2324
from aiohttp import ClientWebSocketResponse, WSMessage, WSMsgType
2425

@@ -98,6 +99,7 @@ class WebSocketClient:
9899
"__heartbeater",
99100
"__shard",
100101
"__presence",
102+
"_zlib",
101103
"_task",
102104
"__heartbeat_event",
103105
"__started",
@@ -148,7 +150,7 @@ def __init__(
148150

149151
self.__closed: Event = Event(loop=self._loop) if version_info < (3, 10) else Event()
150152
self._options: dict = {
151-
"max_msg_size": 1024**2,
153+
"max_msg_size": 0,
152154
"timeout": 60,
153155
"autoclose": False,
154156
"compress": 0,
@@ -180,6 +182,8 @@ def __init__(
180182

181183
self.__stopping: Optional[Task] = None
182184

185+
self._zlib = decompressobj()
186+
183187
@property
184188
def latency(self) -> float:
185189
"""
@@ -827,6 +831,8 @@ async def _reconnect(self, to_resume: bool, code: Optional[int] = 1012) -> None:
827831

828832
self._client = None
829833

834+
self._zlib = decompressobj()
835+
830836
# We need to check about existing heartbeater tasks for edge cases.
831837

832838
if self._task:
@@ -863,6 +869,8 @@ async def __receive_packet(self, ignore_lock: bool = False) -> Optional[Dict[str
863869
:rtype: Optional[Dict[str, Any]]
864870
"""
865871

872+
buffer = bytearray()
873+
866874
while True:
867875

868876
if not ignore_lock:
@@ -874,12 +882,9 @@ async def __receive_packet(self, ignore_lock: bool = False) -> Optional[Dict[str
874882
if packet.type == WSMsgType.CLOSE:
875883
log.debug(f"Disconnecting from gateway = {packet.data}::{packet.extra}")
876884

877-
if packet.data >= 4000: # suppress 4001 because of weird presence errors
885+
if packet.data >= 4000:
878886
# This means that the error code is 4000+, which may signify Discord-provided error codes.
879887

880-
# However, we suppress 4001 because of weird presence errors with change_presence
881-
# The payload is correct, and the presence object persists. /shrug
882-
883888
raise LibraryException(packet.data)
884889

885890
if ignore_lock:
@@ -919,8 +924,20 @@ async def __receive_packet(self, ignore_lock: bool = False) -> Optional[Dict[str
919924
if packet.data is None:
920925
continue # We just loop it over because it could just be processing something.
921926

927+
if isinstance(packet.data, bytes):
928+
buffer.extend(packet.data)
929+
930+
if len(packet.data) < 4 or packet.data[-4:] != b"\x00\x00\xff\xff":
931+
# buffer isn't done we need to wait
932+
continue
933+
934+
msg = self._zlib.decompress(buffer)
935+
msg = msg.decode("utf-8")
936+
else:
937+
msg = packet.data
938+
922939
try:
923-
msg = loads(packet.data)
940+
_msg = loads(msg)
924941
except Exception as e:
925942
import traceback
926943

@@ -929,9 +946,9 @@ async def __receive_packet(self, ignore_lock: bool = False) -> Optional[Dict[str
929946
)
930947
# There's an edge case when the packet's None... or some other deserialisation error.
931948
# Instead of raising an exception, we just log it to debug, so it doesn't annoy end user's console logs.
932-
msg = None
949+
_msg = None
933950

934-
return msg
951+
return _msg
935952

936953
async def _send_packet(self, data: Dict[str, Any]) -> None:
937954
"""
@@ -972,10 +989,11 @@ async def __identify(
972989
"token": self._http.token,
973990
"intents": self._intents.value,
974991
"properties": {
975-
"$os": platform,
976-
"$browser": "interactions.py",
977-
"$device": "interactions.py",
992+
"os": platform,
993+
"browser": "interactions.py",
994+
"device": "interactions.py",
978995
},
996+
"compress": True,
979997
},
980998
}
981999

@@ -1037,5 +1055,4 @@ async def close(self) -> None:
10371055
"""
10381056
if self._client:
10391057
await self._client.close()
1040-
10411058
self.__closed.set()

0 commit comments

Comments
 (0)