Added function to unregister packets on plugin disable. [Dev System Only] (#83)

* Added function to unregister packets on plugin disable.

* Update src/main/java/net/william278/velocitab/packet/PacketRegistration.java

Co-authored-by: William <will27528@gmail.com>

* Added null check + changed style

* Update src/main/java/net/william278/velocitab/Velocitab.java

Co-authored-by: William <will27528@gmail.com>

* Fix indentation problem

---------

Co-authored-by: William <will27528@gmail.com>

build.gradle

@@ -34,6 +34,8 @@ dependencies {
     compileOnly 'net.luckperms:api:5.4'
     compileOnly 'io.github.miniplaceholders:miniplaceholders-api:2.0.0'
     compileOnly 'net.william278:PAPIProxyBridge:1.3'
+    compileOnly 'it.unimi.dsi:fastutil:8.5.12'
+
 
     implementation 'org.apache.commons:commons-text:1.10.0'
     implementation 'net.william278:annotaml:2.0.7'

src/main/java/net/william278/velocitab/Velocitab.java

@@ -23,6 +23,7 @@ import com.google.inject.Inject;
 import com.velocitypowered.api.command.BrigadierCommand;
 import com.velocitypowered.api.event.Subscribe;
 import com.velocitypowered.api.event.proxy.ProxyInitializeEvent;
+import com.velocitypowered.api.event.proxy.ProxyShutdownEvent;
 import com.velocitypowered.api.plugin.Plugin;
 import com.velocitypowered.api.plugin.PluginContainer;
 import com.velocitypowered.api.plugin.PluginDescription;
@@ -91,6 +92,12 @@ public class Velocitab {
         logger.info("Successfully enabled Velocitab");
     }
 
+    @Subscribe
+    public void onProxyShutdown(@NotNull ProxyShutdownEvent event) {
+        disableScoreboardManager();
+        logger.info("Successfully disabled Velocitab");
+    }
+
     @NotNull
     public ProxyServer getServer() {
         return server;
@@ -148,6 +155,12 @@ public class Velocitab {
         }
     }
 
+    private void disableScoreboardManager() {
+        if (scoreboardManager !=null && settings.isSortPlayers()) {
+            scoreboardManager.unregisterPacket();
+        }
+    }
+
     @NotNull
     public Optional<ScoreboardManager> getScoreboardManager() {
         return Optional.ofNullable(scoreboardManager);

src/main/java/net/william278/velocitab/packet/PacketRegistration.java

@@ -23,13 +23,19 @@ import com.velocitypowered.api.network.ProtocolVersion;
 import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import com.velocitypowered.proxy.protocol.ProtocolUtils;
 import com.velocitypowered.proxy.protocol.StateRegistry;
+import io.netty.util.collection.IntObjectMap;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import net.william278.velocitab.Velocitab;
 import org.jetbrains.annotations.NotNull;
 
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
+import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.Objects;
 import java.util.function.Supplier;
 
 // Based on VPacketEvents PacketRegistration API
@@ -76,12 +82,42 @@ public final class PacketRegistration<P extends MinecraftPacket> {
                 final StateRegistry.PacketRegistry packetRegistry = direction == ProtocolUtils.Direction.CLIENTBOUND
                         ? (StateRegistry.PacketRegistry) STATE_REGISTRY$clientBound.invoke(stateRegistry)
                         : (StateRegistry.PacketRegistry) STATE_REGISTRY$serverBound.invoke(stateRegistry);
+
                 PACKET_REGISTRY$register.invoke(
                         packetRegistry,
                         packetClass,
                         packetSupplier,
                         mappings.toArray(StateRegistry.PacketMapping[]::new)
                 );
+
+            } catch (Throwable t) {
+                throw new RuntimeException(t);
+            }
+        }
+
+        @SuppressWarnings("unchecked")
+        public void unregister() {
+            try {
+                final StateRegistry.PacketRegistry packetRegistry = direction == ProtocolUtils.Direction.CLIENTBOUND
+                        ? (StateRegistry.PacketRegistry) STATE_REGISTRY$clientBound.invoke(stateRegistry)
+                        : (StateRegistry.PacketRegistry) STATE_REGISTRY$serverBound.invoke(stateRegistry);
+
+                Map<ProtocolVersion, StateRegistry.PacketRegistry.ProtocolRegistry> versions = (Map<ProtocolVersion, StateRegistry.PacketRegistry.ProtocolRegistry>) PACKET_REGISTRY$versions.invoke(packetRegistry);
+                versions.forEach((protocolVersion, protocolRegistry) -> {
+                    try {
+                        IntObjectMap<Supplier<?>> packetIdToSupplier = (IntObjectMap<Supplier<?>>) PACKET_REGISTRY$packetIdToSupplier.invoke(protocolRegistry);
+                        Object2IntMap<Class<?>> packetClassToId = (Object2IntMap<Class<?>>) PACKET_REGISTRY$packetClassToId.invoke(protocolRegistry);
+                        packetIdToSupplier.keySet().stream()
+                                .filter(supplier -> packetIdToSupplier.get(supplier).get().getClass().equals(packetClass))
+                                .forEach(packetIdToSupplier::remove);
+                        packetClassToId.values().intStream()
+                                .filter(id -> Objects.equals(packetClassToId.getInt(packetClass), id))
+                                .forEach(packetClassToId::removeInt);
+                    } catch (Throwable t) {
+                        throw new RuntimeException(t);
+                    }
+                });
+
             } catch (Throwable t) {
                 throw new RuntimeException(t);
             }
@@ -98,8 +134,14 @@ public final class PacketRegistration<P extends MinecraftPacket> {
         private static final MethodHandle STATE_REGISTRY$clientBound;
         private static final MethodHandle STATE_REGISTRY$serverBound;
         private static final MethodHandle PACKET_REGISTRY$register;
+        private static final MethodHandle PACKET_REGISTRY$packetIdToSupplier;
+        private static final MethodHandle PACKET_REGISTRY$packetClassToId;
+        private static final MethodHandle PACKET_REGISTRY$versions;
         private static final MethodHandle PACKET_MAPPING$map;
 
+
+
+
         static {
             final MethodHandles.Lookup lookup = MethodHandles.lookup();
             try {
@@ -113,6 +155,13 @@ public final class PacketRegistration<P extends MinecraftPacket> {
                 final MethodHandles.Lookup packetRegistryLookup = MethodHandles.privateLookupIn(StateRegistry.PacketRegistry.class, lookup);
                 final MethodType registerType = MethodType.methodType(void.class, Class.class, Supplier.class, StateRegistry.PacketMapping[].class);
                 PACKET_REGISTRY$register = packetRegistryLookup.findVirtual(StateRegistry.PacketRegistry.class, "register", registerType);
+                PACKET_REGISTRY$versions = packetRegistryLookup.findGetter(StateRegistry.PacketRegistry.class, "versions", Map.class);
+
+                final MethodHandles.Lookup protocolRegistryLookup = MethodHandles.privateLookupIn(StateRegistry.PacketRegistry.ProtocolRegistry.class, lookup);
+                PACKET_REGISTRY$packetIdToSupplier = protocolRegistryLookup.findGetter(StateRegistry.PacketRegistry.ProtocolRegistry.class, "packetIdToSupplier", IntObjectMap.class);
+                PACKET_REGISTRY$packetClassToId = protocolRegistryLookup.findGetter(StateRegistry.PacketRegistry.ProtocolRegistry.class, "packetClassToId", Object2IntMap.class);
+
+
             } catch (Throwable e) {
                 throw new RuntimeException(e);
             }

src/main/java/net/william278/velocitab/packet/ScoreboardManager.java

@@ -34,6 +34,7 @@ import static com.velocitypowered.api.network.ProtocolVersion.*;
 
 public class ScoreboardManager {
 
+    private PacketRegistration<UpdateTeamsPacket> packetRegistration;
     private final Velocitab plugin;
     private final Map<UUID, List<String>> createdTeams;
     private final Map<UUID, Map<String, String>> roleMappings;
@@ -98,7 +99,7 @@ public class ScoreboardManager {
 
     public void registerPacket() {
         try {
-            PacketRegistration.of(UpdateTeamsPacket.class)
+            packetRegistration = PacketRegistration.of(UpdateTeamsPacket.class)
                     .direction(ProtocolUtils.Direction.CLIENTBOUND)
                     .packetSupplier(UpdateTeamsPacket::new)
                     .stateRegistry(StateRegistry.PLAY)
@@ -108,12 +109,23 @@ public class ScoreboardManager {
                     .mapping(0x55, MINECRAFT_1_17, false)
                     .mapping(0x58, MINECRAFT_1_19_1, false)
                     .mapping(0x56, MINECRAFT_1_19_3, false)
-                    .mapping(0x5A, MINECRAFT_1_19_4, false)
+                    .mapping(0x5A, MINECRAFT_1_19_4, false);
-                    .register();
+            packetRegistration.register();
         } catch (Throwable e) {
             plugin.log(Level.ERROR, "Failed to register UpdateTeamsPacket", e);
         }
     }
 
+    public void unregisterPacket() {
+        if(packetRegistration==null) {
+            return;
+        }
+        try {
+            packetRegistration.unregister();
+        } catch (Throwable e) {
+            plugin.log(Level.ERROR, "Failed to unregister UpdateTeamsPacket", e);
+        }
+    }
+
 
 }