Fix zlib-stream

sledgang · Aug 13, 2020 · 85e1cad · 85e1cad
1 parent f36997f
commit 85e1cad
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 3 deletions.
diff --git a/source/droid/gateway/compression.d b/source/droid/gateway/compression.d
@@ -17,20 +17,31 @@ class Decompressor {
 }
 
 class ZLibStream : Decompressor {
-    const ulong[] ZLIB_SUFFIX = [0x0, 0x0, 0xFF, 0xFF];
+    const ubyte[] ZLIB_SUFFIX = [0x0, 0x0, 0xFF, 0xFF];
     UnCompress decompressor;
 
+    ubyte[] buffer;
+
     this() {
         decompressor = new UnCompress(HeaderFormat.deflate);
     }
 
+    /*
+     * Reads a zlib stream from the websocket
+     * This will append the data to a buffer,
+     * returning nothing if the data is not a full zlib frame
+     * otherwise, returning the decompressed string.
+     */
     override string read(ubyte[] data) {
+        buffer ~= data;
+
         if (data[$-4..$] != ZLIB_SUFFIX) {
-            throw new DroidException("ZLib-Stream compression enabled but invalid data was recieved!");
+            return "";
         }
 
-        string decompressed = to!string(decompressor.uncompress(data));
+        string decompressed = to!string(decompressor.uncompress(buffer));
         decompressor.flush();
+        buffer = null;
 
         return decompressed;
     }

diff --git a/source/droid/gateway/gateway.d b/source/droid/gateway/gateway.d
@@ -160,6 +160,9 @@ final class Gateway
             else
                 data = ws_.receiveText();
 
+            // The data isn't complete (not a full zlib message, or something borked)
+            if (data == "") return;
+
             const packet = parseMessage(data);
 
             auto opcodeHandler = packet.opcode in OPCODE_MAPPING;

diff --git a/source/droid/gateway/package.d b/source/droid/gateway/package.d
@@ -4,4 +4,5 @@ public {
     import droid.gateway.opcode;
     import droid.gateway.packet;
     import droid.gateway.gateway;
+    import droid.gateway.compression;
 }