From cced369d7e3f8daa4d1c89a6b95080165aa7b9ab Mon Sep 17 00:00:00 2001
From: Tim203 <mctim203@gmail.com>
Date: Fri, 10 Sep 2021 23:12:58 +0200
Subject: [PATCH] Fixed Proxy <-> Server injection

---
 .../inject/bungee/BungeeInjector.java         | 90 ++++++++++++-------
 1 file changed, 60 insertions(+), 30 deletions(-)

diff --git a/bungee/src/main/java/org/geysermc/floodgate/inject/bungee/BungeeInjector.java b/bungee/src/main/java/org/geysermc/floodgate/inject/bungee/BungeeInjector.java
index aa077c60..473bda82 100644
--- a/bungee/src/main/java/org/geysermc/floodgate/inject/bungee/BungeeInjector.java
+++ b/bungee/src/main/java/org/geysermc/floodgate/inject/bungee/BungeeInjector.java
@@ -26,11 +26,10 @@
 package org.geysermc.floodgate.inject.bungee;
 
 import io.netty.channel.Channel;
-import io.netty.channel.ChannelDuplexHandler;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
-import io.netty.channel.ChannelInitializer;
-import io.netty.channel.ChannelPipeline;
+import io.netty.channel.ChannelOutboundHandlerAdapter;
+import io.netty.channel.ChannelPromise;
 import java.lang.reflect.Field;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
@@ -44,6 +43,8 @@ import org.geysermc.floodgate.util.ReflectionUtils;
 
 @RequiredArgsConstructor
 public final class BungeeInjector extends CommonPlatformInjector {
+    private static final String BUNGEE_INIT = "floodgate-bungee-init";
+
     private final FloodgateLogger logger;
     @Getter private boolean injected;
 
@@ -58,7 +59,7 @@ public final class BungeeInjector extends CommonPlatformInjector {
             // (Instead of just replacing the ChannelInitializer which is only called for
             // player <-> proxy)
             BungeeCustomPrepender customPrepender = new BungeeCustomPrepender(
-                    ReflectionUtils.getCastedValue(null, framePrepender), logger
+                    this, ReflectionUtils.getCastedValue(null, framePrepender)
             );
 
             BungeeReflectionUtils.setFieldValue(null, framePrepender, customPrepender);
@@ -82,47 +83,76 @@ public final class BungeeInjector extends CommonPlatformInjector {
         return false;
     }
 
-    public void injectClient(Channel channel, boolean clientToProxy) {
+    void injectClient(Channel channel, boolean clientToProxy) {
+        if (!channel.isOpen()) {
+            return;
+        }
+
+        if (channel.pipeline().get(MinecraftEncoder.class) == null) {
+            logger.debug(
+                    "Minecraft encoder not found while injecting! {}",
+                    String.join(", ", channel.pipeline().names())
+            );
+            return;
+        }
+
         injectAddonsCall(channel, !clientToProxy);
         addInjectedClient(channel);
     }
 
-    private static final String BUNGEE_INIT = "floodgate-bungee-init";
-
     @RequiredArgsConstructor
-    private final class BungeeCustomPrepender extends Varint21LengthFieldPrepender {
+    private static final class BungeeCustomPrepender extends Varint21LengthFieldPrepender {
+        private final BungeeInjector injector;
         private final Varint21LengthFieldPrepender original;
-        private final FloodgateLogger logger;
 
         @Override
         public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
             original.handlerAdded(ctx);
             // The Minecraft encoder being in the pipeline isn't present until later
-            ctx.pipeline().addBefore(PipelineUtils.FRAME_DECODER, BUNGEE_INIT,
-                    new ChannelInboundHandlerAdapter() {
-                        @Override
-                        public void channelRead(ChannelHandlerContext ctx, Object msg)
-                                throws Exception {
-                            if (ctx.channel().pipeline().get(MinecraftEncoder.class) == null) {
-                                logger.debug("Minecraft encoder class not found while " +
-                                        "injecting!");
-                            } else {
-                                ctx.channel().pipeline().addFirst(new BungeeInjectorInitializer());
-                            }
-                            ctx.channel().pipeline().remove(BUNGEE_INIT);
-                            super.channelRead(ctx, msg);
-                        }
-                    });
+
+            if (ctx.channel().parent() != null) {
+                // Client <-> Proxy
+                ctx.pipeline().addBefore(
+                        PipelineUtils.FRAME_DECODER, BUNGEE_INIT,
+                        new BungeeClientToProxyInjectInitializer(injector)
+                );
+            } else {
+                // Proxy <-> Server
+                ctx.pipeline().addLast(
+                        BUNGEE_INIT, new BungeeProxyToServerInjectInitializer(injector)
+                );
+            }
         }
     }
 
-    private final class BungeeInjectorInitializer extends ChannelInitializer<Channel> {
+    @RequiredArgsConstructor
+    private static final class BungeeClientToProxyInjectInitializer
+            extends ChannelInboundHandlerAdapter {
+
+        private final BungeeInjector injector;
+
         @Override
-        protected void initChannel(Channel channel) {
-            if (!channel.isOpen()) {
-                return;
-            }
-            injectClient(channel, channel.parent() != null);
+        public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
+            injector.injectClient(ctx.channel(), true);
+
+            ctx.pipeline().remove(this);
+            super.channelRead(ctx, msg);
+        }
+    }
+
+    @RequiredArgsConstructor
+    private static final class BungeeProxyToServerInjectInitializer
+            extends ChannelOutboundHandlerAdapter {
+
+        private final BungeeInjector injector;
+
+        @Override
+        public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise)
+                throws Exception {
+            injector.injectClient(ctx.channel(), false);
+
+            ctx.pipeline().remove(this);
+            super.write(ctx, msg, promise);
         }
     }
 }