access checks galore!

build.gradle.kts

@@ -5,7 +5,7 @@ plugins {
 
 allprojects {
     group = "pw.dotdash"
-    version = "0.14.0"
+    version = "0.15.0"
 
     apply(plugin = "maven")
 }

director-core/src/main/kotlin/pw/dotdash/director/core/exception/CommandPermissionException.kt

@@ -0,0 +1,3 @@
+package pw.dotdash.director.core.exception
+
+class CommandPermissionException(message: String = "You do not have permission to use this command.") : CommandException(message, true)

director-core/src/main/kotlin/pw/dotdash/director/core/tree/ArgumentCommandTree.kt

@@ -25,6 +25,8 @@ interface ArgumentCommandTree<in S, in P : HList<@UnsafeVariance P>, V, out R> :
 
         override fun setExecutor(executor: (S, HCons<V, P>) -> R): Builder<S, P, V, R>
 
+        override fun setAccessibility(test: (S, HCons<V, P>) -> Boolean): Builder<S, P, V, R>
+
         fun build(): ArgumentCommandTree<S, P, V, R>
     }
 

director-core/src/main/kotlin/pw/dotdash/director/core/tree/ChildCommandTree.kt

@@ -26,6 +26,8 @@ interface ChildCommandTree<in S, in P : HList<@UnsafeVariance P>, out R> : Comma
 
         override fun setExecutor(executor: (S, P) -> R): Builder<S, P, R>
 
+        override fun setAccessibility(test: (S, P) -> Boolean): Builder<S, P, R>
+
         fun build(): ChildCommandTree<S, P, R>
     }
 

director-core/src/main/kotlin/pw/dotdash/director/core/tree/CommandTree.kt

@@ -13,6 +13,8 @@ interface CommandTree<in S, in V : HList<@UnsafeVariance V>, out R> : TreeExecut
 
     fun getUsage(source: S): String
 
+    fun canAccess(source: S, previous: V): Boolean
+
     interface Builder<S, V : HList<V>, R> {
 
         fun addChild(child: ChildCommandTree<S, V, R>): Builder<S, V, R>
@@ -26,5 +28,7 @@ interface CommandTree<in S, in V : HList<@UnsafeVariance V>, out R> : TreeExecut
         fun <NV> setArgument(parameter: Parameter<S, V, NV>, init: ArgumentCommandTree.Builder<S, V, NV, R>.() -> Unit): Builder<S, V, R>
 
         fun setExecutor(executor: (S, V) -> R): Builder<S, V, R>
+
+        fun setAccessibility(test: (S, V) -> Boolean): Builder<S, V, R>
     }
 }

director-core/src/main/kotlin/pw/dotdash/director/core/tree/RootCommandTree.kt

@@ -44,6 +44,8 @@ interface RootCommandTree<in S, V : HList<V>, out R> : CommandTree<S, V, R> {
 
         override fun setExecutor(executor: (S, V) -> R): Builder<S, V, R>
 
+        override fun setAccessibility(test: (S, V) -> Boolean): Builder<S, V, R>
+
         fun build(): RootCommandTree<S, V, R>
     }
 

director-core/src/main/kotlin/pw/dotdash/director/core/tree/simple/SimpleCommandTree.kt

@@ -11,14 +11,17 @@ import pw.dotdash.director.core.lexer.InputTokenizer
 import pw.dotdash.director.core.lexer.QuotedInputTokenizer
 import pw.dotdash.director.core.tree.*
 import pw.dotdash.director.core.util.StartsWithPredicate
-import java.util.function.Consumer
 
 internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
     final override val children: Map<String, SimpleChildCommandTree<S, V, R>>,
     final override val argument: SimpleArgumentCommandTree<S, V, Any?, R>?,
