Fix userconnection availability on join event (PoC) (#2849)

bukkit/src/main/java/com/viaversion/viaversion/bukkit/handlers/BukkitEncodeHandler.java

@@ -100,4 +100,8 @@ public class BukkitEncodeHandler extends MessageToByteEncoder implements ViaCode
             cause.printStackTrace(); // Print if CB doesn't already do it
         }
     }
+
+    public UserConnection getInfo() {
+        return info;
+    }
 }

bukkit/src/main/java/com/viaversion/viaversion/bukkit/listeners/JoinListener.java

@@ -0,0 +1,116 @@
+/*
+ * This file is part of ViaVersion - https://github.com/ViaVersion/ViaVersion
+ * Copyright (C) 2016-2022 ViaVersion and contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package com.viaversion.viaversion.bukkit.listeners;
+
+import com.viaversion.viaversion.api.Via;
+import com.viaversion.viaversion.api.connection.ProtocolInfo;
+import com.viaversion.viaversion.api.connection.UserConnection;
+import com.viaversion.viaversion.bukkit.handlers.BukkitEncodeHandler;
+import com.viaversion.viaversion.bukkit.util.NMSUtil;
+import io.netty.channel.Channel;
+import org.bukkit.entity.Player;
+import org.bukkit.event.EventHandler;
+import org.bukkit.event.EventPriority;
+import org.bukkit.event.Listener;
+import org.bukkit.event.player.PlayerJoinEvent;
+import org.jetbrains.annotations.Nullable;
+
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.util.logging.Level;
+
+public class JoinListener implements Listener {
+
+    private static final Method GET_HANDLE;
+    private static final Field CONNECTION;
+    private static final Field NETWORK_MANAGER;
+    private static final Field CHANNEL;
+
+    static {
+        Method gh = null;
+        Field conn = null, nm = null, ch = null;
+        try {
+            gh = NMSUtil.obc("entity.CraftPlayer").getDeclaredMethod("getHandle");
+            conn = findField(gh.getReturnType(), "PlayerConnection", "ServerGamePacketListenerImpl");
+            nm = findField(conn.getType(), "NetworkManager", "Connection");
+            ch = findField(nm.getType(), "Channel");
+        } catch (NoSuchMethodException | NoSuchFieldException | ClassNotFoundException e) {
+            Via.getPlatform().getLogger().log(
+                    Level.WARNING,
+                    "Couldn't find reflection methods/fields to access Channel from player.\n" +
+                            "Login race condition fixer will be disabled.\n" +
+                            " Some plugins that use ViaAPI on join event may work incorrectly.", e);
+        }
+        GET_HANDLE = gh;
+        CONNECTION = conn;
+        NETWORK_MANAGER = nm;
+        CHANNEL = ch;
+    }
+
+    // Loosely search a field with any name, as long as it matches a type name.
+    private static Field findField(Class<?> cl, String... types) throws NoSuchFieldException {
+        for (Field field : cl.getDeclaredFields()) {
+            for (String type : types)
+                if (field.getType().getSimpleName().equals(type))
+                    return field;
+        }
+        throw new NoSuchFieldException(types[0]);
+    }
+
+    @EventHandler(priority = EventPriority.LOWEST)
+    public void onJoin(PlayerJoinEvent e) {
+        if (CHANNEL == null) return;
+        Player player = e.getPlayer();
+
+        Channel channel;
+        try {
+          channel = getChannel(player);
+        } catch (Exception ex) {
+            Via.getPlatform().getLogger().log(Level.WARNING, ex,
+                    () -> "Could not find Channel for logging-in player " + player.getUniqueId());
+            return;
+        }
+        // The connection has already closed, that was a quick leave
+        if (!channel.isOpen()) return;
+
+        UserConnection user = getUserConnection(channel);
+        if (user == null) {
+            Via.getPlatform().getLogger().log(Level.WARNING,
+                    "Could not find UserConnection for logging-in player {0}",
+                    player.getUniqueId());
+            return;
+        }
+        ProtocolInfo info = user.getProtocolInfo();
+        info.setUuid(player.getUniqueId());
+        info.setUsername(player.getName());
+        Via.getManager().getConnectionManager().onLoginSuccess(user);
+    }
+
+    private @Nullable UserConnection getUserConnection(Channel channel) {
+        BukkitEncodeHandler encoder = channel.pipeline().get(BukkitEncodeHandler.class);
+        return encoder != null ? encoder.getInfo() : null;
+    }
+
+    private Channel getChannel(Player player) throws Exception {
+        Object entityPlayer = GET_HANDLE.invoke(player);
+        Object pc = CONNECTION.get(entityPlayer);
+        Object nm = NETWORK_MANAGER.get(pc);
+        return (Channel) CHANNEL.get(nm);
+    }
+
+}

bukkit/src/main/java/com/viaversion/viaversion/bukkit/platform/BukkitViaLoader.java

@@ -24,6 +24,7 @@ import com.viaversion.viaversion.api.minecraft.item.Item;
 import com.viaversion.viaversion.api.platform.ViaPlatformLoader;
 import com.viaversion.viaversion.api.protocol.version.ProtocolVersion;
 import com.viaversion.viaversion.bukkit.classgenerator.ClassGenerator;
+import com.viaversion.viaversion.bukkit.listeners.JoinListener;
 import com.viaversion.viaversion.bukkit.listeners.UpdateListener;
 import com.viaversion.viaversion.bukkit.listeners.multiversion.PlayerSneakListener;
 import com.viaversion.viaversion.bukkit.listeners.protocol1_15to1_14_4.EntityToggleGlideListener;
@@ -80,6 +81,9 @@ public class BukkitViaLoader implements ViaPlatformLoader {
         // Update Listener
         registerListener(new UpdateListener());
 
+        // Login listener
+        registerListener(new JoinListener());
+
         /* Base Protocol */
         final ViaVersionPlugin plugin = (ViaVersionPlugin) Bukkit.getPluginManager().getPlugin("ViaVersion");
 

