Update to latest Floodgate changes

Co-Authored-By: byquanton <32410361+byquanton@users.noreply.github.com>

src/main/java/org/geysermc/floodgate/addon/data/FabricDataAddon.java

@@ -22,18 +22,19 @@ public final class FabricDataAddon implements InjectorAddon {
     @Named("packetHandler")
     private String packetHandlerName;
 
+    @Inject
+    @Named("kickMessageAttribute")
+    private AttributeKey<String> kickMessageAttribute;
+
     @Inject
     @Named("playerAttribute")
     private AttributeKey<FloodgatePlayer> playerAttribute;
 
     @Override
     public void onInject(Channel channel, boolean toServer) {
-        PacketBlocker blocker = new PacketBlocker();
-        channel.pipeline().addBefore("decoder", "floodgate_packet_blocker", blocker);
-
         channel.pipeline().addBefore(
                 packetHandlerName, "floodgate_data_handler",
-                new FabricDataHandler(config, handshakeHandler, blocker, logger)
+                new FabricDataHandler(handshakeHandler, config, kickMessageAttribute, logger)
         );
     }
 

src/main/java/org/geysermc/floodgate/addon/data/FabricDataHandler.java

@@ -1,133 +1,91 @@
 package org.geysermc.floodgate.addon.data;
 
-import com.google.common.collect.Queues;
+import io.netty.channel.Channel;
+import io.netty.util.AttributeKey;
+import net.minecraft.text.Text;
+import org.geysermc.floodgate.api.logger.FloodgateLogger;
 import org.geysermc.floodgate.mixin.ClientConnectionMixin;
+import org.geysermc.floodgate.mixin.ClientIntentionPacketMixin;
 import org.geysermc.floodgate.mixin_interface.ServerLoginNetworkHandlerSetter;
 import com.mojang.authlib.GameProfile;
 import io.netty.channel.ChannelHandlerContext;
-import io.netty.channel.ChannelInboundHandlerAdapter;
-import io.netty.util.ReferenceCountUtil;
-import lombok.RequiredArgsConstructor;
 import net.minecraft.network.ClientConnection;
 import net.minecraft.network.packet.c2s.handshake.HandshakeC2SPacket;
 import net.minecraft.network.packet.c2s.login.LoginHelloC2SPacket;
 import net.minecraft.server.network.ServerLoginNetworkHandler;
-import org.geysermc.floodgate.api.handshake.HandshakeData;
-import org.geysermc.floodgate.api.logger.FloodgateLogger;
 import org.geysermc.floodgate.api.player.FloodgatePlayer;
 import org.geysermc.floodgate.config.FloodgateConfig;
 import org.geysermc.floodgate.player.FloodgateHandshakeHandler;
-import org.geysermc.floodgate.util.BedrockData;
-import org.geysermc.floodgate.util.Constants;
+import org.geysermc.floodgate.player.FloodgateHandshakeHandler.HandshakeResult;
 
 import java.net.InetSocketAddress;
-import java.util.Queue;
 
-@RequiredArgsConstructor
-public final class FabricDataHandler extends ChannelInboundHandlerAdapter {
-    private final FloodgateConfig config;
-    private final FloodgateHandshakeHandler handshakeHandler;
-    private final PacketBlocker blocker;
+public final class FabricDataHandler extends CommonDataHandler {
     private final FloodgateLogger logger;
-
-    private final Queue<Object> packetQueue = Queues.newConcurrentLinkedQueue();
-
     private ClientConnection networkManager;
     private FloodgatePlayer player;
 
-    /**
-     * As a variable so we can change it without reflection
-     */
-    HandshakeC2SPacket handshakePacket;
-    /**
-     * A boolean to compensate for the above
-     */
-    private boolean packetsBlocked;
-
-    @Override
-    public void channelRead(ChannelHandlerContext ctx, Object packet) {
-        // prevent other packets from being handled while we handle the handshake packet
-        if (packetsBlocked) {
-            packetQueue.add(packet);
-            return;
-        }
-
-        if (packet instanceof HandshakeC2SPacket handshakePacket) {
-            blocker.enable();
-            packetsBlocked = true;
-            this.handshakePacket = handshakePacket;
-
-            networkManager = (ClientConnection) ctx.channel().pipeline().get("packet_handler");
-
-            handshakeHandler.handle(ctx.channel(), handshakePacket.getAddress()).thenApply(result -> {
-                HandshakeData handshakeData = result.getHandshakeData();
-
-                this.handshakePacket = new HandshakeC2SPacket(handshakeData.getHostname(),
-                        handshakePacket.getPort(), handshakePacket.getIntendedState());
-
-                InetSocketAddress newIp = result.getNewIp(ctx.channel());
-                if (newIp != null) {
-                    ((ClientConnectionMixin) networkManager).setAddress(newIp);
-                }
-
-                if (handshakeData.getDisconnectReason() != null) {
-                    ctx.close(); //todo disconnect with message
-                    return true;
-                }
-
-                //todo use kickMessageAttribute and let this be common logic
-
-                switch (result.getResultType()) {
-                    case SUCCESS:
-                        break;
-                    case EXCEPTION:
-                        logger.info(Constants.INTERNAL_ERROR_MESSAGE);
-                        ctx.close();
-                        return true;
-                    case DECRYPT_ERROR:
-                        logger.info(config.getDisconnect().getInvalidKey());
-                        ctx.close();
-                        return true;
-                    case INVALID_DATA_LENGTH:
-                        int dataLength = result.getBedrockData().getDataLength();
-                        logger.info(
-                                config.getDisconnect().getInvalidArgumentsLength(),
-                                BedrockData.EXPECTED_LENGTH, dataLength
-                        );
-                        ctx.close();
-                        return true;
-                    default: // only continue when SUCCESS
-                        return true;
-                }
-
-                player = result.getFloodgatePlayer();
-                return player == null;
-            }).thenAccept(shouldRemove -> {
-                ctx.fireChannelRead(this.handshakePacket);
-                Object queuedPacket;
-                while ((queuedPacket = packetQueue.poll()) != null) {
-                    if (checkLogin(ctx, packet)) {
-                        break;
-                    }
-                    ctx.fireChannelRead(queuedPacket);
-                }
-
-                if (shouldRemove) {
-                    ctx.pipeline().remove(FabricDataHandler.this);
-                }
-                blocker.disable();
-                packetsBlocked = false;
-            });
-            return;
-        }
-
-        if (!checkLogin(ctx, packet)) {
-            ctx.fireChannelRead(packet);
-        }
+    public FabricDataHandler(
+            FloodgateHandshakeHandler handshakeHandler,
+            FloodgateConfig config,
+            AttributeKey<String> kickMessageAttribute, FloodgateLogger logger) {
+        super(handshakeHandler, config, kickMessageAttribute, new PacketBlocker());
+        this.logger = logger;
     }
 
-    private boolean checkLogin(ChannelHandlerContext ctx, Object packet) {
+    @Override
+    protected void setNewIp(Channel channel, InetSocketAddress newIp) {
+        ((ClientConnectionMixin) this.networkManager).setAddress(newIp);
+    }
+
+    @Override
+    protected Object setHostname(Object handshakePacket, String hostname) {
+        // While it would be ideal to simply create a new handshake packet, the packet constructor
+        // does not allow us to set the protocol version
+        ((ClientIntentionPacketMixin) handshakePacket).setAddress(hostname);
+        return handshakePacket;
+    }
+
+    @Override
+    protected boolean shouldRemoveHandler(HandshakeResult result) {
+        player = result.getFloodgatePlayer();
+
+        if (getKickMessage() != null) {
+            // we also have to keep this handler if we want to kick then with a disconnect message
+            return false;
+        } else if (player == null) {
+            // player is not a Floodgate player
+            return true;
+        }
+
+        if (result.getResultType() == FloodgateHandshakeHandler.ResultType.SUCCESS) {
+            logger.info("Floodgate player who is logged in as {} {} joined",
+                    player.getCorrectUsername(), player.getCorrectUniqueId());
+        }
+
+        // Handler will be removed after the login hello packet is handled
+        return false;
+    }
+
+    @Override
+    protected boolean channelRead(Object packet) {
+        if (packet instanceof HandshakeC2SPacket handshakePacket) {
+            ctx.pipeline().addAfter("splitter", "floodgate_packet_blocker", blocker);
+            networkManager = (ClientConnection) ctx.channel().pipeline().get("packet_handler");
+            handle(packet, handshakePacket.getAddress());
+            return false;
+        }
+        return !checkAndHandleLogin(packet);
+    }
+
+    private boolean checkAndHandleLogin(Object packet) {
         if (packet instanceof LoginHelloC2SPacket) {
+            String kickMessage = getKickMessage();
+            if (kickMessage != null) {
+                networkManager.disconnect(Text.of(kickMessage));
+                return true;
+            }
+
             // we have to fake the offline player (login) cycle
             if (!(networkManager.getPacketListener() instanceof ServerLoginNetworkHandler)) {
                 // player is not in the login state, abort

src/main/java/org/geysermc/floodgate/mixin/ClientIntentionPacketMixin.java

@@ -0,0 +1,14 @@
+package org.geysermc.floodgate.mixin;
+
+import net.minecraft.network.packet.c2s.handshake.HandshakeC2SPacket;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Mutable;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+@Mixin(HandshakeC2SPacket.class)
+public interface ClientIntentionPacketMixin {
+
+    @Accessor("address")
+    @Mutable
+    void setAddress(String address);
+}

src/main/java/org/geysermc/floodgate/mixin_interface/ServerLoginNetworkHandlerSetter.java

@@ -1,9 +1,6 @@
 package org.geysermc.floodgate.mixin_interface;
 
 import com.mojang.authlib.GameProfile;
-import net.minecraft.server.network.ServerLoginNetworkHandler;
-
-import java.util.UUID;
 
 public interface ServerLoginNetworkHandlerSetter {
     void setGameProfile(GameProfile profile);

src/main/resources/floodgate.mixins.json

@@ -6,7 +6,8 @@
   "mixins": [
     "ClientConnectionMixin",
     "ServerLoginNetworkHandlerMixin",
-    "ServerNetworkIoMixin"
+    "ServerNetworkIoMixin",
+    "ClientIntentionPacketMixin"
   ],
   "injectors": {
     "defaultRequire": 1