Merge pull request #1155 from booky10/fix/bungee-injection-v2

Support bungee's new combined compressor

Author: booky10

bungeecord/src/main/java/io/github/retrooper/packetevents/handlers/PacketEventsDecoder.java

@@ -27,6 +27,7 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelPipeline;
 import io.netty.handler.codec.MessageToMessageDecoder;
 import net.md_5.bungee.api.connection.ProxiedPlayer;
 import org.jetbrains.annotations.NotNull;
@@ -35,14 +36,20 @@
 
 @ChannelHandler.Sharable
 public class PacketEventsDecoder extends MessageToMessageDecoder<ByteBuf> {
+
     public User user;
     public ProxiedPlayer player;
+    public boolean handledCompression;
 
     public PacketEventsDecoder(User user) {
         this.user = user;
     }
 
     public void read(ChannelHandlerContext ctx, ByteBuf byteBuf, List<Object> output) throws Exception {
+        if (this.tryFixCompressorOrder(ctx, byteBuf)) {
+            return; // skip handling of buffer
+        }
+
         ByteBuf transformed = ctx.alloc().buffer().writeBytes(byteBuf);
         try {
             int firstReaderIndex = transformed.readerIndex();
@@ -55,8 +62,7 @@ public void read(ChannelHandlerContext ctx, ByteBuf byteBuf, List<Object> output
                     ByteBufHelper.clear(packetReceiveEvent.getByteBuf());
                     packetReceiveEvent.getLastUsedWrapper().writeVarInt(packetReceiveEvent.getPacketId());
                     packetReceiveEvent.getLastUsedWrapper().write();
-                }
-                else {
+                } else {
                     transformed.readerIndex(firstReaderIndex);
                 }
                 output.add(transformed.retain());
@@ -88,4 +94,31 @@ public void channelInactive(@NotNull ChannelHandlerContext ctx) throws Exception
         ServerConnectionInitializer.destroyChannel(ctx.channel());
         super.channelInactive(ctx);
     }
+
+    private boolean tryFixCompressorOrder(ChannelHandlerContext ctx, ByteBuf buffer) {
+        if (this.handledCompression) {
+            return false;
+        }
+        ChannelPipeline pipe = ctx.pipeline();
+        List<String> pipeNames = pipe.names();
+        int decompressorIndex = pipeNames.indexOf("decompress");
+        if (decompressorIndex == -1) {
+            return false;
+        }
+        this.handledCompression = true;
+        if (!pipeNames.contains("frame-prepender-compress")) {
+            // before "modern" version, no need to handle this here
+            return false;
+        } else if (decompressorIndex <= pipeNames.indexOf(PacketEvents.DECODER_NAME)) {
+            return false; // order already seems to be correct
+        }
+        // relocate handler - encoder doesn't need relocation
+        PacketEventsDecoder decoder = (PacketEventsDecoder) pipe.remove(PacketEvents.DECODER_NAME);
+        pipe.addAfter("decompress", PacketEvents.DECODER_NAME, decoder);
+
+        // re-fire packet
+        ChannelHandlerContext frameDecoderCtx = pipe.context("frame-decoder");
+        frameDecoderCtx.fireChannelRead(buffer.retain());
+        return true; // skip further handling
+    }
 }

bungeecord/src/main/java/io/github/retrooper/packetevents/handlers/PacketEventsEncoder.java

@@ -30,27 +30,16 @@
 import io.netty.channel.ChannelOutboundHandlerAdapter;
 import io.netty.channel.ChannelPipeline;
 import io.netty.channel.ChannelPromise;
-import io.netty.handler.codec.EncoderException;
-import io.netty.util.Recycler;
 import io.netty.util.ReferenceCountUtil;
-import io.netty.util.concurrent.PromiseCombiner;
 import net.md_5.bungee.api.connection.ProxiedPlayer;
 
 import java.lang.reflect.InvocationTargetException;
-import java.util.ArrayList;
 import java.util.List;
 
 // Thanks to ViaVersion for the compression method.
 @ChannelHandler.Sharable
 public class PacketEventsEncoder extends ChannelOutboundHandlerAdapter {
 
-    private static final Recycler<OutList> OUT_LIST_RECYCLER = new Recycler<OutList>() {
-        @Override
-        protected OutList newObject(Handle<OutList> handle) {
-            return new OutList(handle);
-        }
-    };
-
     public ProxiedPlayer player;
     public User user;
     public boolean handledCompression;
@@ -59,8 +48,8 @@ public PacketEventsEncoder(User user) {
         this.user = user;
     }
 
-    public void read(ChannelHandlerContext ctx, ByteBuf buffer, ChannelPromise promise) throws Exception {
-        boolean doCompression = handleCompressionOrder(ctx, buffer);
+    public void read(ChannelHandlerContext originalCtx, ByteBuf buffer, ChannelPromise promise) {
+        ChannelHandlerContext ctx = this.tryFixCompressorOrder(originalCtx, buffer);
         int firstReaderIndex = buffer.readerIndex();
         PacketSendEvent packetSendEvent = EventCreationUtil.createSendEvent(ctx.channel(), user, player,
                 buffer, false);
@@ -74,11 +63,7 @@ public void read(ChannelHandlerContext ctx, ByteBuf buffer, ChannelPromise promi
             } else {
                 buffer.readerIndex(firstReaderIndex);
             }
-            if (doCompression) {
-                this.recompress(ctx, buffer, promise);
-            } else {
-                ctx.write(buffer, promise);
-            }
+            ctx.write(buffer, promise);
         } else {
             ReferenceCountUtil.release(packetSendEvent.getByteBuf());
         }
@@ -103,79 +88,52 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise)
         }
     }
 
