Add support for WDPE 2.0.2-Snapshot

pom.xml

@@ -39,7 +39,7 @@
         <dependency>
             <groupId>dev.waterdog.waterdogpe</groupId>
             <artifactId>waterdog</artifactId>
-            <version>2.0.1-SNAPSHOT</version>
+            <version>2.0.2-SNAPSHOT</version>
             <scope>provided</scope>
         </dependency>
     </dependencies>

src/main/java/org/nethergames/proxytransport/compression/FrameIdCodec.java

@@ -0,0 +1,22 @@
+package org.nethergames.proxytransport.compression;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.MessageToMessageCodec;
+import org.cloudburstmc.protocol.bedrock.netty.BedrockBatchWrapper;
+
+import java.util.List;
+
+public class FrameIdCodec extends MessageToMessageCodec<ByteBuf, BedrockBatchWrapper> {
+    public static final String NAME = "frame-id-codec";
+
+    @Override
+    protected void encode(ChannelHandlerContext ctx, BedrockBatchWrapper msg, List<Object> out) throws Exception {
+        out.add(msg.getCompressed().retain());
+    }
+
+    @Override
+    protected void decode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> out) throws Exception {
+        out.add(BedrockBatchWrapper.newInstance(msg.readRetainedSlice(msg.readableBytes()), null));
+    }
+}

src/main/java/org/nethergames/proxytransport/compression/ProxyTransportAlgorithm.java

@@ -0,0 +1,11 @@
+package org.nethergames.proxytransport.compression;
+
+import org.cloudburstmc.protocol.bedrock.data.CompressionAlgorithm;
+
+public enum ProxyTransportAlgorithm implements CompressionAlgorithm {
+    ZSTD;
+
+    private ProxyTransportAlgorithm() {
+
+    }
+}

src/main/java/org/nethergames/proxytransport/compression/ProxyTransportCompressionCodec.java

@@ -0,0 +1,77 @@
+package org.nethergames.proxytransport.compression;
+
+import dev.waterdog.waterdogpe.network.connection.codec.compression.ProxiedCompressionCodec;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import org.cloudburstmc.protocol.bedrock.data.CompressionAlgorithm;
+import org.cloudburstmc.protocol.bedrock.data.PacketCompressionAlgorithm;
+import org.cloudburstmc.protocol.bedrock.netty.BedrockBatchWrapper;
+import org.cloudburstmc.protocol.bedrock.netty.codec.compression.BatchCompression;
+import org.cloudburstmc.protocol.bedrock.netty.codec.compression.CompressionStrategy;
+
+import java.util.List;
+
+public class ProxyTransportCompressionCodec extends ProxiedCompressionCodec {
+    private final boolean prefixed;
+    private final ZstdCompression zstdCompression = new ZstdCompression();
+
+    public ProxyTransportCompressionCodec(CompressionStrategy strategy, boolean prefixed) {
+        super(strategy, prefixed);
+        this.prefixed = prefixed;
+    }
+
+    protected void encode(ChannelHandlerContext ctx, BedrockBatchWrapper msg, List<Object> out) throws Exception {
+        if (msg.getCompressed() == null && msg.getUncompressed() == null) {
+            throw new IllegalStateException("Batch was not encoded before");
+        } else if (msg.getCompressed() != null && !msg.isModified()) { // already compressed
+            if (!this.prefixed) { // we need to prefix the compressed data
+                CompositeByteBuf buf = ctx.alloc().compositeDirectBuffer(2);
+                buf.addComponent(true, ctx.alloc().ioBuffer(1).writeByte(getCompressionHeader(msg.getAlgorithm())));
+                buf.addComponent(true, msg.getCompressed().retainedSlice());
+                msg.setCompressed(buf, msg.getAlgorithm());
+            }
+
+            this.onPassedThrough(ctx, msg);
+            out.add(msg.retain());
+        } else {
+            BatchCompression compression = this.getStrategy().getCompression(msg);
+            if (!compression.getAlgorithm().equals(PacketCompressionAlgorithm.NONE)) {
+                compression = this.zstdCompression;
+            }
+
+            ByteBuf compressed = compression.encode(ctx, msg.getUncompressed());
+
+            try {
+                ByteBuf outBuf;
+
+                outBuf = ctx.alloc().ioBuffer(1 + compressed.readableBytes());
+                outBuf.writeByte(this.getCompressionHeader(compression.getAlgorithm()));
+                outBuf.writeBytes(compressed);
+
+                msg.setCompressed(outBuf, compression.getAlgorithm());
+            } finally {
+                compressed.release();
+            }
+
+            this.onCompressed(ctx, msg);
+            out.add(msg.retain());
+        }
+    }
+
+    protected byte getCompressionHeader0(CompressionAlgorithm algorithm) {
+        if (algorithm instanceof ProxyTransportAlgorithm) {
+            return -2;
+        }
+
+        return super.getCompressionHeader0(algorithm);
+    }
+
+    protected CompressionAlgorithm getCompressionAlgorithm0(byte header) {
+        if (header == -2) {
+            return ProxyTransportAlgorithm.ZSTD;
+        }
+
+        return super.getCompressionAlgorithm0(header);
+    }
+}

src/main/java/org/nethergames/proxytransport/compression/ZstdCompression.java

