From bcc652ede8715ca4c2a22c4328267e6622caeddd Mon Sep 17 00:00:00 2001
From: themode <themode@outlook.fr>
Date: Thu, 27 Jan 2022 15:29:16 +0100
Subject: [PATCH] Add tests for incomplete reads

---
 .../server/network/SocketReadTest.java        | 60 ++++++++++++++++++-
 1 file changed, 58 insertions(+), 2 deletions(-)

diff --git a/src/test/java/net/minestom/server/network/SocketReadTest.java b/src/test/java/net/minestom/server/network/SocketReadTest.java
index d69752571..99b64c26f 100644
--- a/src/test/java/net/minestom/server/network/SocketReadTest.java
+++ b/src/test/java/net/minestom/server/network/SocketReadTest.java
@@ -2,6 +2,7 @@ package net.minestom.server.network;
 
 import net.minestom.server.network.packet.client.play.ClientPluginMessagePacket;
 import net.minestom.server.utils.PacketUtils;
+import net.minestom.server.utils.Utils;
 import net.minestom.server.utils.binary.BinaryBuffer;
 import net.minestom.server.utils.binary.BinaryReader;
 import net.minestom.server.utils.binary.PooledBuffers;
@@ -10,8 +11,7 @@ import org.junit.jupiter.params.provider.ValueSource;
 
 import java.util.zip.DataFormatException;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.*;
 
 public class SocketReadTest {
 
@@ -62,4 +62,60 @@ public class SocketReadTest {
             assertEquals(2000, readPacket.data().length);
         }
     }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    public void insufficientLength(boolean compressed) throws DataFormatException {
+        // Write a complete packet then the next packet length without any payload
+
+        var packet = new ClientPluginMessagePacket("channel", new byte[2000]);
+
+        var buffer = PooledBuffers.packetBuffer();
+        PacketUtils.writeFramedPacket(buffer, 0x0A, packet, compressed ? 256 : 0);
+        Utils.writeVarInt(buffer, 200); // incomplete 200 bytes packet
+
+        var wrapper = BinaryBuffer.wrap(buffer);
+        wrapper.reset(0, buffer.position());
+
+        var result = PacketUtils.readPackets(wrapper, compressed);
+        var remaining = result.remaining();
+        assertNotNull(remaining);
+        assertEquals(Utils.getVarIntSize(200), remaining.readableBytes());
+
+        var packets = result.packets();
+        assertEquals(1, packets.size());
+        var rawPacket = packets.get(0);
+        assertEquals(0x0A, rawPacket.id());
+        var readPacket = new ClientPluginMessagePacket(new BinaryReader(rawPacket.payload()));
+        assertEquals("channel", readPacket.channel());
+        assertEquals(2000, readPacket.data().length);
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    public void incomplete(boolean compressed) throws DataFormatException {
+        // Write a complete packet and incomplete var-int length for the next packet
+
+        var packet = new ClientPluginMessagePacket("channel", new byte[2000]);
+
+        var buffer = PooledBuffers.packetBuffer();
+        PacketUtils.writeFramedPacket(buffer, 0x0A, packet, compressed ? 256 : 0);
+        buffer.put((byte) -85); // incomplete var-int length
+
+        var wrapper = BinaryBuffer.wrap(buffer);
+        wrapper.reset(0, buffer.position());
+
+        var result = PacketUtils.readPackets(wrapper, compressed);
+        var remaining = result.remaining();
+        assertNotNull(remaining);
+        assertEquals(1, remaining.readableBytes());
+
+        var packets = result.packets();
+        assertEquals(1, packets.size());
+        var rawPacket = packets.get(0);
+        assertEquals(0x0A, rawPacket.id());
+        var readPacket = new ClientPluginMessagePacket(new BinaryReader(rawPacket.payload()));
+        assertEquals("channel", readPacket.channel());
+        assertEquals(2000, readPacket.data().length);
+    }
 }