19
19
from sys import platform , version_info
20
20
from time import perf_counter
21
21
from typing import TYPE_CHECKING , Any , Dict , List , Optional , Tuple , Type , Union
22
+ from zlib import decompressobj
22
23
23
24
from aiohttp import ClientWebSocketResponse , WSMessage , WSMsgType
24
25
@@ -98,6 +99,7 @@ class WebSocketClient:
98
99
"__heartbeater" ,
99
100
"__shard" ,
100
101
"__presence" ,
102
+ "_zlib" ,
101
103
"_task" ,
102
104
"__heartbeat_event" ,
103
105
"__started" ,
@@ -148,7 +150,7 @@ def __init__(
148
150
149
151
self .__closed : Event = Event (loop = self ._loop ) if version_info < (3 , 10 ) else Event ()
150
152
self ._options : dict = {
151
- "max_msg_size" : 1024 ** 2 ,
153
+ "max_msg_size" : 0 ,
152
154
"timeout" : 60 ,
153
155
"autoclose" : False ,
154
156
"compress" : 0 ,
@@ -180,6 +182,8 @@ def __init__(
180
182
181
183
self .__stopping : Optional [Task ] = None
182
184
185
+ self ._zlib = decompressobj ()
186
+
183
187
@property
184
188
def latency (self ) -> float :
185
189
"""
@@ -827,6 +831,8 @@ async def _reconnect(self, to_resume: bool, code: Optional[int] = 1012) -> None:
827
831
828
832
self ._client = None
829
833
834
+ self ._zlib = decompressobj ()
835
+
830
836
# We need to check about existing heartbeater tasks for edge cases.
831
837
832
838
if self ._task :
@@ -863,6 +869,8 @@ async def __receive_packet(self, ignore_lock: bool = False) -> Optional[Dict[str
863
869
:rtype: Optional[Dict[str, Any]]
864
870
"""
865
871
872
+ buffer = bytearray ()
873
+
866
874
while True :
867
875
868
876
if not ignore_lock :
@@ -874,12 +882,9 @@ async def __receive_packet(self, ignore_lock: bool = False) -> Optional[Dict[str
874
882
if packet .type == WSMsgType .CLOSE :
875
883
log .debug (f"Disconnecting from gateway = { packet .data } ::{ packet .extra } " )
876
884
877
- if packet .data >= 4000 : # suppress 4001 because of weird presence errors
885
+ if packet .data >= 4000 :
878
886
# This means that the error code is 4000+, which may signify Discord-provided error codes.
879
887
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
-
883
888
raise LibraryException (packet .data )
884
889
885
890
if ignore_lock :
@@ -919,8 +924,20 @@ async def __receive_packet(self, ignore_lock: bool = False) -> Optional[Dict[str
919
924
if packet .data is None :
920
925
continue # We just loop it over because it could just be processing something.
921
926
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
+
922
939
try :
923
- msg = loads (packet . data )
940
+ _msg = loads (msg )
924
941
except Exception as e :
925
942
import traceback
926
943
@@ -929,9 +946,9 @@ async def __receive_packet(self, ignore_lock: bool = False) -> Optional[Dict[str
929
946
)
930
947
# There's an edge case when the packet's None... or some other deserialisation error.
931
948
# 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
933
950
934
- return msg
951
+ return _msg
935
952
936
953
async def _send_packet (self , data : Dict [str , Any ]) -> None :
937
954
"""
@@ -972,10 +989,11 @@ async def __identify(
972
989
"token" : self ._http .token ,
973
990
"intents" : self ._intents .value ,
974
991
"properties" : {
975
- "$ os" : platform ,
976
- "$ browser" : "interactions.py" ,
977
- "$ device" : "interactions.py" ,
992
+ "os" : platform ,
993
+ "browser" : "interactions.py" ,
994
+ "device" : "interactions.py" ,
978
995
},
996
+ "compress" : True ,
979
997
},
980
998
}
981
999
@@ -1037,5 +1055,4 @@ async def close(self) -> None:
1037
1055
"""
1038
1056
if self ._client :
1039
1057
await self ._client .close ()
1040
-
1041
1058
self .__closed .set ()
0 commit comments