From b6d5875f5f72dcb8eed009b1715988cd26b0d86c Mon Sep 17 00:00:00 2001
From: UnfamiliarLegacy <74633542+UnfamiliarLegacy@users.noreply.github.com>
Date: Fri, 28 Jun 2024 04:52:20 +0200
Subject: [PATCH] Deal properly with partially received packets

---
 .../gearth/protocol/memory/Rc4Obtainer.java   |  2 +
 .../shockwave/buffers/ShockwaveOutBuffer.java | 15 ++++---
 G-Earth/src/test/java/TestRc4Shockwave.java   | 41 +++++++++++++++++++
 3 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/G-Earth/src/main/java/gearth/protocol/memory/Rc4Obtainer.java b/G-Earth/src/main/java/gearth/protocol/memory/Rc4Obtainer.java
index 728f578..97f87b5 100644
--- a/G-Earth/src/main/java/gearth/protocol/memory/Rc4Obtainer.java
+++ b/G-Earth/src/main/java/gearth/protocol/memory/Rc4Obtainer.java
@@ -121,6 +121,8 @@ public class Rc4Obtainer {
     }
 
     private boolean onSendFirstEncryptedMessage(EncryptedPacketHandler packetHandler, List<byte[]> potentialRC4tables) {
+        logger.info("Attempting to brute force RC4 table");
+
         if (potentialRC4tables == null || potentialRC4tables.isEmpty()) {
             return false;
         }
diff --git a/G-Earth/src/main/java/gearth/protocol/packethandler/shockwave/buffers/ShockwaveOutBuffer.java b/G-Earth/src/main/java/gearth/protocol/packethandler/shockwave/buffers/ShockwaveOutBuffer.java
index 18df6df..70582ef 100644
--- a/G-Earth/src/main/java/gearth/protocol/packethandler/shockwave/buffers/ShockwaveOutBuffer.java
+++ b/G-Earth/src/main/java/gearth/protocol/packethandler/shockwave/buffers/ShockwaveOutBuffer.java
@@ -17,6 +17,7 @@ public class ShockwaveOutBuffer extends PayloadBuffer {
     public static final int PACKET_SIZE_MIN = PACKET_HEADER_SIZE + PACKET_LENGTH_SIZE;
     public static final int PACKET_SIZE_MIN_ENCRYPTED = PACKET_HEADER_SIZE + PACKET_LENGTH_SIZE_ENCRYPTED;
 
+    private int previousLength = 0;
     private RC4Cipher cipher;
 
     @Override
@@ -39,12 +40,15 @@ public class ShockwaveOutBuffer extends PayloadBuffer {
             int length;
 
             if (this.cipher != null) {
-                final byte[] decData = this.cipher.decipher(buffer, 0, PACKET_LENGTH_SIZE_ENCRYPTED);
-                final int decDataLen = Base64Encoding.decode(new byte[]{decData[1], decData[2], decData[3]});
+                if (previousLength  == 0) {
+                    final byte[] decData = this.cipher.decipher(buffer, 0, PACKET_LENGTH_SIZE_ENCRYPTED);
 
-                // TODO: Store length in a variable for if we don't have enough bytes.
-
-                length = decDataLen;
+                    // When a packet has been received that we can't fully read, we need to store the decrypted length.
+                    // Otherwise, we would keep decrypting the same bytes and mutating the rc4 state, messing up the entire state.
+                    length = previousLength = Base64Encoding.decode(new byte[]{decData[1], decData[2], decData[3]});
+                } else {
+                    length = previousLength;
+                }
             } else {
                 length = Base64Encoding.decode(new byte[]{buffer[0], buffer[1], buffer[2]});
             }
@@ -58,6 +62,7 @@ public class ShockwaveOutBuffer extends PayloadBuffer {
             out.add(Arrays.copyOfRange(buffer, packetLengthSize, endPos));
 
             buffer = Arrays.copyOfRange(buffer, endPos, buffer.length);
+            previousLength = 0;
         }
 
         return out.toArray(new byte[0][]);
diff --git a/G-Earth/src/test/java/TestRc4Shockwave.java b/G-Earth/src/test/java/TestRc4Shockwave.java
index df68846..cea24f3 100644
--- a/G-Earth/src/test/java/TestRc4Shockwave.java
+++ b/G-Earth/src/test/java/TestRc4Shockwave.java
@@ -7,6 +7,7 @@ import gearth.protocol.memory.habboclient.HabboClientFactory;
 import gearth.protocol.memory.habboclient.external.MemoryClient;
 import gearth.protocol.packethandler.EncryptedPacketHandler;
 import gearth.protocol.packethandler.shockwave.ShockwavePacketOutgoingHandler;
+import gearth.protocol.packethandler.shockwave.buffers.ShockwaveOutBuffer;
 import org.bouncycastle.util.encoders.Hex;
 import org.junit.jupiter.api.Test;
 import org.mockito.MockedStatic;
@@ -210,4 +211,44 @@ public class TestRc4Shockwave {
         assertEquals(h3Q, cipher.getQ());
         assertEquals(h3J, cipher.getJ());
     }
+
+    @Test
+    public void testSplitPackets() {
+        final RC4Cipher c = new RC4Base64(
+                Hex.decode("d102ecab2e8d0a851000a393a483de68f2182f879b884bb6be77595701ffc0900db9f00f415332fd9fe35dcc8ceaa5c480214cc8ee661627e23736795b444a3d5a3f721cfbd4b370d7daaae7f8b1769a78569c3065d6c7607438752c39ced51992634e15a850efcd1a5fd0b782cb6254f4c598f36f485e35f7bf86a9fc03344d8ae9fe07db1355bc7d7f2de8e42052b424e0954f8f71df4612b0c95c7c252aacc394edbb7e45b540f6cadcf18443678ebd513169580e913a3b6b0664290c11fac173a26e1f09ada66199b2b805d8d9d38117331b0447262b1d8b08ddcf7b42e63e97289d23896d1ea70b147af5e19ed2aeebaf6c3cbac2a16a96e522c6f949a0"),
+                152,
+                79
+        );
+
+        final ShockwaveOutBuffer buffer = new ShockwaveOutBuffer();
+
+        buffer.setCipher(c);
+
+        // State.
+        int received = 0;
+
+        // Packet 1.
+        buffer.push(Hex.decode("4b78316b61774f426838576553476f48355478684946554c7368525749373856326a784b723762"));
+        received += buffer.receive().length;
+
+        buffer.push(Hex.decode("7866797279636c313743324a34714a392f5361584e656a6a6b6e53657550476b7638684637302b354d42484d373350434c434a3977"));
+        received += buffer.receive().length;
+
+        // Packet 2.
+        buffer.push(Hex.decode("3634616845413754796c6371324d346c"));
+        received += buffer.receive().length;
+
+        buffer.push(Hex.decode("546f616b62484d55346a5943694c4b386c4a"));
+        received += buffer.receive().length;
+
+        // Packet 3.
+        buffer.push(Hex.decode("4750"));
+        received += buffer.receive().length;
+
+        buffer.push(Hex.decode("41364d67417330"));
+        received += buffer.receive().length;
+
+        assertEquals(3, received);
+        assertTrue(buffer.isEmpty());
+    }
 }