From 79006ac40f0c83d2bc6ef5b7a296948ef003cfb1 Mon Sep 17 00:00:00 2001
From: Zach Levis <zml@aoeu.xyz>
Date: Fri, 1 Jan 2021 15:14:34 -0800
Subject: [PATCH] fabric: Add registration environment meta

---
 .../fabric/FabricCommandManager.java              | 15 ++++++++++++++-
 .../fabric/FabricCommandRegistrationHandler.java  | 11 +++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/cloud-minecraft/cloud-fabric/src/main/java/cloud/commandframework/fabric/FabricCommandManager.java b/cloud-minecraft/cloud-fabric/src/main/java/cloud/commandframework/fabric/FabricCommandManager.java
index da40e9b0..f3944475 100644
--- a/cloud-minecraft/cloud-fabric/src/main/java/cloud/commandframework/fabric/FabricCommandManager.java
+++ b/cloud-minecraft/cloud-fabric/src/main/java/cloud/commandframework/fabric/FabricCommandManager.java
@@ -33,6 +33,7 @@ import cloud.commandframework.execution.AsynchronousCommandExecutionCoordinator;
 import cloud.commandframework.execution.CommandExecutionCoordinator;
 import cloud.commandframework.meta.CommandMeta;
 import cloud.commandframework.meta.SimpleCommandMeta;
+import net.minecraft.server.command.CommandManager.RegistrationEnvironment;
 import net.minecraft.server.command.CommandOutput;
 import net.minecraft.server.command.ServerCommandSource;
 import net.minecraft.text.LiteralText;
@@ -43,6 +44,17 @@ import org.checkerframework.checker.nullness.qual.NonNull;
 import java.util.function.Function;
 
 public class FabricCommandManager<C> extends CommandManager<C> implements BrigadierManagerHolder<C> {
+
+    /**
+     * A meta attribute specifying which environments a command should be registered in.
+     *
+     * <p>The default value is {@link RegistrationEnvironment#ALL}.</p>
+     */
+    public static final CommandMeta.Key<RegistrationEnvironment> META_REGISTRATION_ENVIRONMENT = CommandMeta.Key.of(
+            RegistrationEnvironment.class,
+            "cloud:registration-environment"
+    );
+
     private final Function<ServerCommandSource, C> commandSourceMapper;
     private final Function<C, ServerCommandSource> backwardsCommandSourceMapper;
     private final CloudBrigadierManager<C, ServerCommandSource> brigadierManager;
@@ -100,7 +112,7 @@ public class FabricCommandManager<C> extends CommandManager<C> implements Brigad
                                 null,
                                 null
                         )),
-                        this.getCaptionRegistry()
+                        this
                 ));
 
         ((FabricCommandRegistrationHandler<C>) this.getCommandRegistrationHandler()).initialize(this);
@@ -149,6 +161,7 @@ public class FabricCommandManager<C> extends CommandManager<C> implements Brigad
         return this.brigadierManager;
     }
 
+    /* transition state to prevent further registration */
     final void registrationCalled() {
         this.transitionOrThrow(RegistrationState.REGISTERING, RegistrationState.AFTER_REGISTRATION);
     }
diff --git a/cloud-minecraft/cloud-fabric/src/main/java/cloud/commandframework/fabric/FabricCommandRegistrationHandler.java b/cloud-minecraft/cloud-fabric/src/main/java/cloud/commandframework/fabric/FabricCommandRegistrationHandler.java
index a3a42921..c48c59b6 100644
--- a/cloud-minecraft/cloud-fabric/src/main/java/cloud/commandframework/fabric/FabricCommandRegistrationHandler.java
+++ b/cloud-minecraft/cloud-fabric/src/main/java/cloud/commandframework/fabric/FabricCommandRegistrationHandler.java
@@ -32,6 +32,7 @@ import com.mojang.brigadier.tree.CommandNode;
 import com.mojang.brigadier.tree.RootCommandNode;
 import net.fabricmc.fabric.api.command.v1.CommandRegistrationCallback;
 import net.minecraft.server.command.CommandManager;
+import net.minecraft.server.command.CommandManager.RegistrationEnvironment;
 import net.minecraft.server.command.ServerCommandSource;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 import org.checkerframework.checker.nullness.qual.NonNull;
@@ -65,6 +66,16 @@ public final class FabricCommandRegistrationHandler<C> implements CommandRegistr
     private void registerAllCommands(final CommandDispatcher<ServerCommandSource> dispatcher, final boolean isDedicated) {
         this.commandManager.registrationCalled();
         for (final Command<C> command : this.registeredCommands) {
+            /* Only register commands in the declared environment */
+            final RegistrationEnvironment env = command.getCommandMeta().getOrDefault(
+                FabricCommandManager.META_REGISTRATION_ENVIRONMENT,
+                RegistrationEnvironment.ALL
+            );
+
+            if ((env == RegistrationEnvironment.INTEGRATED && isDedicated)
+                    || (env == RegistrationEnvironment.DEDICATED && !isDedicated)) {
+                continue;
+            }
             this.registerCommand(dispatcher.getRoot(), command);
         }
     }