@@ -0,0 +1,67 @@
+package org.nethergames.proxytransport.compression;
+
+import com.github.luben.zstd.Zstd;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.util.ReferenceCountUtil;
+import org.cloudburstmc.protocol.bedrock.data.CompressionAlgorithm;
+import org.cloudburstmc.protocol.bedrock.netty.codec.compression.BatchCompression;
+
+import java.nio.ByteBuffer;
+
+public class ZstdCompression implements BatchCompression {
+    private int level = -1;
+
+    public ByteBuf encode(ChannelHandlerContext ctx, ByteBuf msg) throws Exception {
+        ByteBuf direct;
+        if (!msg.isDirect() || msg instanceof CompositeByteBuf) {
+            direct = ctx.alloc().ioBuffer(msg.readableBytes());
+            direct.writeBytes(msg);
+        } else {
+            direct = msg;
+        }
+
+        ByteBuf output = ctx.alloc().directBuffer();
+        try {
+            int uncompressedLength = direct.readableBytes();
+            int maxLength = (int) Zstd.compressBound(uncompressedLength);
+
+            output.ensureWritable(maxLength);
+
+            int compressedLength;
+            if (direct.hasMemoryAddress()) {
+                compressedLength = (int) Zstd.compressUnsafe(output.memoryAddress(), maxLength, direct.memoryAddress() + direct.readerIndex(), uncompressedLength, this.level);
+            } else {
+                ByteBuffer sourceNio = direct.nioBuffer(direct.readerIndex(), direct.readableBytes());
+                ByteBuffer targetNio = output.nioBuffer(0, maxLength);
+
+                compressedLength = Zstd.compress(targetNio, sourceNio, this.level);
+            }
+
+            output.writerIndex(compressedLength);
+            return output.retain();
+        } finally {
+            ReferenceCountUtil.release(output);
+            if (direct != msg) {
+                ReferenceCountUtil.release(direct);
+            }
+        }
+    }
+
+    public ByteBuf decode(ChannelHandlerContext channelHandlerContext, ByteBuf byteBuf) throws Exception {
+        throw new UnsupportedOperationException("Zstd is not supported");
+    }
+
+    public CompressionAlgorithm getAlgorithm() {
+        return ProxyTransportAlgorithm.ZSTD;
+    }
+
+    public void setLevel(int level) {
+        this.level = level;
+    }
+
+    public int getLevel() {
+        return level;
+    }
+}

src/main/java/org/nethergames/proxytransport/impl/TransportChannelInitializer.java

@@ -5,6 +5,7 @@
 import dev.waterdog.waterdogpe.network.connection.client.ClientConnection;
 import dev.waterdog.waterdogpe.network.connection.codec.batch.BedrockBatchDecoder;
 import dev.waterdog.waterdogpe.network.connection.codec.batch.BedrockBatchEncoder;
+import dev.waterdog.waterdogpe.network.connection.codec.compression.CompressionType;
 import dev.waterdog.waterdogpe.network.connection.codec.packet.BedrockPacketCodec;
 import dev.waterdog.waterdogpe.network.serverinfo.ServerInfo;
 import dev.waterdog.waterdogpe.player.ProxiedPlayer;
@@ -16,13 +17,14 @@
 import org.cloudburstmc.netty.channel.raknet.RakChannel;
 import org.cloudburstmc.netty.channel.raknet.config.RakChannelOption;
 import org.cloudburstmc.netty.channel.raknet.config.RakMetrics;
-import org.nethergames.proxytransport.compression.ZstdCompressionCodec;
+import org.cloudburstmc.protocol.bedrock.netty.codec.compression.CompressionCodec;
+import org.nethergames.proxytransport.compression.FrameIdCodec;
+import org.nethergames.proxytransport.compression.ProxyTransportCompressionCodec;
 import org.nethergames.proxytransport.integration.CustomClientEventHandler;
 
 import static dev.waterdog.waterdogpe.network.connection.codec.initializer.ProxiedSessionInitializer.*;
 
 public class TransportChannelInitializer extends ChannelInitializer<Channel> {
-    private static final int ZSTD_COMPRESSION_LEVEL = 3;
 
     private final ProxiedPlayer player;
     private final ServerInfo serverInfo;
@@ -39,6 +41,7 @@ public TransportChannelInitializer(ProxiedPlayer player, ServerInfo serverInfo,
     @Override
     protected void initChannel(Channel channel) {
         int rakVersion = this.player.getProtocol().getRaknetVersion();
+        CompressionType compression = this.player.getProxy().getConfiguration().getCompression();
 
         channel.attr(PacketDirection.ATTRIBUTE).set(PacketDirection.FROM_SERVER);
 
@@ -55,14 +58,14 @@ protected void initChannel(Channel channel) {
                 .addLast(FRAME_DECODER, new LengthFieldBasedFrameDecoder(Integer.MAX_VALUE, 0, 4, 0, 4))
                 .addLast(FRAME_ENCODER, new LengthFieldPrepender(4));
 
-
-        ClientConnection connection = this.createConnection(channel);
         channel.pipeline()
-                .addLast(ZstdCompressionCodec.NAME, new ZstdCompressionCodec(ZSTD_COMPRESSION_LEVEL, connection))
+                .addLast(FrameIdCodec.NAME, new FrameIdCodec())
+                .addLast(CompressionCodec.NAME, new ProxyTransportCompressionCodec(getCompressionStrategy(compression, rakVersion, true), false))
                 .addLast(BedrockBatchDecoder.NAME, BATCH_DECODER)
                 .addLast(BedrockBatchEncoder.NAME, new BedrockBatchEncoder())
                 .addLast(BedrockPacketCodec.NAME, getPacketCodec(rakVersion));
 
+        ClientConnection connection = this.createConnection(channel);
         if (connection instanceof ChannelHandler handler) { // For reference: This will take care of the packets received being handled.
             channel.pipeline().addLast(ClientConnection.NAME, handler);
         }