Restrict instance manager context for commands

server/src/main/java/com/soulfiremc/server/ServerCommandManager.java

@@ -75,6 +75,7 @@
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Function;
 import java.util.function.ToIntFunction;
+import java.util.stream.Collectors;
 
 import static com.mojang.brigadier.CommandDispatcher.ARGUMENT_SEPARATOR;
 import static com.soulfiremc.server.brigadier.ServerBrigadierHelper.*;
@@ -86,6 +87,8 @@
 public class ServerCommandManager implements PlatformCommandManager<ServerCommandSource> {
   private static final ThreadLocal<Map<String, String>> COMMAND_CONTEXT =
     ThreadLocal.withInitial(Object2ObjectOpenHashMap::new);
+  private static final String INSTANCE_IDS_KEY = "instance_ids";
+  private static final String BOT_NAMES_KEY = "bot_names";
   @Getter
   private final CommandDispatcher<ServerCommandSource> dispatcher = new CommandDispatcher<>();
   private final SoulFireServer soulFireServer;
@@ -100,6 +103,42 @@ private static String getCurrentUsername() {
     }
   }
 
+  public static void putInstanceIds(List<UUID> ids) {
+    if (hasInstanceIds()) {
+      throw new IllegalStateException("Instance ids already set!");
+    }
+
+    COMMAND_CONTEXT.get().put(INSTANCE_IDS_KEY, ids.stream().map(UUID::toString).collect(Collectors.joining(",")));
+  }
+
+  public static void putBotNames(List<String> names) {
+    if (hasBotNames()) {
+      throw new IllegalStateException("Bot names already set!");
+    }
+
+    COMMAND_CONTEXT.get().put(BOT_NAMES_KEY, String.join(",", names));
+  }
+
+  public static boolean hasBotNames() {
+    return COMMAND_CONTEXT.get().containsKey(BOT_NAMES_KEY);
+  }
+
+  public static boolean hasInstanceIds() {
+    return COMMAND_CONTEXT.get().containsKey(INSTANCE_IDS_KEY);
+  }
+
+  public static List<String> getBotNames() {
+    return List.of(COMMAND_CONTEXT.get().get(BOT_NAMES_KEY).split(","));
+  }
+
+  public static List<UUID> getInstanceIds() {
+    return Arrays.stream(COMMAND_CONTEXT.get().get(INSTANCE_IDS_KEY).split(",")).map(UUID::fromString).toList();
+  }
+
+  public static void clearContext() {
+    COMMAND_CONTEXT.get().clear();
+  }
+
   @PostConstruct
   public void postConstruct() {
     // Help
@@ -815,7 +854,7 @@ public void postConstruct() {
               helpRedirect(
                 "Instead of running a command for selected bots, run it for a specific list of bots. Use a comma to separate the names",
                 c -> {
-                  COMMAND_CONTEXT.get().put("bot_names", StringArgumentType.getString(c, "bot_names"));
+                  putBotNames(List.of(StringArgumentType.getString(c, "bot_names").split(",")));
                   return Collections.singleton(c.getSource());
                 }),
               false)));
@@ -841,7 +880,11 @@ public void postConstruct() {
               helpRedirect(
                 "Instead of running a command for selected attacks, run it for a specific list of attacks. Use a comma to separate the ids",
                 c -> {
-                  COMMAND_CONTEXT.get().put("attack_ids", StringArgumentType.getString(c, "attack_ids"));
+                  putInstanceIds(Arrays.stream(StringArgumentType.getString(c, "attack_ids").split(","))
+                    .map(UUIDHelper::tryParseUniqueId)
+                    .filter(Optional::isPresent)
+                    .map(Optional::get)
+                    .toList());
                   return Collections.singleton(c.getSource());
                 }),
               false)));
