Implement play state compression and handle compression being set twice (#607)

* Implement play state compression and handle compression being set twice

* Fix 0 threshold; Code cleanup

---------

Co-authored-by: FlorianMichael <[email protected]>

Author: Beaness

common/src/main/java/com/viaversion/viarewind/protocol/v1_8to1_7_6_10/Protocol1_8To1_7_6_10.java

@@ -81,18 +81,8 @@ public void register() {
                 map(Types.BYTE_ARRAY_PRIMITIVE, Types.SHORT_BYTE_ARRAY); // verification token
             }
         });
-        this.registerClientbound(State.LOGIN, ClientboundLoginPackets.LOGIN_COMPRESSION, new PacketHandlers() {
-            @Override
-            public void register() {
-                handler(wrapper -> {
-                    final int threshold = wrapper.read(Types.VAR_INT);
-
-                    Via.getManager().getProviders().get(CompressionHandlerProvider.class).onHandleLoginCompressionPacket(wrapper.user(), threshold);
-                    wrapper.cancel();
-                });
-            }
-        });
-        this.cancelClientbound(ClientboundPackets1_8.SET_COMPRESSION); // unused
+        this.registerClientbound(State.LOGIN, ClientboundLoginPackets.LOGIN_COMPRESSION, this::handleCompression);
+        this.registerClientbound(ClientboundPackets1_8.SET_COMPRESSION, null, this::handleCompression);
         this.registerClientbound(ClientboundPackets1_8.KEEP_ALIVE, new PacketHandlers() {
             @Override
             public void register() {
@@ -116,6 +106,13 @@ public void register() {
         });
     }
 
+    private void handleCompression(final PacketWrapper wrapper) {
+        wrapper.cancel();
+        final int threshold = wrapper.read(Types.VAR_INT);
+
+        Via.getManager().getProviders().get(CompressionHandlerProvider.class).setCompressionThreshold(wrapper.user(), threshold);
+    }
+
     @Override
     public void transform(Direction direction, State state, PacketWrapper packetWrapper) throws CancelException {
         Via.getManager().getProviders().get(CompressionHandlerProvider.class).onTransformPacket(packetWrapper.user());

common/src/main/java/com/viaversion/viarewind/protocol/v1_8to1_7_6_10/provider/CompressionHandlerProvider.java

@@ -24,19 +24,19 @@
 
 public abstract class CompressionHandlerProvider implements Provider {
 
-    public abstract void onHandleLoginCompressionPacket(UserConnection user, int threshold);
+    public abstract void setCompressionThreshold(UserConnection user, int threshold);
 
     public abstract void onTransformPacket(UserConnection user);
 
     public abstract ChannelHandler getEncoder(int threshold);
 
     public abstract ChannelHandler getDecoder(int threshold);
 
-    public boolean isCompressionEnabled(UserConnection user) {
+    public boolean isRemoveCompression(UserConnection user) {
         return user.get(CompressionStatusTracker.class).removeCompression;
     }
 
-    public void setCompressionEnabled(UserConnection user, boolean enabled) {
+    public void setRemoveCompression(UserConnection user, boolean enabled) {
         user.get(CompressionStatusTracker.class).removeCompression = enabled;
     }
 }

common/src/main/java/com/viaversion/viarewind/protocol/v1_8to1_7_6_10/provider/compression/CompressionDecoder.java

@@ -29,12 +29,16 @@
 public class CompressionDecoder extends MessageToMessageDecoder<ByteBuf> {
     private final Inflater inflater = new Inflater();
 
-    private final int threshold;
+    private int threshold;
 
     public CompressionDecoder(final int threshold) {
         this.threshold = threshold;
     }
 
+    public void setThreshold(final int threshold) {
+        this.threshold = threshold;
+    }
+
     @Override
     protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) throws Exception {
         if (!in.isReadable()) return;

common/src/main/java/com/viaversion/viarewind/protocol/v1_8to1_7_6_10/provider/compression/CompressionEncoder.java

@@ -27,12 +27,16 @@
 public class CompressionEncoder extends MessageToByteEncoder<ByteBuf> {
     private final Deflater deflater = new Deflater();
 
-    private final int threshold;
+    private int threshold;
 
     public CompressionEncoder(final int threshold) {
         this.threshold = threshold;
     }
 
+    public void setThreshold(final int threshold) {
+        this.threshold = threshold;
+    }
+
     @Override
     protected void encode(ChannelHandlerContext ctx, ByteBuf in, ByteBuf out) throws Exception {
         int frameLength = in.readableBytes();

common/src/main/java/com/viaversion/viarewind/protocol/v1_8to1_7_6_10/provider/compression/TrackingCompressionHandlerProvider.java

@@ -29,19 +29,33 @@
 public class TrackingCompressionHandlerProvider extends CompressionHandlerProvider {
 
     @Override
-    public void onHandleLoginCompressionPacket(UserConnection user, int threshold) {
+    public void setCompressionThreshold(UserConnection user, int threshold) {
         final ChannelPipeline pipeline = user.getChannel().pipeline();
-        if (user.isClientSide()) {
-            pipeline.addBefore(Via.getManager().getInjector().getEncoderName(), ViaRewind.getPlatform().compressHandlerName(), getEncoder(threshold));
-            pipeline.addBefore(Via.getManager().getInjector().getDecoderName(), ViaRewind.getPlatform().decompressHandlerName(), getDecoder(threshold));
+        if (!user.isClientSide() || threshold < 0) {
+            setRemoveCompression(user, true); // We need to remove compression for 1.7 clients
+            return;
+        }
+
+        final String compressHandlerName = ViaRewind.getPlatform().compressHandlerName();
+        final CompressionEncoder encoder = (CompressionEncoder) pipeline.get(compressHandlerName);
+        if (encoder != null) {
+            encoder.setThreshold(threshold);
+        } else {
+            pipeline.addBefore(Via.getManager().getInjector().getEncoderName(), compressHandlerName, getEncoder(threshold));
+        }
+
+        final String decompressHandlerName = ViaRewind.getPlatform().decompressHandlerName();
+        final CompressionDecoder decoder = (CompressionDecoder) pipeline.get(decompressHandlerName);
+        if (decoder != null) {
+            decoder.setThreshold(threshold);
         } else {
-            setCompressionEnabled(user, true); // We need to remove compression for 1.7 clients
+            pipeline.addBefore(Via.getManager().getInjector().getDecoderName(), decompressHandlerName, getDecoder(threshold));
         }
     }
 
     @Override
     public void onTransformPacket(UserConnection user) {
-        if (isCompressionEnabled(user)) {
+        if (isRemoveCompression(user)) {
             final ChannelPipeline pipeline = user.getChannel().pipeline();
 
             String compressor = null;
@@ -58,7 +72,7 @@ public void onTransformPacket(UserConnection user) {
                 throw new IllegalStateException("Couldn't remove compression for 1.7!");
             }
 
-            setCompressionEnabled(user, false);
+            setRemoveCompression(user, false);
         }
     }