Fragments: All parts should have NEW_PROTOCOL flag if original Packet has it

Fragments: fix last packet detection

src/main/java/com/github/manevolent/ts3j/protocol/packet/fragment/Fragments.java

@@ -66,7 +66,8 @@ public final class Fragments {
                 int flush = Math.min(maxFragmentSize, size - offs);
 
                 boolean first = offs == 0;
-                boolean last = flush < maxFragmentSize;
+                boolean hasNext = offs + flush < size;
+                boolean last = !hasNext;
 
                 Packet piece = new Packet(packet.getRole());
                 piece.setHeader(packet.getHeader().clone());
@@ -74,6 +75,12 @@ public final class Fragments {
                 if (!first) // Only first packet has flags
                     piece.getHeader().setPacketFlags(HeaderFlag.NONE.getIndex());
 
+                // All parts should have NEW_PROTOCOL flag if original Packet has it
+                piece.getHeader().setPacketFlag(
+                        HeaderFlag.NEW_PROTOCOL,
+                        packet.getHeader().getPacketFlag(HeaderFlag.NEW_PROTOCOL)
+                );
+
                 // First and last packet get FRAGMENTED flag
                 piece.getHeader().setPacketFlag(HeaderFlag.FRAGMENTED, first || last);
 

src/test/java/com/github/manevolent/ts3j/FragmentTest.java

@@ -3,23 +3,113 @@ package com.github.manevolent.ts3j;
 import com.github.manevolent.ts3j.protocol.Packet;
 import com.github.manevolent.ts3j.protocol.ProtocolRole;
 import com.github.manevolent.ts3j.protocol.header.ClientPacketHeader;
+import com.github.manevolent.ts3j.protocol.header.HeaderFlag;
 import com.github.manevolent.ts3j.protocol.packet.PacketBody2Command;
 import com.github.manevolent.ts3j.protocol.packet.PacketBodyType;
 import com.github.manevolent.ts3j.protocol.packet.fragment.Fragments;
 import com.github.manevolent.ts3j.protocol.packet.fragment.PacketReassembly;
 import com.github.manevolent.ts3j.protocol.socket.client.AbstractTeamspeakClientSocket;
 import com.github.manevolent.ts3j.util.Pair;
+import com.github.manevolent.ts3j.util.QuickLZ;
 import com.github.manevolent.ts3j.util.Ts3Debugging;
 import junit.framework.TestCase;
 import org.bouncycastle.util.encoders.Base64;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Random;
 
 public class FragmentTest extends TestCase {
 
     public static void main(String[] args) throws Exception {
-        new FragmentTest().testParser();
+        FragmentTest fragmentTest = new FragmentTest();
+
+        fragmentTest.testParser();
+        fragmentTest.testFragmentedFlag();
+        fragmentTest.testNewProtocolFlag();
+    }
+
+    private List<Packet> getSplitPackets(boolean isNewProtocol) {
+        ClientPacketHeader header = new ClientPacketHeader();
+        header.setType(PacketBodyType.COMMAND);
+        if (isNewProtocol) {
+            header.setPacketFlag(HeaderFlag.NEW_PROTOCOL, true);
+        }
+
+        Packet largePacket = new Packet(ProtocolRole.CLIENT, header);
+
+        largePacket.setBody(new PacketBody2Command(
+                ProtocolRole.CLIENT,
+                String.join("", Collections.nCopies(1000, "a!")) + "!"
+        ));
+
+        return Fragments.split(largePacket);
+    }
+
+    /**
+     * All parts should have NEW_PROTOCOL flag if original Packet has it
+     */
+    public void testNewProtocolFlag() {
+        for (Packet piece : getSplitPackets(false)) {
+            assertFalse(
+                    "Unexpected NEW_PROTOCOL flag",
+                    piece.getHeader().getPacketFlag(HeaderFlag.NEW_PROTOCOL)
+            );
+        }
+
+        for (Packet piece : getSplitPackets(true)) {
+            assertTrue(
+                    "NEW_PROTOCOL flag expected",
+                    piece.getHeader().getPacketFlag(HeaderFlag.NEW_PROTOCOL)
+            );
+        }
+    }
+
+    /**
+     * First and last fragmented packets should have FRAGMENTED flag.
+     */
+    public void testFragmentedFlag() {
+        Ts3Debugging.setEnabled(true);
+        ClientPacketHeader header = new ClientPacketHeader();
+
+        header.setType(PacketBodyType.COMMAND);
+        header.setPacketFlag(HeaderFlag.NEW_PROTOCOL, true);
+        header.setPacketFlag(HeaderFlag.COMPRESSED, true);
+
+        Packet largePacket = new Packet(ProtocolRole.CLIENT, header);
+
+        /*
+         * We need to construct a packet which should be split
+         * to exactly two packets of MAXIMUM_PACKET_SIZE size after compression.
+         */
+
+        int targetCompressedBodyLength = 2 * (Fragments.MAXIMUM_PACKET_SIZE - largePacket.getHeader().getSize());
+
+
+        //this may be fragile if MAXIMUM_PACKET_SIZE or split implementation will change
+        String payload = String.join("", Collections.nCopies(482, "a!")) + "!";
+
+        assertEquals(
+                "Unexpected compressed body length",
+                targetCompressedBodyLength,
+                QuickLZ.compress(payload.getBytes(), 1).length
+        );
+
+        largePacket.setBody(new PacketBody2Command(ProtocolRole.CLIENT, payload));
+
+        List<Packet> pieces = Fragments.split(largePacket);
+
+        assertEquals(2, pieces.size());
+
+        assertTrue(
+                "First packet should have FRAGMENTED flag set.",
+                pieces.get(0).getHeader().getPacketFlag(HeaderFlag.FRAGMENTED)
+        );
+
+        assertTrue(
+                "Last packet should have FRAGMENTED flag set.",
+                pieces.get(1).getHeader().getPacketFlag(HeaderFlag.FRAGMENTED)
+        );
     }
 
     public void testParser() throws Exception {