-    private boolean handleCompressionOrder(ChannelHandlerContext ctx, ByteBuf buffer) {
+    private ChannelHandlerContext tryFixCompressorOrder(ChannelHandlerContext ctx, ByteBuf buffer) {
+        if (this.handledCompression) {
+            return ctx;
+        }
         ChannelPipeline pipe = ctx.pipeline();
-        if (handledCompression) {
-            return false;
+        List<String> pipeNames = pipe.names();
+        if (pipeNames.contains("frame-prepender-compress")) {
+            // "modern" version, no need to handle this here
+            this.handledCompression = true;
+            return ctx;
         }
-        int encoderIndex = pipe.names().indexOf("compress");
-        if (encoderIndex == -1) {
-            return false;
+        int compressorIndex = pipeNames.indexOf("compress");
+        if (compressorIndex == -1) {
+            return ctx;
         }
-        if (encoderIndex > pipe.names().indexOf(PacketEvents.ENCODER_NAME)) {
-            // Need to decompress this packet due to bad order
-            ChannelHandler decompressor = pipe.get("decompress");
-            try {
-                ByteBuf decompressed = (ByteBuf) CustomPipelineUtil.callPacketDecodeByteBuf(decompressor, ctx, buffer).get(0);
-                if (buffer != decompressed) {
-                    try {
-                        buffer.clear().writeBytes(decompressed);
-                    } finally {
-                        decompressed.release();
-                    }
-                }
-                //Relocate handlers
-                PacketEventsDecoder decoder = (PacketEventsDecoder) pipe.remove(PacketEvents.DECODER_NAME);
-                PacketEventsEncoder encoder = (PacketEventsEncoder) pipe.remove(PacketEvents.ENCODER_NAME);
-                pipe.addAfter("decompress", PacketEvents.DECODER_NAME, decoder);
-                pipe.addAfter("compress", PacketEvents.ENCODER_NAME, encoder);
-                //System.out.println("Pipe: " + ChannelHelper.pipelineHandlerNamesAsString(ctx.channel()));
-                handledCompression = true;
-                return true;
-            } catch (InvocationTargetException e) {
-                e.printStackTrace();
-            }
+        this.handledCompression = true;
+        if (compressorIndex <= pipeNames.indexOf(PacketEvents.ENCODER_NAME)) {
+            return ctx; // order already seems to be correct
         }
-        return false;
+        // relocate handlers
+        PacketEventsDecoder decoder = (PacketEventsDecoder) pipe.remove(PacketEvents.DECODER_NAME);
+        PacketEventsEncoder encoder = (PacketEventsEncoder) pipe.remove(PacketEvents.ENCODER_NAME);
+        pipe.addAfter("decompress", PacketEvents.DECODER_NAME, decoder);
+        pipe.addAfter("compress", PacketEvents.ENCODER_NAME, encoder);
+
+        // manually decompress packet and update context,
+        // so we don't need to additionally manually re-compress the packet
+        this.decompress(pipe, buffer);
+        return pipe.context(PacketEvents.ENCODER_NAME);
     }
 
-    private void recompress(ChannelHandlerContext ctx, ByteBuf buffer, ChannelPromise promise) {
-        OutList outWrapper = OUT_LIST_RECYCLER.get();
-        List<Object> out = outWrapper.list;
-        try {
-            ChannelHandler compressor = ctx.pipeline().get("compress");
-            CustomPipelineUtil.callPacketEncodeByteBuf(compressor, ctx, buffer, out);
+    private void decompress(ChannelPipeline pipe, ByteBuf buffer) {
+        ChannelHandler decompressor = pipe.get("decompress");
+        ChannelHandlerContext decompressorCtx = pipe.context("decompress");
 
-            int len = out.size();
-            if (len == 1) {
-                // should be the only case which
-                // happens on vanilla bungeecord
-                ctx.write(out.get(0), promise);
-            } else {
-                // copied from MessageToMessageEncoder#writePromiseCombiner
-                PromiseCombiner combiner = new PromiseCombiner(ctx.executor());
-                for (int i = 0; i < len; i++) {
-                    combiner.add(ctx.write(out.get(i)));
-                }
-                combiner.finish(promise);
+        ByteBuf decompressed = null;
+        try {
+            decompressed = (ByteBuf) CustomPipelineUtil.callPacketDecodeByteBuf(
+                    decompressor, decompressorCtx, buffer).get(0);
+            if (buffer != decompressed) {
+                buffer.clear().writeBytes(decompressed);
             }
         } catch (InvocationTargetException exception) {
-            throw new EncoderException("Error while recompressing bytebuf " + buffer.readableBytes(), exception);
+            throw new RuntimeException(exception);
         } finally {
-            out.clear();
-            outWrapper.handle.recycle(outWrapper);
-            buffer.release();
-        }
-    }
-
-    private static final class OutList {
-
-        // the default bungee compressor only produces one output bytebuf
-        private final List<Object> list = new ArrayList<>(1);
-        private final Recycler.Handle<OutList> handle;
-
-        public OutList(Recycler.Handle<OutList> handle) {
-            this.handle = handle;
+            ReferenceCountUtil.release(decompressed);
         }
     }
 }

bungeecord/src/main/java/io/github/retrooper/packetevents/injector/BungeePipelineInjector.java

@@ -19,7 +19,6 @@
 
 import com.github.retrooper.packetevents.PacketEvents;
 import com.github.retrooper.packetevents.injector.ChannelInjector;
-import com.github.retrooper.packetevents.protocol.ConnectionState;
 import com.github.retrooper.packetevents.protocol.player.User;
 import com.github.retrooper.packetevents.util.reflection.Reflection;
 import io.github.retrooper.packetevents.handlers.PacketEventsDecoder;
@@ -29,10 +28,8 @@
 import io.netty.channel.ChannelInitializer;
 import net.md_5.bungee.api.ProxyServer;
 import net.md_5.bungee.api.connection.ProxiedPlayer;
-import org.jetbrains.annotations.NotNull;
 
 import java.lang.reflect.Field;
-import java.lang.reflect.Method;
 import java.util.Set;
 
 //Thanks to ViaVersion for helping us design this injector.
@@ -81,10 +78,8 @@ public void injectChannel(Channel channel) {
         }
 
         try {
-            Field f = bootstrapAcceptor.getClass().getDeclaredField("childHandler");
-            f.setAccessible(true);
-            f.set(bootstrapAcceptor, newInitializer);
-        } catch (IllegalAccessException | NoSuchFieldException e) {
+            initializerField.set(bootstrapAcceptor, newInitializer);
+        } catch (IllegalAccessException e) {
             throw new RuntimeException(e);
         }
 

bungeecord/src/main/java/io/github/retrooper/packetevents/injector/CustomPipelineUtil.java

@@ -149,7 +149,7 @@ public static void callPacketEncodeByteBuf(Object encoder, Object ctx, Object ms
     }
 
     public static List<Object> callPacketDecodeByteBuf(Object decoder, Object ctx, Object msg) throws InvocationTargetException {
-        List<Object> output = new ArrayList<>();
+        List<Object> output = new ArrayList<>(1);
         if (BUNGEE_PACKET_DECODE_BYTEBUF == null) {
             try {
                 BUNGEE_PACKET_DECODE_BYTEBUF = decoder.getClass().getDeclaredMethod("decode", ChannelHandlerContext.class,