@@ -899,11 +942,9 @@ private int forEveryAttack(
 
     var resultCode = Command.SINGLE_SUCCESS;
     for (var instanceManager : soulFireServer.instances().values()) {
-      if (COMMAND_CONTEXT.get().containsKey("attack_ids")
-        && Arrays.stream(COMMAND_CONTEXT.get().get("attack_ids").split(","))
-        .map(UUIDHelper::tryParseUniqueId)
-        .filter(Optional::isPresent)
-        .noneMatch(i -> i.orElseThrow().equals(instanceManager.id()))) {
+      if (hasInstanceIds() && getInstanceIds()
+        .stream()
+        .noneMatch(instanceManager.id()::equals)) {
         continue;
       }
 
@@ -960,9 +1001,9 @@ private int forEveryBot(
       instanceManager -> {
         var resultCode = Command.SINGLE_SUCCESS;
         for (var bot : instanceManager.botConnections().values()) {
-          if (COMMAND_CONTEXT.get().containsKey("bot_names")
-            && Arrays.stream(COMMAND_CONTEXT.get().get("bot_names").split(","))
-            .noneMatch(s -> s.equals(bot.accountName()))) {
+          if (hasBotNames() && getBotNames()
+            .stream()
+            .noneMatch(bot.accountName()::equals)) {
             continue;
           }
 
@@ -1003,7 +1044,7 @@ public int execute(String command, ServerCommandSource source) {
       source.sendWarn(e.getMessage());
       return Command.SINGLE_SUCCESS;
     } finally {
-      COMMAND_CONTEXT.get().clear();
+      clearContext();
     }
   }
 
@@ -1018,7 +1059,7 @@ public List<String> getCompletionSuggestions(String command, ServerCommandSource
         .map(Suggestion::getText)
         .toList();
     } finally {
-      COMMAND_CONTEXT.get().clear();
+      clearContext();
     }
   }
 

server/src/main/java/com/soulfiremc/server/grpc/CommandServiceImpl.java

@@ -27,6 +27,8 @@
 import lombok.extern.slf4j.Slf4j;
 
 import javax.inject.Inject;
+import java.util.List;
+import java.util.UUID;
 
 @Slf4j
 @RequiredArgsConstructor(onConstructor_ = @Inject)
@@ -39,6 +41,7 @@ public void executeCommand(
     ServerRPCConstants.USER_CONTEXT_KEY.get().hasPermissionOrThrow(Permissions.COMMAND_EXECUTION);
 
     try {
+      ServerCommandManager.putInstanceIds(List.of(UUID.fromString(request.getInstanceId())));
       var code = serverCommandManager.execute(request.getCommand(), ServerRPCConstants.USER_CONTEXT_KEY.get());
 
       responseObserver.onNext(CommandResponse.newBuilder().setCode(code).build());

server/src/main/java/com/soulfiremc/server/plugins/ChatControl.java

@@ -33,6 +33,8 @@
 import net.lenni0451.lambdaevents.EventHandler;
 import org.pf4j.Extension;
 
+import java.util.List;
+
 @Extension
 public class ChatControl extends InternalPlugin {
   public ChatControl() {
@@ -64,6 +66,8 @@ public static void onChat(ChatMessageReceiveEvent event) {
     plainMessage = plainMessage.substring(prefixIndex);
     var command = plainMessage.substring(prefix.length());
     connection.logger().info("[ChatControl] Executing command: \"{}\"", command);
+
+    ServerCommandManager.putInstanceIds(List.of(connection.instanceManager().id()));
     var code = connection.instanceManager()
       .soulFireServer()
       .injector()

server/src/main/java/com/soulfiremc/server/plugins/POVServer.java

@@ -993,6 +993,8 @@ public void packetSent(Session botSession, Packet packet) {
                 if (message.startsWith(prefix)) {
                   var command = message.substring(prefix.length());
                   var source = new PovServerUser(clientSession, clientSession.getFlag(MinecraftConstants.PROFILE_KEY).getName());
+
+                  ServerCommandManager.putInstanceIds(List.of(instanceManager.id()));
                   var code = instanceManager
                     .soulFireServer()
                     .injector()