More fixes

NetherGamesMC · Feb 10, 2024 · 43bdd97 · 43bdd97
1 parent 19eeae2
commit 43bdd97
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 21 deletions.
diff --git a/src/main/java/org/nethergames/proxytransport/compression/ProxyTransportCompressionCodec.java b/src/main/java/org/nethergames/proxytransport/compression/ProxyTransportCompressionCodec.java
@@ -5,7 +5,9 @@
 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;
@@ -23,28 +25,31 @@ protected void encode(ChannelHandlerContext ctx, BedrockBatchWrapper msg, List<O
         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) {
-                this.onPassedThrough(ctx, msg);
-                out.add(msg.retain());
-            } else { // we need to prefix the compressed data
+            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());
-
-                this.onPassedThrough(ctx, msg);
-                out.add(buf.retain());
+                msg.setCompressed(buf, msg.getAlgorithm());
             }
+
+            this.onPassedThrough(ctx, msg);
+            out.add(msg.retain());
         } else {
-            ByteBuf compressed = this.zstdCompression.encode(ctx, msg.getUncompressed());
+            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(this.zstdCompression.getAlgorithm()));
+                outBuf.writeByte(this.getCompressionHeader(compression.getAlgorithm()));
                 outBuf.writeBytes(compressed);
 
-                msg.setCompressed(outBuf, this.zstdCompression.getAlgorithm());
+                msg.setCompressed(outBuf, compression.getAlgorithm());
             } finally {
                 compressed.release();
             }
@@ -59,14 +64,14 @@ protected byte getCompressionHeader0(CompressionAlgorithm algorithm) {
             return -2;
         }
 
-        return -1;
+        return super.getCompressionHeader0(algorithm);
     }
 
     protected CompressionAlgorithm getCompressionAlgorithm0(byte header) {
         if (header == -2) {
             return ProxyTransportAlgorithm.ZSTD;
         }
 
-        return null;
+        return super.getCompressionAlgorithm0(header);
     }
 }
diff --git a/src/main/java/org/nethergames/proxytransport/impl/TransportChannelInitializer.java b/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.client.ClientPacketQueue;
 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;
@@ -56,14 +57,13 @@ 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(ProxyTransportCompressionCodec.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);
         }

diff --git a/src/main/java/org/nethergames/proxytransport/utils/CompressionType.java b/src/main/java/org/nethergames/proxytransport/utils/CompressionType.java