Add experimental removal of commands when unloading plugins

Velocity/src/main/java/net/frankheijden/serverutils/velocity/ServerUtils.java

@@ -3,6 +3,9 @@ package net.frankheijden.serverutils.velocity;
 import co.aikar.commands.CommandCompletions;
 import co.aikar.commands.VelocityCommandCompletionContext;
 import co.aikar.commands.VelocityCommandManager;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Multimaps;
 import com.google.inject.Inject;
 import com.google.inject.name.Named;
 import com.velocitypowered.api.event.Subscribe;
@@ -19,6 +22,7 @@ import net.frankheijden.serverutils.velocity.commands.CommandPlugins;
 import net.frankheijden.serverutils.velocity.commands.CommandServerUtils;
 import net.frankheijden.serverutils.velocity.entities.VelocityPlugin;
 import net.frankheijden.serverutils.velocity.managers.VelocityPluginManager;
+import net.frankheijden.serverutils.velocity.reflection.RVelocityCommandManager;
 import org.bstats.velocity.Metrics;
 import org.slf4j.Logger;
 
@@ -56,6 +60,20 @@ public class ServerUtils {
     @Named("serverutils")
     private PluginContainer pluginContainer;
 
+    private final Multimap<String, String> pluginCommands = Multimaps.synchronizedSetMultimap(HashMultimap.create());
+
+    /**
+     * Initialises ServerUtils.
+     */
+    @Inject
+    public ServerUtils(ProxyServer proxy) {
+        RVelocityCommandManager.proxyRegistrars(
+                proxy,
+                getClass().getClassLoader(),
+                (container, meta) -> pluginCommands.putAll(container.getDescription().getId(), meta.getAliases())
+        );
+    }
+
     /**
      * Initialises and enables ServerUtils.
      */
@@ -108,6 +126,10 @@ public class ServerUtils {
         return plugin;
     }
 
+    public Multimap<String, String> getPluginCommands() {
+        return pluginCommands;
+    }
+
     public void reload() {
         new Config("config.toml", CONFIG_RESOURCE);
         new Messenger("messages.toml", MESSAGES_RESOURCE);

Velocity/src/main/java/net/frankheijden/serverutils/velocity/managers/VelocityPluginManager.java

@@ -216,9 +216,12 @@ public class VelocityPluginManager extends AbstractPluginManager<PluginContainer
             task.cancel();
         }
 
-        // TODO: unload commands of plugin
+        String pluginId = plugin.getDescription().getId();
+        for (String alias : ServerUtils.getInstance().getPluginCommands().removeAll(pluginId)) {
+            proxy.getCommandManager().unregister(alias);
+        }
 
-        RVelocityPluginManager.getPlugins(proxy.getPluginManager()).remove(plugin.getDescription().getId());
+        RVelocityPluginManager.getPlugins(proxy.getPluginManager()).remove(pluginId);
         RVelocityPluginManager.getPluginInstances(proxy.getPluginManager()).remove(pluginInstance);
 
         List<Closeable> closeables = new ArrayList<>();

Velocity/src/main/java/net/frankheijden/serverutils/velocity/reflection/RVelocityCommandManager.java

@@ -1,9 +1,22 @@
 package net.frankheijden.serverutils.velocity.reflection;
 
 import com.mojang.brigadier.CommandDispatcher;
+import com.velocitypowered.api.command.Command;
 import com.velocitypowered.api.command.CommandManager;
+import com.velocitypowered.api.command.CommandMeta;
 import com.velocitypowered.api.command.CommandSource;
+import com.velocitypowered.api.plugin.PluginContainer;
+import com.velocitypowered.api.proxy.ProxyServer;
 import dev.frankheijden.minecraftreflection.MinecraftReflection;
+import dev.frankheijden.minecraftreflection.Reflection;
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.BiConsumer;
+import net.frankheijden.serverutils.velocity.utils.ReflectionUtils;
 
 public class RVelocityCommandManager {
 
@@ -15,4 +28,82 @@ public class RVelocityCommandManager {
     public static CommandDispatcher<CommandSource> getDispatcher(CommandManager manager) {
         return reflection.get(manager, "dispatcher");
     }
+
+    /**
+     * Proxies the registrars.
+     */
+    @SuppressWarnings("rawtypes")
+    public static void proxyRegistrars(
+            ProxyServer proxy,
+            ClassLoader loader,
+            BiConsumer<PluginContainer, CommandMeta> registrationConsumer
+    ) {
+        List<Object> proxiedRegistrars = new ArrayList<>();
+
+        Class<?> commandRegistrarClass;
+        try {
+            commandRegistrarClass = Class.forName("com.velocitypowered.proxy.command.registrar.CommandRegistrar");
+        } catch (ClassNotFoundException ex) {
+            ex.printStackTrace();
+            return;
+        }
+
+        for (Object registrar : (List) reflection.get(proxy.getCommandManager(), "registrars")) {
+            proxiedRegistrars.add(Proxy.newProxyInstance(
+                    loader,
+                    new Class[]{ commandRegistrarClass },
+                    new CommandRegistrarInvocationHandler(
+                            proxy,
+                            registrar,
+                            registrationConsumer
+                    )
+            ));
+        }
+
+        Field registrarsField = Reflection.getAccessibleField(reflection.getClazz(), "registrars");
+        ReflectionUtils.doPrivilegedWithUnsafe(unsafe -> {
+            long offset = unsafe.objectFieldOffset(registrarsField);
+            unsafe.putObject(proxy.getCommandManager(), offset, proxiedRegistrars);
+        });
+    }
+
+    public static final class CommandRegistrarInvocationHandler implements InvocationHandler {
+
+        private final ProxyServer proxy;
+        private final Object commandRegistrar;
+        private final BiConsumer<PluginContainer, CommandMeta> registrationConsumer;
+
+        /**
+         * Constructs  a new {@link CommandRegistrarInvocationHandler}.
+         */
+        public CommandRegistrarInvocationHandler(
+                ProxyServer proxy,
+                Object commandRegistrar,
+                BiConsumer<PluginContainer, CommandMeta> registrationConsumer
+        ) {
+            this.proxy = proxy;
+            this.commandRegistrar = commandRegistrar;
+            this.registrationConsumer = registrationConsumer;
+        }
+
+        @Override
+        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+            Object obj = method.invoke(commandRegistrar, args);
+            if (method.getName().equals("register")) {
+                handleRegisterMethod((CommandMeta) args[0], (Command) args[1]);
+            }
+            return obj;
+        }
+
+        private void handleRegisterMethod(CommandMeta commandMeta, Command command) {
+            ClassLoader classLoader = command.getClass().getClassLoader();
+
+            for (PluginContainer container : proxy.getPluginManager().getPlugins()) {
+                if (container.getInstance().filter(i -> i.getClass().getClassLoader() == classLoader).isPresent()) {
+                    registrationConsumer.accept(container, commandMeta);
+                    break;
+                }
+            }
+        }
+    }
 }

Velocity/src/main/java/net/frankheijden/serverutils/velocity/utils/ReflectionUtils.java

@@ -0,0 +1,41 @@
+package net.frankheijden.serverutils.velocity.utils;
+
+import dev.frankheijden.minecraftreflection.Reflection;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.security.AccessController;
+import java.security.PrivilegedAction;
+import java.util.function.Consumer;
+import sun.misc.Unsafe;
+
+public class ReflectionUtils {
+
+    private static MethodHandle theUnsafeFieldMethodHandle;
+
+    static {
+        try {
+            theUnsafeFieldMethodHandle = MethodHandles.lookup().unreflectGetter(Reflection.getAccessibleField(
+                    Unsafe.class,
+                    "theUnsafe"
+            ));
+        } catch (Throwable th) {
+            th.printStackTrace();
+        }
+    }
+
+    private ReflectionUtils() {}
+
+    /**
+     * Performs a privileged action while accessing {@link Unsafe}.
+     */
+    public static void doPrivilegedWithUnsafe(Consumer<Unsafe> privilegedAction) {
+        AccessController.doPrivileged((PrivilegedAction<Void>) () -> {
+            try {
+                privilegedAction.accept((Unsafe) theUnsafeFieldMethodHandle.invoke());
+            } catch (Throwable th) {
+                th.printStackTrace();
+            }
+            return null;
+        });
+    }
+}