common/src/main/java/com/viaversion/viaversion/connection/ConnectionManagerImpl.java

@@ -20,6 +20,7 @@ package com.viaversion.viaversion.connection;
 import com.viaversion.viaversion.api.Via;
 import com.viaversion.viaversion.api.connection.ConnectionManager;
 import com.viaversion.viaversion.api.connection.UserConnection;
+import io.netty.channel.Channel;
 import io.netty.channel.ChannelFutureListener;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -37,17 +38,29 @@ public class ConnectionManagerImpl implements ConnectionManager {
     @Override
     public void onLoginSuccess(UserConnection connection) {
         Objects.requireNonNull(connection, "connection is null!");
-        connections.add(connection);
+        Channel channel = connection.getChannel();
+
+        // This user has already disconnected...
+        if (channel != null && !channel.isOpen()) return;
+
+        boolean newlyAdded = connections.add(connection);
 
         if (isFrontEnd(connection)) {
             UUID id = connection.getProtocolInfo().getUuid();
-            if (clients.put(id, connection) != null) {
+            UserConnection previous = clients.put(id, connection);
+            if (previous != null && previous != connection) {
                 Via.getPlatform().getLogger().warning("Duplicate UUID on frontend connection! (" + id + ")");
             }
         }
 
-        if (connection.getChannel() != null) {
+        if (channel != null) {
-            connection.getChannel().closeFuture().addListener((ChannelFutureListener) future -> onDisconnect(connection));
+            // We managed to add a user that had already disconnected!
+            // Let's clean up the mess here and now
+            if (!channel.isOpen()) {
+                onDisconnect(connection);
+            } else if (newlyAdded) { // Setup to clean-up on disconnect
+                channel.closeFuture().addListener((ChannelFutureListener) future -> onDisconnect(connection));
+            }
         }
     }