Implement zlib streaming for the gateway.

discord/gateway.py

@@ -186,6 +186,8 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
         # ws related stuff
         self.session_id = None
         self.sequence = None
+        self._zlib = zlib.decompressobj()
+        self._buffer = bytearray()
 
     @classmethod
     @asyncio.coroutine
@@ -312,8 +314,18 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
         self._dispatch('socket_raw_receive', msg)
 
         if isinstance(msg, bytes):
-            msg = zlib.decompress(msg, 15, 10490000) # This is 10 MiB
-            msg = msg.decode('utf-8')
+            self._buffer.extend(msg)
+
+            if len(msg) >= 4:
+                suffix = int.from_bytes(msg[-4:], byteorder='big')
+                if suffix == 0xFFFF:
+                    msg = self._zlib.decompress(self._buffer)
+                    msg = msg.decode('utf-8')
+                    self._buffer = bytearray()
+                else:
+                    return
+            else:
+                return
 
         msg = json.loads(msg)
 

discord/http.py

@@ -739,21 +739,29 @@ class HTTPClient:
         return self.request(Route('GET', '/oauth2/applications/@me'))
 
     @asyncio.coroutine
-    def get_gateway(self):
+    def get_gateway(self, *, encoding='json', v=6, zlib=True):
         try:
             data = yield from self.request(Route('GET', '/gateway'))
         except HTTPException as e:
             raise GatewayNotFound() from e
-        return data.get('url') + '?encoding=json&v=6'
+        if zlib:
+            value = '{0}?encoding={1}&v={2}&compress=zlib-stream'
+        else:
+            value = '{0}?encoding={1}&v={2}'
+        return value.format(data['url'], encoding, v)
 
     @asyncio.coroutine
-    def get_bot_gateway(self):
+    def get_bot_gateway(self, *, encoding='json', v=6, zlib=True):
         try:
             data = yield from self.request(Route('GET', '/gateway/bot'))
         except HTTPException as e:
             raise GatewayNotFound() from e
+
+        if zlib:
+            value = '{0}?encoding={1}&v={2}&compress=zlib-stream'
         else:
-            return data['shards'], data['url'] + '?encoding=json&v=6'
+            value = '{0}?encoding={1}&v={2}'
+        return data['shards'], value.format(data['url'], encoding, v)
 
     def get_user_info(self, user_id):
         return self.request(Route('GET', '/users/{user_id}', user_id=user_id))