diff --git a/server/src/main/java/com/soulfiremc/server/ServerCommandManager.java b/server/src/main/java/com/soulfiremc/server/ServerCommandManager.java index 4fe52dd07..a93949a8f 100644 --- a/server/src/main/java/com/soulfiremc/server/ServerCommandManager.java +++ b/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 { private static final ThreadLocal> 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 dispatcher = new CommandDispatcher<>(); private final SoulFireServer soulFireServer; @@ -100,6 +103,42 @@ private static String getCurrentUsername() { } } + public static void putInstanceIds(List 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 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 getBotNames() { + return List.of(COMMAND_CONTEXT.get().get(BOT_NAMES_KEY).split(",")); + } + + public static List 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 getCompletionSuggestions(String command, ServerCommandSource .map(Suggestion::getText) .toList(); } finally { - COMMAND_CONTEXT.get().clear(); + clearContext(); } } diff --git a/server/src/main/java/com/soulfiremc/server/grpc/CommandServiceImpl.java b/server/src/main/java/com/soulfiremc/server/grpc/CommandServiceImpl.java index bd0d748a4..452d30c64 100644 --- a/server/src/main/java/com/soulfiremc/server/grpc/CommandServiceImpl.java +++ b/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()); diff --git a/server/src/main/java/com/soulfiremc/server/plugins/ChatControl.java b/server/src/main/java/com/soulfiremc/server/plugins/ChatControl.java index 6a7f58e3d..6afeda9f4 100644 --- a/server/src/main/java/com/soulfiremc/server/plugins/ChatControl.java +++ b/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() diff --git a/server/src/main/java/com/soulfiremc/server/plugins/POVServer.java b/server/src/main/java/com/soulfiremc/server/plugins/POVServer.java index 688d79aec..237599498 100644 --- a/server/src/main/java/com/soulfiremc/server/plugins/POVServer.java +++ b/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()