Fix reference count issues (#165)

* Fix reference count issues

* Only remove the handler after we get the Handshake packet

* Same thing for the ServerDataHandlers. Optimized imports as well

* Further removal optimizations in backend handlers and Spigot

* Removed unused imports

Co-authored-by: Tim203 <mctim203@gmail.com>
Co-authored-by: Camotoy <20743703+Camotoy@users.noreply.github.com>

bungee/src/main/java/org/geysermc/floodgate/addon/data/BungeeDataAddon.java

@@ -35,7 +35,6 @@ import org.geysermc.floodgate.api.logger.FloodgateLogger;
 import org.geysermc.floodgate.api.player.FloodgatePlayer;
 import org.geysermc.floodgate.config.ProxyFloodgateConfig;
 import org.geysermc.floodgate.player.FloodgateHandshakeHandler;
-import org.geysermc.floodgate.util.Utils;
 
 public class BungeeDataAddon implements InjectorAddon {
     @Inject private FloodgateHandshakeHandler handshakeHandler;
@@ -62,10 +61,12 @@ public class BungeeDataAddon implements InjectorAddon {
     @Override
     public void onInject(Channel channel, boolean toServer) {
         if (toServer) {
-            channel.pipeline().addAfter(
-                    packetEncoder, "floodgate_data_handler",
-                    new BungeeServerDataHandler(config, api, playerAttribute)
-            );
+            if (config.isSendFloodgateData()) {
+                channel.pipeline().addAfter(
+                        packetEncoder, "floodgate_data_handler",
+                        new BungeeServerDataHandler(api, playerAttribute)
+                );
+            }
             return;
         }
         channel.pipeline().addBefore(
@@ -76,7 +77,6 @@ public class BungeeDataAddon implements InjectorAddon {
 
     @Override
     public void onLoginDone(Channel channel) {
-        onRemoveInject(channel);
     }
 
     @Override
@@ -89,7 +89,6 @@ public class BungeeDataAddon implements InjectorAddon {
 
     @Override
     public void onRemoveInject(Channel channel) {
-        Utils.removeHandler(channel.pipeline(), "floodgate_data_handler");
     }
 
     @Override

bungee/src/main/java/org/geysermc/floodgate/addon/data/BungeeProxyDataHandler.java

@@ -30,7 +30,6 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.util.AttributeKey;
-import io.netty.util.ReferenceCountUtil;
 import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
 import lombok.RequiredArgsConstructor;
@@ -65,22 +64,17 @@ public class BungeeProxyDataHandler extends ChannelInboundHandlerAdapter {
     private final ProxyFloodgateConfig config;
     private final FloodgateHandshakeHandler handler;
     private final AttributeKey<String> kickMessageAttribute;
-    private boolean done;
 
     @Override
     public void channelRead(ChannelHandlerContext ctx, Object msg) {
-        ReferenceCountUtil.retain(msg);
-        if (done || !(msg instanceof PacketWrapper)) {
-            ctx.fireChannelRead(msg);
-            return;
-        }
+        if (msg instanceof PacketWrapper) {
+            DefinedPacket packet = ((PacketWrapper) msg).packet;
 
-        DefinedPacket packet = ((PacketWrapper) msg).packet;
-
-        // we're only interested in the Handshake packet
-        if (packet instanceof Handshake) {
-            handleHandshake(ctx, (Handshake) packet);
-            done = true;
+            // we're only interested in the Handshake packet
+            if (packet instanceof Handshake) {
+                handleHandshake(ctx, (Handshake) packet);
+                ctx.pipeline().remove(this);
+            }
         }
 
         ctx.fireChannelRead(msg);

bungee/src/main/java/org/geysermc/floodgate/addon/data/BungeeServerDataHandler.java

@@ -28,11 +28,10 @@ package org.geysermc.floodgate.addon.data;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import io.netty.channel.ChannelHandlerContext;
-import io.netty.handler.codec.MessageToMessageEncoder;
+import io.netty.channel.ChannelOutboundHandlerAdapter;
+import io.netty.channel.ChannelPromise;
 import io.netty.util.AttributeKey;
-import io.netty.util.ReferenceCountUtil;
 import java.lang.reflect.Field;
-import java.util.List;
 import lombok.RequiredArgsConstructor;
 import net.md_5.bungee.ServerConnector;
 import net.md_5.bungee.UserConnection;
@@ -41,14 +40,13 @@ import net.md_5.bungee.netty.HandlerBoss;
 import net.md_5.bungee.protocol.packet.Handshake;
 import org.geysermc.floodgate.api.ProxyFloodgateApi;
 import org.geysermc.floodgate.api.player.FloodgatePlayer;
-import org.geysermc.floodgate.config.ProxyFloodgateConfig;
 import org.geysermc.floodgate.player.FloodgatePlayerImpl;
 import org.geysermc.floodgate.util.BedrockData;
 import org.geysermc.floodgate.util.ReflectionUtils;
 
 @SuppressWarnings("ConstantConditions")
 @RequiredArgsConstructor
-public class BungeeServerDataHandler extends MessageToMessageEncoder<Object> {
+public class BungeeServerDataHandler extends ChannelOutboundHandlerAdapter {
     private static final Field HANDLER;
     private static final Field USER_CONNECTION;
     private static final Field CHANNEL_WRAPPER;
@@ -65,51 +63,40 @@ public class BungeeServerDataHandler extends MessageToMessageEncoder<Object> {
         checkNotNull(CHANNEL_WRAPPER, "ChannelWrapper field cannot be null");
     }
 
-    private final ProxyFloodgateConfig config;
     private final ProxyFloodgateApi api;
     private final AttributeKey<FloodgatePlayer> playerAttribute;
-    private boolean done;
 
     @Override
-    protected void encode(ChannelHandlerContext ctx, Object packet, List<Object> out) {
-        ReferenceCountUtil.retain(packet);
-        if (done) {
-            out.add(packet);
-            return;
+    public void write(ChannelHandlerContext ctx, Object packet, ChannelPromise promise)
+            throws Exception {
+        if (packet instanceof Handshake) {
+            // get the Proxy <-> Player channel from the Proxy <-> Server channel
+            HandlerBoss handlerBoss = ctx.pipeline().get(HandlerBoss.class);
+            ServerConnector connector = ReflectionUtils.getCastedValue(handlerBoss, HANDLER);
+            UserConnection connection = ReflectionUtils.getCastedValue(connector, USER_CONNECTION);
+            ChannelWrapper wrapper = ReflectionUtils.getCastedValue(connection, CHANNEL_WRAPPER);
+
+            FloodgatePlayer player = wrapper.getHandle().attr(playerAttribute).get();
+
+            if (player != null) {
+                BedrockData data = player.as(FloodgatePlayerImpl.class).toBedrockData();
+                String encryptedData = api.createEncryptedDataString(data);
+
+                Handshake handshake = (Handshake) packet;
+                String address = handshake.getHost();
+
+                // our data goes before all the other data
+                int addressFinished = address.indexOf('\0');
+                String originalAddress = address.substring(0, addressFinished);
+                String remaining = address.substring(addressFinished);
+
+                handshake.setHost(originalAddress + '\0' + encryptedData + remaining);
+                // Bungeecord will add his data after our data
+            }
+
+            ctx.pipeline().remove(this);
         }
 
-        // passes the information through to the connecting server if enabled
-        if (!(packet instanceof Handshake) || !config.isSendFloodgateData()) {
-            done = true;
-            out.add(packet);
-            return;
-        }
-
-        // get the Proxy <-> Player channel from the Proxy <-> Server channel
-        HandlerBoss handlerBoss = ctx.pipeline().get(HandlerBoss.class);
-        ServerConnector connector = ReflectionUtils.getCastedValue(handlerBoss, HANDLER);
-        UserConnection connection = ReflectionUtils.getCastedValue(connector, USER_CONNECTION);
-        ChannelWrapper wrapper = ReflectionUtils.getCastedValue(connection, CHANNEL_WRAPPER);
-
-        FloodgatePlayer player = wrapper.getHandle().attr(playerAttribute).get();
-
-        if (player != null) {
-            BedrockData data = player.as(FloodgatePlayerImpl.class).toBedrockData();
-            String encryptedData = api.createEncryptedDataString(data);
-
-            Handshake handshake = (Handshake) packet;
-            String address = handshake.getHost();
-
-            // our data goes before all the other data
-            int addressFinished = address.indexOf('\0');
-            String originalAddress = address.substring(0, addressFinished);
-            String remaining = address.substring(addressFinished);
-
-            handshake.setHost(originalAddress + '\0' + encryptedData + remaining);
-            // Bungeecord will add his data after our data
-        }
-
-        done = true;
-        out.add(packet);
+        ctx.write(packet, promise);
     }
 }

spigot/src/main/java/org/geysermc/floodgate/addon/data/SpigotDataAddon.java

@@ -35,7 +35,6 @@ 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.Utils;
 
 public final class SpigotDataAddon implements InjectorAddon {
     @Inject private FloodgateHandshakeHandler handshakeHandler;
@@ -61,7 +60,6 @@ public final class SpigotDataAddon implements InjectorAddon {
 
     @Override
     public void onLoginDone(Channel channel) {
-        onRemoveInject(channel);
     }
 
     @Override
@@ -74,7 +72,6 @@ public final class SpigotDataAddon implements InjectorAddon {
 
     @Override
     public void onRemoveInject(Channel channel) {
-        Utils.removeHandler(channel.pipeline(), "floodgate_data_handler");
     }
 
     @Override

spigot/src/main/java/org/geysermc/floodgate/addon/data/SpigotDataHandler.java

@@ -31,7 +31,6 @@ import static org.geysermc.floodgate.util.ReflectionUtils.setValue;
 import com.mojang.authlib.GameProfile;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
-import io.netty.util.ReferenceCountUtil;
 import java.net.InetSocketAddress;
 import lombok.RequiredArgsConstructor;
 import org.geysermc.floodgate.api.handshake.HandshakeData;
@@ -52,21 +51,14 @@ public final class SpigotDataHandler extends ChannelInboundHandlerAdapter {
     private final FloodgateLogger logger;
     private Object networkManager;
     private FloodgatePlayer player;
-    private boolean bungeeData;
-    private boolean done;
 
     @Override
     public void channelRead(ChannelHandlerContext ctx, Object packet) throws Exception {
-        ReferenceCountUtil.retain(packet);
-        // we're done but we're not yet removed from the connection
-        if (done) {
-            ctx.fireChannelRead(packet);
-            return;
-        }
-
         boolean isHandshake = ClassNames.HANDSHAKE_PACKET.isInstance(packet);
         boolean isLogin = ClassNames.LOGIN_START_PACKET.isInstance(packet);
 
+        boolean bungeeData = false;
+
         try {
             if (isHandshake) {
                 networkManager = ctx.channel().pipeline().get("packet_handler");
@@ -121,48 +113,43 @@ public final class SpigotDataHandler extends ChannelInboundHandlerAdapter {
                     setValue(networkManager, "spoofedUUID", player.getCorrectUniqueId());
                 }
             } else if (isLogin) {
-                if (!bungeeData) {
-                    // we have to fake the offline player (login) cycle
-                    Object loginListener = ClassNames.PACKET_LISTENER.get(networkManager);
+                // we have to fake the offline player (login) cycle
+                Object loginListener = ClassNames.PACKET_LISTENER.get(networkManager);
 
-                    // check if the server is actually in the Login state
-                    if (!ClassNames.LOGIN_LISTENER.isInstance(loginListener)) {
-                        // player is not in the login state, abort
-                        return;
-                    }
-
-                    // set the player his GameProfile, we can't change the username without this
-                    GameProfile gameProfile = new GameProfile(
-                            player.getCorrectUniqueId(), player.getCorrectUsername()
-                    );
-                    setValue(loginListener, ClassNames.LOGIN_PROFILE, gameProfile);
-
-                    // just like on Spigot:
-
-                    // LoginListener#initUUID
-                    // new LoginHandler().fireEvents();
-
-                    // and the tick of LoginListener will do the rest
-
-                    ClassNames.INIT_UUID.invoke(loginListener);
-
-                    Object loginHandler =
-                            ClassNames.LOGIN_HANDLER_CONSTRUCTOR.newInstance(loginListener);
-                    ClassNames.FIRE_LOGIN_EVENTS.invoke(loginHandler);
+                // check if the server is actually in the Login state
+                if (!ClassNames.LOGIN_LISTENER.isInstance(loginListener)) {
+                    // player is not in the login state, abort
+                    return;
                 }
+
+                // set the player his GameProfile, we can't change the username without this
+                GameProfile gameProfile = new GameProfile(
+                        player.getCorrectUniqueId(), player.getCorrectUsername()
+                );
+                setValue(loginListener, ClassNames.LOGIN_PROFILE, gameProfile);
+
+                // just like on Spigot:
+
+                // LoginListener#initUUID
+                // new LoginHandler().fireEvents();
+
+                // and the tick of LoginListener will do the rest
+
+                ClassNames.INIT_UUID.invoke(loginListener);
+
+                Object loginHandler =
+                        ClassNames.LOGIN_HANDLER_CONSTRUCTOR.newInstance(loginListener);
+                ClassNames.FIRE_LOGIN_EVENTS.invoke(loginHandler);
             }
         } finally {
             // don't let the packet through if the packet is the login packet
-            // because we want to skip the login cycle
-            if (isLogin) {
-                ReferenceCountUtil.release(packet, 2);
-            } else {
+            if (!isLogin) {
                 ctx.fireChannelRead(packet);
             }
 
-            if (isHandshake && bungeeData || isLogin && !bungeeData || player == null) {
-                // we're done, we'll just wait for the loginSuccessCall
-                done = true;
+            if (isHandshake && bungeeData || isLogin || player == null) {
+                // We're done
+                ctx.pipeline().remove(this);
             }
         }
     }

velocity/src/main/java/org/geysermc/floodgate/addon/data/VelocityDataAddon.java

@@ -36,7 +36,6 @@ import org.geysermc.floodgate.api.logger.FloodgateLogger;
 import org.geysermc.floodgate.api.player.FloodgatePlayer;
 import org.geysermc.floodgate.config.ProxyFloodgateConfig;
 import org.geysermc.floodgate.player.FloodgateHandshakeHandler;
-import org.geysermc.floodgate.util.Utils;
 
 public final class VelocityDataAddon implements InjectorAddon {
     @Inject private FloodgateHandshakeHandler handshakeHandler;
@@ -64,10 +63,12 @@ public final class VelocityDataAddon implements InjectorAddon {
     @Override
     public void onInject(Channel channel, boolean toServer) {
         if (toServer) {
-            channel.pipeline().addAfter(
-                    packetEncoder, "floodgate_data_handler",
-                    new VelocityServerDataHandler(config, api, proxy)
-            );
+            if (config.isSendFloodgateData()) {
+                channel.pipeline().addAfter(
+                        packetEncoder, "floodgate_data_handler",
+                        new VelocityServerDataHandler(api, proxy)
+                );
+            }
             return;
         }
         // The handler is already added so we should add our handler before it
@@ -79,7 +80,6 @@ public final class VelocityDataAddon implements InjectorAddon {
 
     @Override
     public void onLoginDone(Channel channel) {
-        onRemoveInject(channel);
     }
 
     @Override
@@ -92,7 +92,6 @@ public final class VelocityDataAddon implements InjectorAddon {
 
     @Override
     public void onRemoveInject(Channel channel) {
-        Utils.removeHandler(channel.pipeline(), "floodgate_data_handler");
     }
 
     @Override

velocity/src/main/java/org/geysermc/floodgate/addon/data/VelocityProxyDataHandler.java

@@ -34,7 +34,6 @@ import static org.geysermc.floodgate.util.ReflectionUtils.setValue;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.util.AttributeKey;
-import io.netty.util.ReferenceCountUtil;
 import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
 import lombok.RequiredArgsConstructor;
@@ -74,21 +73,17 @@ public final class VelocityProxyDataHandler extends ChannelInboundHandlerAdapter
     private final FloodgateHandshakeHandler handshakeHandler;
     private final AttributeKey<String> kickMessageAttribute;
     private final FloodgateLogger logger;
-    private boolean done;
 
     @Override
     public void channelRead(ChannelHandlerContext ctx, Object msg) {
-        ReferenceCountUtil.retain(msg);
         // we're only interested in the Handshake packet.
         // it should be the first packet but you never know
-        if (done || !HANDSHAKE_PACKET.isInstance(msg)) {
-            ctx.fireChannelRead(msg);
-            return;
+        if (HANDSHAKE_PACKET.isInstance(msg)) {
+            handleClientToProxy(ctx, msg);
+            ctx.pipeline().remove(this);
         }
 
-        handleClientToProxy(ctx, msg);
         ctx.fireChannelRead(msg);
-        done = true;
     }
 
     private void handleClientToProxy(ChannelHandlerContext ctx, Object packet) {

velocity/src/main/java/org/geysermc/floodgate/addon/data/VelocityServerDataHandler.java

@@ -37,19 +37,17 @@ import static org.geysermc.floodgate.util.ReflectionUtils.setValue;
 import com.velocitypowered.api.proxy.Player;
 import com.velocitypowered.api.proxy.ProxyServer;
 import io.netty.channel.ChannelHandlerContext;
-import io.netty.handler.codec.MessageToMessageEncoder;
-import io.netty.util.ReferenceCountUtil;
+import io.netty.channel.ChannelOutboundHandlerAdapter;
+import io.netty.channel.ChannelPromise;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
-import java.util.List;
 import org.geysermc.floodgate.api.ProxyFloodgateApi;
 import org.geysermc.floodgate.api.player.FloodgatePlayer;
-import org.geysermc.floodgate.config.ProxyFloodgateConfig;
 import org.geysermc.floodgate.player.FloodgatePlayerImpl;
 import org.geysermc.floodgate.util.BedrockData;
 
 @SuppressWarnings("ConstantConditions")
-public final class VelocityServerDataHandler extends MessageToMessageEncoder<Object> {
+public final class VelocityServerDataHandler extends ChannelOutboundHandlerAdapter {
     private static final Class<?> HANDSHAKE_PACKET;
     private static final Field HANDSHAKE_ADDRESS;
     private static final Method GET_ASSOCIATION;
@@ -79,16 +77,12 @@ public final class VelocityServerDataHandler extends MessageToMessageEncoder<Obj
         checkNotNull(GET_PLAYER, "getPlayer in VelocityServerConnection cannot be null");
     }
 
-    private final ProxyFloodgateConfig config;
     private final ProxyFloodgateApi api;
     private final boolean isModernForwarding;
     //private final AttributeKey<FloodgatePlayer> playerAttribute;
-    private boolean done;
 
-    public VelocityServerDataHandler(ProxyFloodgateConfig config,
-                                     ProxyFloodgateApi api,
+    public VelocityServerDataHandler(ProxyFloodgateApi api,
                                      ProxyServer proxy) {
-        this.config = config;
         this.api = api;
 
         Enum<?> forwardingMode = castedInvoke(proxy.getConfiguration(), GET_FORWARDING_MODE);
@@ -96,62 +90,51 @@ public final class VelocityServerDataHandler extends MessageToMessageEncoder<Obj
     }
 
     @Override
-    protected void encode(ChannelHandlerContext ctx, Object packet, List<Object> out) {
-        ReferenceCountUtil.retain(packet);
-        if (done) {
-            out.add(packet);
-            return;
+    public void write(ChannelHandlerContext ctx, Object packet, ChannelPromise promise)
+            throws Exception {
+        if (HANDSHAKE_PACKET.isInstance(packet)) {
+            String address = getCastedValue(packet, HANDSHAKE_ADDRESS);
+
+            // get the FloodgatePlayer from the ConnectedPlayer
+            Object minecraftConnection = ctx.pipeline().get("handler");
+            Object association = invoke(minecraftConnection, GET_ASSOCIATION);
+            Player velocityPlayer = castedInvoke(association, GET_PLAYER);
+
+            //noinspection ConstantConditions
+            FloodgatePlayer player = api.getPlayer(velocityPlayer.getUniqueId());
+
+            //todo use something similar to what's written below for a more direct approach
+
+            // get the Proxy <-> Player channel from the Proxy <-> Server channel
+            //MinecraftConnection minecraftConnection = ctx.pipeline().get("handler");
+            //((VelocityServerConnection) minecraftConnection.association).proxyPlayer.connection.channel
+
+            //FloodgatePlayer player = playerChannel.attr(playerAttribute).get();
+            if (player != null) {
+                // Player is a Floodgate player
+                BedrockData data = player.as(FloodgatePlayerImpl.class).toBedrockData();
+                String encryptedData = api.createEncryptedDataString(data);
+
+                // use the same system that we use on bungee, our data goes before all the other data
+                int addressFinished = address.indexOf('\0');
+                String originalAddress;
+                String remaining;
+                if (isModernForwarding && addressFinished == -1) {
+                    // There is no additional data to hook onto
+                    originalAddress = address;
+                    remaining = "";
+                } else {
+                    originalAddress = address.substring(0, addressFinished);
+                    remaining = address.substring(addressFinished);
+                }
+
+                setValue(packet, HANDSHAKE_ADDRESS, originalAddress + '\0' + encryptedData
+                        + remaining);
+            }
+
+            ctx.pipeline().remove(this);
         }
 
-        if (!HANDSHAKE_PACKET.isInstance(packet) || !config.isSendFloodgateData()) {
-            done = true;
-            out.add(packet);
-            return;
-        }
-
-        String address = getCastedValue(packet, HANDSHAKE_ADDRESS);
-
-        // get the FloodgatePlayer from the ConnectedPlayer
-        Object minecraftConnection = ctx.pipeline().get("handler");
-        Object association = invoke(minecraftConnection, GET_ASSOCIATION);
-        Player velocityPlayer = castedInvoke(association, GET_PLAYER);
-
-        //noinspection ConstantConditions
-        FloodgatePlayer player = api.getPlayer(velocityPlayer.getUniqueId());
-
-        //todo use something similar to what's written below for a more direct approach
-
-        // get the Proxy <-> Player channel from the Proxy <-> Server channel
-        //MinecraftConnection minecraftConnection = ctx.pipeline().get("handler");
-        //((VelocityServerConnection) minecraftConnection.association).proxyPlayer.connection.channel
-
-        //FloodgatePlayer player = playerChannel.attr(playerAttribute).get();
-
-        // player is not a Floodgate player
-        if (player == null) {
-            out.add(packet);
-            return;
-        }
-
-        BedrockData data = player.as(FloodgatePlayerImpl.class).toBedrockData();
-        String encryptedData = api.createEncryptedDataString(data);
-
-        // use the same system that we use on bungee, our data goes before all the other data
-        int addressFinished = address.indexOf('\0');
-        String originalAddress;
-        String remaining;
-        if (isModernForwarding && addressFinished == -1) {
-            // There is no additional data to hook onto
-            originalAddress = address;
-            remaining = "";
-        } else {
-            originalAddress = address.substring(0, addressFinished);
-            remaining = address.substring(addressFinished);
-        }
-
-        setValue(packet, HANDSHAKE_ADDRESS, originalAddress + '\0' + encryptedData + remaining);
-
-        done = true;
-        out.add(packet);
+        ctx.write(packet, promise);
     }
 }