From ac6b0e560962bbc23a40f14f4b4098e53c81b544 Mon Sep 17 00:00:00 2001
From: Dries C <coenendries3@gmail.com>
Date: Sat, 10 Feb 2024 18:59:19 +0100
Subject: [PATCH] More fixes

---
 .../ProxyTransportCompressionCodec.java          | 16 +++++++---------
 .../impl/TransportChannelInitializer.java        |  4 ++--
 .../proxytransport/utils/CompressionType.java    |  7 -------
 3 files changed, 9 insertions(+), 18 deletions(-)
 delete mode 100644 src/main/java/org/nethergames/proxytransport/utils/CompressionType.java

diff --git a/src/main/java/org/nethergames/proxytransport/compression/ProxyTransportCompressionCodec.java b/src/main/java/org/nethergames/proxytransport/compression/ProxyTransportCompressionCodec.java
index 075569c..67e6d74 100644
--- a/src/main/java/org/nethergames/proxytransport/compression/ProxyTransportCompressionCodec.java
+++ b/src/main/java/org/nethergames/proxytransport/compression/ProxyTransportCompressionCodec.java
@@ -23,17 +23,15 @@ 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());
 
@@ -59,7 +57,7 @@ protected byte getCompressionHeader0(CompressionAlgorithm algorithm) {
             return -2;
         }
 
-        return -1;
+        return super.getCompressionHeader0(algorithm);
     }
 
     protected CompressionAlgorithm getCompressionAlgorithm0(byte header) {
@@ -67,6 +65,6 @@ protected CompressionAlgorithm getCompressionAlgorithm0(byte header) {
             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
index c565844..090bd40 100644
--- 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
deleted file mode 100644
index 62fb838..0000000
--- a/src/main/java/org/nethergames/proxytransport/utils/CompressionType.java
+++ /dev/null
@@ -1,7 +0,0 @@
-package org.nethergames.proxytransport.utils;
-
-public enum CompressionType {
-    METHOD_ZLIB,
-    METHOD_ZSTD,
-    METHOD_SNAPPY;
-}
\ No newline at end of file