1 files changed, 35 insertions, 0 deletions
diff --git a/code/cogs/socketfix.py b/code/cogs/socketfix.py
new file mode 100644
index 0000000..9dfbf29
--- /dev/null
+++ b/code/cogs/socketfix.py
@@ -0,0 +1,35 @@
+import zlib
+import discord
+
+from discord.ext import commands
+
+class SocketFix(commands.Cog):
+    def __init__(self, bot):
+        self.bot = bot
+
+        self._zlib = zlib.decompressobj()
+        self._buffer = bytearray()
+
+    @commands.Cog.listener()
+    async def on_socket_raw_receive(self, msg):
+        if type(msg) is bytes:
+            self._buffer.extend(msg)
+
+            if len(msg) < 4 or msg[-4:] != b'\x00\x00\xff\xff':
+                return
+
+            try:
+                msg = self._zlib.decompress(self._buffer)
+            except Exception:
+                self._buffer = bytearray()  # Reset buffer on fail just in case...
+                return
+
+            msg = msg.decode('utf-8')
+            self._buffer = bytearray()
+
+        msg = discord.utils._from_json(msg)
+        self.bot.dispatch('on_socket_response', msg)
+
+
+async def setup(bot):
+    await bot.add_cog(SocketFix(bot))
+\ No newline at end of file