Fix wrong handling if the client and child server have the same protocol id

bungee/src/main/java/us/myles/ViaVersion/BungeePlugin.java

@@ -174,6 +174,7 @@ public class BungeePlugin extends Plugin implements ViaPlatform, Listener {
     public void onServerConnect(ServerConnectEvent e) throws NoSuchFieldException, IllegalAccessException {
         UserConnection user = Via.getManager().getConnection(e.getPlayer().getUniqueId());
         if (!user.has(BungeeStorage.class)) {
+            System.out.println("new storage");
             user.put(new BungeeStorage(user, e.getPlayer()));
         }
 
@@ -181,16 +182,14 @@ public class BungeePlugin extends Plugin implements ViaPlatform, Listener {
         List<Pair<Integer, Protocol>> protocols = ProtocolRegistry.getProtocolPath(user.get(ProtocolInfo.class).getProtocolVersion(), protocolId);
 
         // Check if ViaVersion can support that version
-        if (protocols != null) {
+        try {
-            try {
+            Object pendingConnection = ReflectionUtil.invoke(e.getPlayer(), "getPendingConnection");
-                Object pendingConnection = ReflectionUtil.invoke(e.getPlayer(), "getPendingConnection");
+            Object handshake = ReflectionUtil.invoke(pendingConnection, "getHandshake");
-                Object handshake = ReflectionUtil.invoke(pendingConnection, "getHandshake");
+            Method setProtocol = handshake.getClass().getDeclaredMethod("setProtocolVersion", int.class);
-                Method setProtocol = handshake.getClass().getDeclaredMethod("setProtocolVersion", int.class);
+            setProtocol.invoke(handshake, protocols == null ? user.get(ProtocolInfo.class).getProtocolVersion() : protocolId);
-                setProtocol.invoke(handshake, protocolId);
+            System.out.println("Changed server protocol id " + protocolId + " clientProtocol:" + user.get(ProtocolInfo.class).getProtocolVersion() + " path:" + protocols);
-                System.out.println("Changed server protocol id " + protocolId + " clientProtocol:" + user.get(ProtocolInfo.class).getProtocolVersion() + " path:" + protocols);
+        } catch (NoSuchMethodException | InvocationTargetException e1) {
-            } catch (NoSuchMethodException | InvocationTargetException e1) {
+            e1.printStackTrace();
-                e1.printStackTrace();
-            }
         }
     }
 

bungee/src/main/java/us/myles/ViaVersion/bungee/handlers/BungeeEncodeHandler.java

@@ -99,7 +99,7 @@ public class BungeeEncodeHandler extends MessageToMessageEncoder<ByteBuf> {
             if (player.getServer() != null) {
                 if (player.getServer() != null && !player.getServer().getInfo().getName().equals(storage.getCurrentServer())) {
 
-                    System.out.println("Server change " + player.getServer().getInfo().getName());
+                    System.out.println("Server change " + player.getServer().getInfo().getName() + " curr" + storage.getCurrentServer());
                     String serverName = player.getServer().getInfo().getName();
 
                     storage.setCurrentServer(serverName);
@@ -118,31 +118,36 @@ public class BungeeEncodeHandler extends MessageToMessageEncoder<ByteBuf> {
                     List<Pair<Integer, Protocol>> protocols = ProtocolRegistry.getProtocolPath(info.getProtocolVersion(), protocolId);
                     System.out.println(info.getProtocolVersion() + ">" + protocolId + " " + protocols);
                     ProtocolPipeline pipeline = viaConnection.get(ProtocolInfo.class).getPipeline();
-                    if (protocols != null) {
-                        viaConnection.clearStoredObjects();
-                        pipeline.cleanPipes();
 
+                    viaConnection.clearStoredObjects();
+                    pipeline.cleanPipes();
+
+                    if (protocols != null)
                         for (Pair<Integer, Protocol> prot : protocols) {
                             pipeline.add(prot.getValue());
                         }
-                        viaConnection.put(info);
-                        viaConnection.setActive(true);
-
-                        // Init all protocols TODO check if this can get moved up to the previous for loop, and doesn't require the pipeline to already exist.
-                        for (Pair<Integer, Protocol> protocol : protocols) {
-                            protocol.getValue().init(viaConnection);
-                        }
 
 
-                        Object wrapper = ReflectionUtil.get(player, "ch", Object.class);
+                    viaConnection.put(info);
-                        wrapper.getClass().getDeclaredMethod("setVersion", int.class).invoke(wrapper, protocolId);
+                    viaConnection.put(storage);
-                        ReflectionUtil.invoke(player, "init");
+
-                    } else {
+                    viaConnection.setActive(protocols != null);
-                        viaConnection.setActive(false);
+
+                    // Init all protocols TODO check if this can get moved up to the previous for loop, and doesn't require the pipeline to already exist.
+                    for (Protocol protocol : pipeline.pipes()) {
+                        protocol.init(viaConnection);
                     }
+
+                    Object wrapper = ReflectionUtil.get(player, "ch", Object.class);
+                    wrapper.getClass().getDeclaredMethod("setVersion", int.class).invoke(wrapper, protocolId);
+//                    ReflectionUtil.invoke(player, "init");
+
+                    Object entityMap = Class.forName("net.md_5.bungee.entitymap.EntityMap").getDeclaredMethod("getEntityMap", int.class).invoke(null, protocolId);
+                    ReflectionUtil.set(player, "entityRewrite", entityMap);
+
+                    System.out.println("VERSION:" + ((net.md_5.bungee.UserConnection) player).getPendingConnection().getVersion());
                 }
             }
         }
     }
-
 }

bungee/src/main/java/us/myles/ViaVersion/bungee/service/ProtocolDetectorService.java

@@ -29,13 +29,20 @@ public class ProtocolDetectorService implements Runnable {
     @Override
     public void run() {
         for (final Map.Entry<String, ServerInfo> lists : plugin.getProxy().getServers().entrySet()) {
-            lists.getValue().ping(new Callback<ServerPing>() {
+            updateProtocolInfo(lists.getKey(), lists.getValue());
-                @Override
-                public void done(ServerPing serverPing, Throwable throwable) {
-                    if (throwable == null)
-                        protocolIds.put(lists.getKey(), serverPing.getVersion().getProtocol());
-                }
-            });
         }
     }
+
+    private void updateProtocolInfo(final String key, ServerInfo value) {
+        value.ping(new Callback<ServerPing>() {
+            @Override
+            public void done(ServerPing serverPing, Throwable throwable) {
+                if (throwable == null)
+                    protocolIds.put(key, serverPing.getVersion().getProtocol());
+                else
+                    throwable.printStackTrace();
+            }
+        });
+    }
+
 }