-    final override val executor: ((S, V) -> R)?
+    final override val executor: ((S, V) -> R)?,
+    private val accessibility: ((S, V) -> Boolean)?
 ) : CommandTree<S, V, R> {
 
+    override fun canAccess(source: S, previous: V): Boolean =
+        this.accessibility == null || this.accessibility.invoke(source, previous)
+
     override fun execute(source: S, tokens: CommandTokens, previous: V): R {
         return execute(source, tokens, previous, ArrayList())
     }
@@ -36,11 +39,16 @@ internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
             if (this.argument != null) {
                 try {
                     val parsed: Any? = this.argument.parameter.value.parse(source, tokens, previous)
-                    usageParts += this.argument.parameter.getUsage(source)
+                    val next: HCons<Any?, V> = HCons(parsed, previous)
+
+                    if (this.argument.canAccess(source, next)) {
+                        usageParts += this.argument.parameter.getUsage(source)
+                        return this.argument.execute(source, tokens, next, usageParts)
+                    }
 
-                    return this.argument.execute(source, tokens, HCons(parsed, previous), usageParts)
+                    // Otherwise ignore access errors here.
                 } catch (ignored: ArgumentParseException) {
-                    // Ignore parsing errors since we don't have any tokens left.
+                    // Ignore parsing errors since we don't have any tokens left anyways.
                     tokens.snapshot = snapshot
                 } catch (e: CommandException) {
                     throw e.wrap(usageParts)
@@ -66,8 +74,11 @@ internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
         val child: SimpleChildCommandTree<S, V, R>? = this.children[alias]
 
         if (child != null) {
-            usageParts += alias
+            if (!child.canAccess(source, previous)) {
+                throw tokens.createError("You do not have access to this subcommand.").wrap(usageParts)
+            }
 
+            usageParts += alias
             return child.execute(source, tokens, previous, usageParts)
         }
 
@@ -76,15 +87,20 @@ internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
         if (this.argument != null) {
             try {
                 val parsed: Any? = this.argument.parameter.value.parse(source, tokens, previous)
-                usageParts += this.argument.parameter.getUsage(source)
+                val next: HCons<Any?, V> = HCons(parsed, previous)
+
+                if (!this.argument.canAccess(source, next)) {
+                    throw tokens.createError("You do not have access to this argument.").wrap(usageParts)
+                }
 
-                return this.argument.execute(source, tokens, HCons(parsed, previous), usageParts)
+                usageParts += this.argument.parameter.getUsage(source)
+                return this.argument.execute(source, tokens, next, usageParts)
             } catch (e: CommandException) {
                 throw e.wrap(usageParts)
             }
         }
 
-        throw tokens.createError("Invalid subcommand.").wrap(usageParts)
+        throw tokens.createError("Too many arguments!").wrap(usageParts)
     }
 
     protected fun complete(source: S, tokens: CommandTokens, previous: V, usageParts: MutableList<String>): List<String> {
@@ -94,11 +110,16 @@ internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
             if (this.argument != null) {
                 try {
                     val parsed: Any? = this.argument.parameter.value.parse(source, tokens, previous)
-                    usageParts += this.argument.parameter.getUsage(source)
+                    val next: HCons<Any?, V> = HCons(parsed, previous)
+
+                    if (this.argument.canAccess(source, next)) {
+                        usageParts += this.argument.parameter.getUsage(source)
+                        return this.argument.complete(source, tokens, next, usageParts)
+                    }
 
-                    return this.argument.complete(source, tokens, HCons(parsed, previous), usageParts)
+                    // Otherwise ignore access errors here.
                 } catch (ignored: ArgumentParseException) {
-                    // Ignore parsing errors since we don't have any tokens left.
+                    // Ignore parsing errors since we don't have any tokens left anyways.
                     tokens.snapshot = snapshot
                 } catch (e: CommandException) {
                     throw e.wrap(usageParts)
@@ -112,11 +133,12 @@ internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
         val child: SimpleChildCommandTree<S, V, R>? = this.children[alias]
 
         if (tokens.hasNext()) {
-            if (child != null) {
+            if (child != null && child.canAccess(source, previous)) {
                 usageParts += alias
-
                 return child.complete(source, tokens, previous, usageParts)
             }
+
+            // Otherwise ignore access errors since we're completing.
         } else {
             return this.subCompletions(source, tokens, previous, usageParts).filter(StartsWithPredicate(alias))
         }
@@ -127,10 +149,15 @@ internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
         if (this.argument != null) {
             try {
                 val parsed: Any? = this.argument.parameter.value.parse(source, tokens, previous)
-                usageParts += this.argument.parameter.getUsage(source)
+                val next: HCons<Any?, V> = HCons(parsed, previous)
+
+                if (this.argument.canAccess(source, next)) {
+                    // Argument successfully parsed; complete its subtree.
+                    usageParts += this.argument.parameter.getUsage(source)
+                    return this.argument.complete(source, tokens, next, usageParts)
+                }
 
-                // Argument successfully parsed; complete its subtree.
-                return this.argument.complete(source, tokens, HCons(parsed, previous), usageParts)
+                // Otherwise ignore access errors since we're completing.
             } catch (e: CommandException) {
                 // Failed to parse argument; rollback.
                 tokens.snapshot = snapshot
@@ -141,9 +168,13 @@ internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
     }
 
     private fun subCompletions(source: S, tokens: CommandTokens, previous: V, usageParts: MutableList<String>): List<String> {
-        val result = ArrayList<String>()
+        val result = HashSet<String>()
 
-        result += this.children.keys
+        for (child: SimpleChildCommandTree<S, V, R> in this.children.values) {
+            if (child.canAccess(source, previous)) {
+                result += child.aliases
+            }
+        }
 
         if (this.argument != null) {
             try {
@@ -153,7 +184,7 @@ internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
             }
         }
 
-        return result
+        return result.toList()
     }
 
     private fun CommandException.wrap(usageParts: List<String>): TreeCommandException =
@@ -182,6 +213,7 @@ internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
         protected val children = HashMap<String, SimpleChildCommandTree<S, V, R>>()
         protected var argument: SimpleArgumentCommandTree<S, V, in Any?, R>? = null
         protected var executor: ((S, V) -> R)? = null
+        protected var accessibility: ((S, V) -> Boolean)? = null
 
         override fun addChild(child: ChildCommandTree<S, V, R>): B {
             require(child is SimpleChildCommandTree) { "Child trees must be made with ChildCommandTree.builder()" }
@@ -213,6 +245,11 @@ internal sealed class SimpleCommandTree<S, V : HList<V>, R>(
             this.executor = executor
             return this as B
         }
+
+        override fun setAccessibility(test: (S, V) -> Boolean): B {
+            this.accessibility = test
+            return this as B
+        }
     }
 }
 
@@ -224,8 +261,9 @@ internal class SimpleRootCommandTree<S, V : HList<V>, R>(
     override val extendedDescription: String?,
     children: Map<String, SimpleChildCommandTree<S, V, R>>,
     argument: SimpleArgumentCommandTree<S, V, in Any?, R>?,
-    executor: ((S, V) -> R)?
-) : SimpleCommandTree<S, V, R>(children, argument, executor), RootCommandTree<S, V, R> {
+    executor: ((S, V) -> R)?,
+    accessibility: ((S, V) -> Boolean)?
+) : SimpleCommandTree<S, V, R>(children, argument, executor, accessibility), RootCommandTree<S, V, R> {
 
     class Builder<S, V : HList<V>, R> :
         SimpleCommandTree.Builder<Builder<S, V, R>, S, V, R>(),
@@ -275,7 +313,8 @@ internal class SimpleRootCommandTree<S, V : HList<V>, R>(
             extendedDescription = this.extendedDescription,
             children = this.children,
             argument = this.argument,
-            executor = this.executor
+            executor = this.executor,
+            accessibility = this.accessibility
         )
     }
 }
@@ -284,8 +323,9 @@ internal class SimpleChildCommandTree<S, P : HList<P>, R>(
     override val aliases: List<String>,
     children: Map<String, SimpleChildCommandTree<S, P, R>>,
     argument: SimpleArgumentCommandTree<S, P, in Any?, R>?,
-    executor: ((S, P) -> R)?
-) : SimpleCommandTree<S, P, R>(children, argument, executor), ChildCommandTree<S, P, R> {
+    executor: ((S, P) -> R)?,
+    accessibility: ((S, P) -> Boolean)?
+) : SimpleCommandTree<S, P, R>(children, argument, executor, accessibility), ChildCommandTree<S, P, R> {
 
     class Builder<S, P : HList<P>, R> :
         SimpleCommandTree.Builder<Builder<S, P, R>, S, P, R>(),
@@ -307,7 +347,8 @@ internal class SimpleChildCommandTree<S, P : HList<P>, R>(
             aliases = checkNotNull(this.aliases),
             children = this.children,
             argument = this.argument,
-            executor = this.executor
+            executor = this.executor,
+            accessibility = this.accessibility
         )
     }
 }
@@ -316,8 +357,9 @@ internal class SimpleArgumentCommandTree<S, P : HList<P>, V, R>(
     override val parameter: Parameter<S, P, V>,
     children: Map<String, SimpleChildCommandTree<S, HCons<V, P>, R>>,
     argument: SimpleArgumentCommandTree<S, HCons<V, P>, in Any?, R>?,
-    executor: ((S, HCons<V, P>) -> R)?
-) : SimpleCommandTree<S, HCons<V, P>, R>(children, argument, executor), ArgumentCommandTree<S, P, V, R> {
+    executor: ((S, HCons<V, P>) -> R)?,
+    accessibility: ((S, HCons<V, P>) -> Boolean)?
+) : SimpleCommandTree<S, HCons<V, P>, R>(children, argument, executor, accessibility), ArgumentCommandTree<S, P, V, R> {
 
     class Builder<S, P : HList<P>, V, R> :
         SimpleCommandTree.Builder<Builder<S, P, V, R>, S, HCons<V, P>, R>(),
@@ -334,7 +376,8 @@ internal class SimpleArgumentCommandTree<S, P : HList<P>, V, R>(
             parameter = checkNotNull(this.parameter),
             children = this.children,
             argument = this.argument,
-            executor = this.executor
+            executor = this.executor,
+            accessibility = this.accessibility
         )
     }
 }

director-sponge/src/main/kotlin/pw/dotdash/director/sponge/CommandTreeCallable.kt

@@ -1,9 +1,6 @@
 package pw.dotdash.director.sponge
 
-import org.spongepowered.api.command.CommandCallable
-import org.spongepowered.api.command.CommandException
-import org.spongepowered.api.command.CommandResult
-import org.spongepowered.api.command.CommandSource
+import org.spongepowered.api.command.*
 import org.spongepowered.api.text.Text
 import org.spongepowered.api.text.Text.NEW_LINE
 import org.spongepowered.api.text.format.TextColors.RED
@@ -21,10 +18,9 @@ import pw.dotdash.director.core.tree.RootCommandTree
 import java.util.*
 
 
-class CommandTreeCallable<T : HList<T>> @JvmOverloads constructor(
-    private val root: RootCommandTree<CommandSource, T, CommandResult>,
-    private val initial: T,
-    private val permission: String? = null
+class CommandTreeCallable<V : HList<V>>(
+    private val root: RootCommandTree<CommandSource, V, CommandResult>,
+    private val initial: V
 ) : CommandCallable {
 
     companion object {
@@ -36,14 +32,17 @@ class CommandTreeCallable<T : HList<T>> @JvmOverloads constructor(
 
         @JvmStatic
         @JvmName("of")
-        @JvmOverloads
-        operator fun invoke(root: RootCommandTree<CommandSource, HNil, CommandResult>, permission: String? = null): CommandTreeCallable<HNil> =
-            CommandTreeCallable(root, HNil, permission)
+        operator fun invoke(root: RootCommandTree<CommandSource, HNil, CommandResult>): CommandTreeCallable<HNil> =
+            CommandTreeCallable(root, HNil)
     }
 
     private val rootAlias: Text = Text.of(YELLOW, this.root.aliases.first())
 
     override fun process(source: CommandSource, arguments: String): CommandResult {
+        if (!this.testPermission(source)) {
+            throw CommandPermissionException()
+        }
+
         try {
             val tokens = SimpleCommandTokens(arguments, this.root.tokenizer.tokenize(arguments, false).toMutableList())
             return this.root.execute(source, tokens, this.initial)
@@ -53,6 +52,10 @@ class CommandTreeCallable<T : HList<T>> @JvmOverloads constructor(
     }
 
     override fun getSuggestions(source: CommandSource, arguments: String, targetPosition: Location<World>?): List<String> {
+        if (!this.testPermission(source)) {
+            return emptyList()
+        }
+
         try {
             val tokens = SimpleCommandTokens(arguments, this.root.tokenizer.tokenize(arguments, true).toMutableList())
             return this.root.complete(source, tokens, this.initial)
@@ -88,7 +91,7 @@ class CommandTreeCallable<T : HList<T>> @JvmOverloads constructor(
     }
 
     override fun testPermission(source: CommandSource): Boolean =
-        this.permission == null || source.hasPermission(this.permission)
+        this.root.canAccess(source, this.initial)
 
     override fun getShortDescription(source: CommandSource): Optional<Text> =
         Optional.ofNullable(this.root.description?.let(Text::of))

director-sponge/src/main/kotlin/pw/dotdash/director/sponge/SpongeCommandTree.kt

@@ -7,6 +7,7 @@ import pw.dotdash.director.core.HNil
 import pw.dotdash.director.core.Parameter
 import pw.dotdash.director.core.tree.ArgumentCommandTree
 import pw.dotdash.director.core.tree.ChildCommandTree
+import pw.dotdash.director.core.tree.CommandTree
 import pw.dotdash.director.core.tree.RootCommandTree
 
 object SpongeCommandTree {
@@ -46,4 +47,9 @@ object SpongeCommandTree {
     @JvmStatic
     fun <P : HList<P>, V> argument(parameter: Parameter<CommandSource, P, V>): ArgumentCommandTree.Builder<CommandSource, P, V, CommandResult> =
         ArgumentCommandTree.builder(parameter)
+}
+
+fun <B : CommandTree.Builder<S, *, *>, S : CommandSource> B.setPermission(permission: String): B {
+    this.setAccessibility { source, _ -> source.hasPermission(permission) }
+    return this
 }