diff --git a/src/main/java/com/mojang/brigadier/CommandDispatcher.java b/src/main/java/com/mojang/brigadier/CommandDispatcher.java index c82d9d3..9d39604 100644 --- a/src/main/java/com/mojang/brigadier/CommandDispatcher.java +++ b/src/main/java/com/mojang/brigadier/CommandDispatcher.java @@ -19,6 +19,7 @@ import java.util.stream.Collectors; public class CommandDispatcher { public static final SimpleCommandExceptionType ERROR_UNKNOWN_COMMAND = new SimpleCommandExceptionType("command.unknown", "Unknown command"); + public static final SimpleCommandExceptionType ERROR_IMPERMISSIBLE = new SimpleCommandExceptionType("command.impermissible", "Command not allowed"); public static final String ARGUMENT_SEPARATOR = " "; private static final String USAGE_OPTIONAL_OPEN = "["; private static final String USAGE_OPTIONAL_CLOSE = "]"; @@ -45,8 +46,13 @@ public class CommandDispatcher { private CommandContext parseNodes(CommandNode node, String command, CommandContextBuilder contextBuilder) throws CommandException { CommandException exception = null; + final S source = contextBuilder.getSource(); for (CommandNode child : node.getChildren()) { + if (!child.canUse(source)) { + exception = ERROR_IMPERMISSIBLE.create(); + continue; + } try { CommandContextBuilder context = contextBuilder.copy(); String remaining = child.parse(command, context); @@ -113,7 +119,11 @@ public class CommandDispatcher { } private Set findSuggestions(CommandNode node, String command, CommandContextBuilder contextBuilder, Set result) { + final S source = contextBuilder.getSource(); for (CommandNode child : node.getChildren()) { + if (!child.canUse(source)) { + continue; + } try { CommandContextBuilder context = contextBuilder.copy(); String remaining = child.parse(command, context); diff --git a/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java b/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java index a72ce7a..5b20952 100644 --- a/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java +++ b/src/main/java/com/mojang/brigadier/builder/ArgumentBuilder.java @@ -5,10 +5,12 @@ import com.mojang.brigadier.tree.CommandNode; import com.mojang.brigadier.tree.RootCommandNode; import java.util.Collection; +import java.util.function.Predicate; public abstract class ArgumentBuilder> { private final RootCommandNode arguments = new RootCommandNode<>(); private Command command; + private Predicate requirement = s -> true; protected abstract T getThis(); @@ -30,5 +32,14 @@ public abstract class ArgumentBuilder> { return command; } + public T requires(Predicate requirement) { + this.requirement = requirement; + return getThis(); + } + + public Predicate getRequirement() { + return requirement; + } + public abstract CommandNode build(); } diff --git a/src/main/java/com/mojang/brigadier/builder/LiteralArgumentBuilder.java b/src/main/java/com/mojang/brigadier/builder/LiteralArgumentBuilder.java index f540cab..75c292b 100644 --- a/src/main/java/com/mojang/brigadier/builder/LiteralArgumentBuilder.java +++ b/src/main/java/com/mojang/brigadier/builder/LiteralArgumentBuilder.java @@ -25,7 +25,7 @@ public class LiteralArgumentBuilder extends ArgumentBuilder build() { - LiteralCommandNode result = new LiteralCommandNode<>(getLiteral(), getCommand()); + LiteralCommandNode result = new LiteralCommandNode<>(getLiteral(), getCommand(), getRequirement()); for (CommandNode argument : getArguments()) { result.addChild(argument); diff --git a/src/main/java/com/mojang/brigadier/builder/RequiredArgumentBuilder.java b/src/main/java/com/mojang/brigadier/builder/RequiredArgumentBuilder.java index 7aad5af..b5615ca 100644 --- a/src/main/java/com/mojang/brigadier/builder/RequiredArgumentBuilder.java +++ b/src/main/java/com/mojang/brigadier/builder/RequiredArgumentBuilder.java @@ -31,7 +31,7 @@ public class RequiredArgumentBuilder extends ArgumentBuilder build() { - ArgumentCommandNode result = new ArgumentCommandNode<>(getName(), getType(), getCommand()); + ArgumentCommandNode result = new ArgumentCommandNode<>(getName(), getType(), getCommand(), getRequirement()); for (CommandNode argument : getArguments()) { result.addChild(argument); diff --git a/src/main/java/com/mojang/brigadier/context/CommandContextBuilder.java b/src/main/java/com/mojang/brigadier/context/CommandContextBuilder.java index 7b51349..862327f 100644 --- a/src/main/java/com/mojang/brigadier/context/CommandContextBuilder.java +++ b/src/main/java/com/mojang/brigadier/context/CommandContextBuilder.java @@ -16,6 +16,10 @@ public class CommandContextBuilder { this.source = source; } + public S getSource() { + return source; + } + public CommandContextBuilder withArgument(String name, ParsedArgument argument) { this.arguments.put(name, argument); return this; diff --git a/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java b/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java index dd1c400..3df7b65 100644 --- a/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java +++ b/src/main/java/com/mojang/brigadier/tree/ArgumentCommandNode.java @@ -7,6 +7,7 @@ import com.mojang.brigadier.context.ParsedArgument; import com.mojang.brigadier.exceptions.CommandException; import java.util.Set; +import java.util.function.Predicate; public class ArgumentCommandNode extends CommandNode { private static final String USAGE_ARGUMENT_OPEN = "<"; @@ -15,8 +16,8 @@ public class ArgumentCommandNode extends CommandNode { private final String name; private final CommandArgumentType type; - public ArgumentCommandNode(String name, CommandArgumentType type, Command command) { - super(command); + public ArgumentCommandNode(String name, CommandArgumentType type, Command command, Predicate requirement) { + super(command, requirement); this.name = name; this.type = type; } diff --git a/src/main/java/com/mojang/brigadier/tree/CommandNode.java b/src/main/java/com/mojang/brigadier/tree/CommandNode.java index 30474ac..b81ba4c 100644 --- a/src/main/java/com/mojang/brigadier/tree/CommandNode.java +++ b/src/main/java/com/mojang/brigadier/tree/CommandNode.java @@ -8,13 +8,16 @@ import com.mojang.brigadier.exceptions.CommandException; import java.util.Collection; import java.util.Map; import java.util.Set; +import java.util.function.Predicate; public abstract class CommandNode { private final Map> children = Maps.newLinkedHashMap(); private Command command; + private Predicate requirement; - protected CommandNode(Command command) { + protected CommandNode(Command command, Predicate requirement) { this.command = command; + this.requirement = requirement; } public Command getCommand() { @@ -25,6 +28,10 @@ public abstract class CommandNode { return children.values(); } + public boolean canUse(S source) { + return requirement.test(source); + } + public void addChild(CommandNode node) { CommandNode child = children.get(node.getMergeKey()); if (child != null) { diff --git a/src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java b/src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java index 51ed4a6..a12f81d 100644 --- a/src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java +++ b/src/main/java/com/mojang/brigadier/tree/LiteralCommandNode.java @@ -7,14 +7,15 @@ import com.mojang.brigadier.exceptions.CommandException; import com.mojang.brigadier.exceptions.ParameterizedCommandExceptionType; import java.util.Set; +import java.util.function.Predicate; public class LiteralCommandNode extends CommandNode { public static final ParameterizedCommandExceptionType ERROR_INCORRECT_LITERAL = new ParameterizedCommandExceptionType("argument.literal.incorrect", "Expected literal ${expected}", "expected"); private final String literal; - public LiteralCommandNode(String literal, Command command) { - super(command); + public LiteralCommandNode(String literal, Command command, Predicate requirement) { + super(command, requirement); this.literal = literal; } diff --git a/src/main/java/com/mojang/brigadier/tree/RootCommandNode.java b/src/main/java/com/mojang/brigadier/tree/RootCommandNode.java index 1596715..16677e3 100644 --- a/src/main/java/com/mojang/brigadier/tree/RootCommandNode.java +++ b/src/main/java/com/mojang/brigadier/tree/RootCommandNode.java @@ -7,7 +7,7 @@ import java.util.Set; public class RootCommandNode extends CommandNode { public RootCommandNode() { - super(null); + super(null, c -> true); } @Override diff --git a/src/test/java/com/mojang/brigadier/CommandDispatcherCompletionsTest.java b/src/test/java/com/mojang/brigadier/CommandDispatcherCompletionsTest.java index 7956497..e6e2dc1 100644 --- a/src/test/java/com/mojang/brigadier/CommandDispatcherCompletionsTest.java +++ b/src/test/java/com/mojang/brigadier/CommandDispatcherCompletionsTest.java @@ -35,6 +35,7 @@ public class CommandDispatcherCompletionsTest { public void testCommand() throws Exception { subject.register(literal("foo")); subject.register(literal("bar")); + subject.register(literal("baz").requires(s -> false)); assertThat(subject.getCompletionSuggestions("", source), equalTo(new String[] {"foo", "bar"})); assertThat(subject.getCompletionSuggestions("f", source), equalTo(new String[] {"foo"})); assertThat(subject.getCompletionSuggestions("b", source), equalTo(new String[] {"bar"})); @@ -43,7 +44,7 @@ public class CommandDispatcherCompletionsTest { @Test public void testSubCommand() throws Exception { - subject.register(literal("foo").then(literal("abc")).then(literal("def"))); + subject.register(literal("foo").then(literal("abc")).then(literal("def")).then(literal("ghi").requires(s -> false))); subject.register(literal("bar")); assertThat(subject.getCompletionSuggestions("", source), equalTo(new String[] {"foo", "bar"})); assertThat(subject.getCompletionSuggestions("f", source), equalTo(new String[] {"foo"})); diff --git a/src/test/java/com/mojang/brigadier/CommandDispatcherTest.java b/src/test/java/com/mojang/brigadier/CommandDispatcherTest.java index 131cf34..6f3295d 100644 --- a/src/test/java/com/mojang/brigadier/CommandDispatcherTest.java +++ b/src/test/java/com/mojang/brigadier/CommandDispatcherTest.java @@ -99,6 +99,19 @@ public class CommandDispatcherTest { } } + @Test + public void testExecuteImpermissibleCommand() throws Exception { + subject.register(literal("foo").requires(s -> false)); + + try { + subject.execute("foo", source); + fail(); + } catch (CommandException ex) { + assertThat(ex.getType(), is(CommandDispatcher.ERROR_IMPERMISSIBLE)); + assertThat(ex.getData(), is(Collections.emptyMap())); + } + } + @Test public void testExecuteUnknownSubcommand() throws Exception { subject.register(literal("foo").executes(command)); diff --git a/src/test/java/com/mojang/brigadier/CommandDispatcherUsagesTest.java b/src/test/java/com/mojang/brigadier/CommandDispatcherUsagesTest.java index a05446e..bbe16d0 100644 --- a/src/test/java/com/mojang/brigadier/CommandDispatcherUsagesTest.java +++ b/src/test/java/com/mojang/brigadier/CommandDispatcherUsagesTest.java @@ -42,6 +42,19 @@ public class CommandDispatcherUsagesTest { } } + @Test + public void testInaccessibleCommand() throws Exception { + subject.register(literal("foo").requires(s -> false)); + + try { + subject.getUsage("foo", source); + fail(); + } catch (CommandException ex) { + assertThat(ex.getType(), is(CommandDispatcher.ERROR_IMPERMISSIBLE)); + assertThat(ex.getData(), is(Collections.emptyMap())); + } + } + @Test public void testSubcommandUsage() throws Exception { subject.register(