Implement query rules in file based system access control

presto-docs/src/main/sphinx/security/built-in-system-access-control.rst

@@ -71,10 +71,11 @@ contents:
 The config file is specified in JSON format.
 
 * It contains the rules defining which catalog can be accessed by which user (see Catalog Rules below).
+* The query rules specifying which queries can be managed by which user (see Query Rules below).
 * The impersonation rules specify which user impersonations are allowed (see Impersonation Rules below).
 * The principal rules specifying what principals can identify as what users (see Principal Rules below).
 
-This plugin currently only supports catalog access control rules, impersonation and principal
+This plugin currently supports catalog access, query, impersonation. and principal
 rules. If you want to limit access on a system level in any other way, you
 must implement a custom SystemAccessControl plugin
 (see :doc:`/develop/system-access-control`).
@@ -144,6 +145,32 @@ and deny all other access, you can use the following rules:
       ]
     }
 
+.. _query_rules:
+
+Query Rules
+-----------
+
+These rules control the ability of a user to execute, view, or kill a query. The user is
+granted or denied access, based on the first matching rule read from top to bottom. If no
+rules are specified, all users are allowed to execute queries, and to view or kill queries
+owned by any user. If no rule matches, query management is denied. Each rule is composed
+of the following fields:
+
+* ``user`` (optional): regex to match against user name. Defaults to ``.*``.
+* ``owner`` (optional): regex to match against the query owner name. Defaults to ``.*``.
+* ``allow`` (required): set of query permissions granted to user. Values: ``execute``, ``view``, ``kill``
+
+.. note::
+
+    Users always have permission to view or kill their own queries.
+
+For example, if you want to allow the user ``admin`` full query access, allow the user ``alice``
+to execute and kill queries, any user to execute queries, and deny all other access, you can use
+the following rules:
+
+.. literalinclude:: query-access.json
+    :language: json
+
 .. _impersonation_rules:
 
 Impersonation Rules

presto-docs/src/main/sphinx/security/query-access.json

@@ -0,0 +1,20 @@
+{
+  "catalogs": [
+    {
+      "allow": true
+    }
+  ],
+  "queries": [
+    {
+      "user": "admin",
+      "allow": ["execute", "kill", "view"]
+    },
+    {
+      "user": "alice",
+      "allow": ["execute", "kill"]
+    },
+    {
+      "allow": ["execute"]
+    }
+  ]
+}

presto-plugin-toolkit/src/main/java/io/prestosql/plugin/base/security/FileBasedSystemAccessControl.java

@@ -40,6 +40,7 @@
 
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Suppliers.memoizeWithExpiration;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static io.prestosql.plugin.base.security.CatalogAccessControlRule.AccessMode.ALL;
 import static io.prestosql.plugin.base.security.CatalogAccessControlRule.AccessMode.READ_ONLY;
 import static io.prestosql.plugin.base.security.FileBasedAccessControlConfig.SECURITY_CONFIG_FILE;
@@ -67,6 +68,7 @@
 import static io.prestosql.spi.security.AccessDeniedException.denyRenameView;
 import static io.prestosql.spi.security.AccessDeniedException.denyRevokeTablePrivilege;
 import static io.prestosql.spi.security.AccessDeniedException.denySetUser;
+import static io.prestosql.spi.security.AccessDeniedException.denyViewQuery;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
@@ -79,15 +81,18 @@ public class FileBasedSystemAccessControl
     public static final String NAME = "file";
 
     private final List<CatalogAccessControlRule> catalogRules;
+    private final Optional<List<QueryAccessRule>> queryAccessRules;
     private final Optional<List<ImpersonationRule>> impersonationRules;
     private final Optional<List<PrincipalUserMatchRule>> principalUserMatchRules;
 
     private FileBasedSystemAccessControl(
             List<CatalogAccessControlRule> catalogRules,
+            Optional<List<QueryAccessRule>> queryAccessRules,
             Optional<List<ImpersonationRule>> impersonationRules,
             Optional<List<PrincipalUserMatchRule>> principalUserMatchRules)
     {
         this.catalogRules = catalogRules;
+        this.queryAccessRules = queryAccessRules;
         this.impersonationRules = impersonationRules;
         this.principalUserMatchRules = principalUserMatchRules;
     }
@@ -152,7 +157,11 @@ private SystemAccessControl create(String configFileName)
                     Optional.of(Pattern.compile(".*")),
                     Optional.of(Pattern.compile("system"))));
 
-            return new FileBasedSystemAccessControl(catalogRulesBuilder.build(), rules.getImpersonationRules(), rules.getPrincipalUserMatchRules());
+            return new FileBasedSystemAccessControl(
+                    catalogRulesBuilder.build(),
+                    rules.getQueryAccessRules(),
+                    rules.getImpersonationRules(),
+                    rules.getPrincipalUserMatchRules());
         }
     }
 
@@ -213,22 +222,59 @@ public void checkCanSetUser(Optional<Principal> principal, String userName)
     @Override
     public void checkCanExecuteQuery(SystemSecurityContext context)
     {
+        if (!queryAccessRules.isPresent()) {
+            return;
+        }
+        if (!canAccessQuery(context.getIdentity(), QueryAccessRule.AccessMode.EXECUTE)) {
+            denyViewQuery();
+        }
     }
 
     @Override
     public void checkCanViewQueryOwnedBy(SystemSecurityContext context, String queryOwner)
     {
+        if (!queryAccessRules.isPresent()) {
+            return;
+        }
+        if (!canAccessQuery(context.getIdentity(), QueryAccessRule.AccessMode.VIEW)) {
+            denyViewQuery();
+        }
     }
 
     @Override
     public Set<String> filterViewQueryOwnedBy(SystemSecurityContext context, Set<String> queryOwners)
     {
-        return queryOwners;
+        if (!queryAccessRules.isPresent()) {
+            return queryOwners;
+        }
+        Identity identity = context.getIdentity();
+        return queryOwners.stream()
+                .filter(owner -> canAccessQuery(identity, QueryAccessRule.AccessMode.VIEW))
+                .collect(toImmutableSet());
     }
 
     @Override
     public void checkCanKillQueryOwnedBy(SystemSecurityContext context, String queryOwner)
     {
+        if (!queryAccessRules.isPresent()) {
+            return;
+        }
+        if (!canAccessQuery(context.getIdentity(), QueryAccessRule.AccessMode.KILL)) {
+            denyViewQuery();
+        }
+    }
+
+    private boolean canAccessQuery(Identity identity, QueryAccessRule.AccessMode requiredAccess)
+    {
+        if (queryAccessRules.isPresent()) {
+            for (QueryAccessRule rule : queryAccessRules.get()) {
+                Optional<Set<QueryAccessRule.AccessMode>> accessMode = rule.match(identity.getUser());
+                if (accessMode.isPresent()) {
+                    return accessMode.get().contains(requiredAccess);
+                }
+            }
+        }
+        return false;
     }
 
     @Override

presto-plugin-toolkit/src/main/java/io/prestosql/plugin/base/security/FileBasedSystemAccessControlRules.java

@@ -23,16 +23,19 @@
 public class FileBasedSystemAccessControlRules
 {
     private final List<CatalogAccessControlRule> catalogRules;
+    private final Optional<List<QueryAccessRule>> queryAccessRules;
     private final Optional<List<ImpersonationRule>> impersonationRules;
     private final Optional<List<PrincipalUserMatchRule>> principalUserMatchRules;
 
     @JsonCreator
     public FileBasedSystemAccessControlRules(
             @JsonProperty("catalogs") Optional<List<CatalogAccessControlRule>> catalogRules,
+            @JsonProperty("queries") Optional<List<QueryAccessRule>> queryAccessRules,
             @JsonProperty("impersonation") Optional<List<ImpersonationRule>> impersonationRules,
             @JsonProperty("principals") Optional<List<PrincipalUserMatchRule>> principalUserMatchRules)
     {
         this.catalogRules = catalogRules.map(ImmutableList::copyOf).orElse(ImmutableList.of());
+        this.queryAccessRules = queryAccessRules.map(ImmutableList::copyOf);
         this.principalUserMatchRules = principalUserMatchRules.map(ImmutableList::copyOf);
         this.impersonationRules = impersonationRules.map(ImmutableList::copyOf);
     }
@@ -42,6 +45,11 @@ public List<CatalogAccessControlRule> getCatalogRules()
         return catalogRules;
     }
 
+    public Optional<List<QueryAccessRule>> getQueryAccessRules()
+    {
+        return queryAccessRules;
+    }
+
     public Optional<List<ImpersonationRule>> getImpersonationRules()
     {
         return impersonationRules;

presto-plugin-toolkit/src/main/java/io/prestosql/plugin/base/security/QueryAccessRule.java

@@ -0,0 +1,100 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.prestosql.plugin.base.security;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonValue;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.Locale;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.regex.Pattern;
+
+import static com.google.common.base.MoreObjects.toStringHelper;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static java.util.Arrays.stream;
+import static java.util.Objects.requireNonNull;
+import static java.util.function.Function.identity;
+
+public class QueryAccessRule
+{
+    private final Set<AccessMode> allow;
+    private final Optional<Pattern> userRegex;
+
+    @JsonCreator
+    public QueryAccessRule(
+            @JsonProperty("allow") Set<AccessMode> allow,
+            @JsonProperty("user") Optional<Pattern> userRegex)
+    {
+        this.allow = ImmutableSet.copyOf(requireNonNull(allow, "allow is null"));
+        this.userRegex = requireNonNull(userRegex, "userRegex is null");
+    }
+
+    public Optional<Set<AccessMode>> match(String user)
+    {
+        if (userRegex.map(regex -> regex.matcher(user).matches()).orElse(true)) {
+            return Optional.of(allow);
+        }
+        return Optional.empty();
+    }
+
+    @Override
+    public String toString()
+    {
+        return toStringHelper(this)
+                .omitNullValues()
+                .add("allow", allow)
+                .add("userRegex", userRegex.orElse(null))
+                .toString();
+    }
+
+    public enum AccessMode
+    {
+        EXECUTE("execute"),
+        VIEW("view"),
+        KILL("kill");
+
+        private static final Map<String, AccessMode> modeByName = stream(AccessMode.values()).collect(toImmutableMap(AccessMode::toString, identity()));
+
+        private final String stringValue;
+
+        AccessMode(String stringValue)
+        {
+            this.stringValue = requireNonNull(stringValue, "stringValue is null");
+        }
+
+        @JsonValue
+        @Override
+        public String toString()
+        {
+            return stringValue;
+        }
+
+        @JsonCreator
+        public static AccessMode fromJson(Object value)
+        {
+            if (value instanceof String) {
+                AccessMode accessMode = modeByName.get(((String) value).toLowerCase(Locale.US));
+                if (accessMode != null) {
+                    return accessMode;
+                }
+            }
+
+            throw new IllegalArgumentException("Unknown " + AccessMode.class.getSimpleName() + ": " + value);
+        }
+    }
+}

presto-plugin-toolkit/src/test/java/io/prestosql/plugin/base/security/TestFileBasedSystemAccessControl.java

@@ -36,6 +36,7 @@
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.assertj.core.util.Files.newTemporaryFile;
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
 
 public class TestFileBasedSystemAccessControl
 {
@@ -96,6 +97,65 @@ public void testCanSetUserOperations()
         accessControlNoPatterns.checkCanSetUser(kerberosValidAlice.getPrincipal(), kerberosValidAlice.getUser());
     }
 
+    @Test
+    public void testQuery()
+    {
+        SystemAccessControl accessControlManager = newFileBasedSystemAccessControl("query.json");
+
+        accessControlManager.checkCanExecuteQuery(new SystemSecurityContext(admin));
+        accessControlManager.checkCanViewQueryOwnedBy(new SystemSecurityContext(admin), "any");
+        assertEquals(accessControlManager.filterViewQueryOwnedBy(new SystemSecurityContext(admin), ImmutableSet.of("a", "b")), ImmutableSet.of("a", "b"));
+        accessControlManager.checkCanKillQueryOwnedBy(new SystemSecurityContext(admin), "any");
+
+        accessControlManager.checkCanExecuteQuery(new SystemSecurityContext(alice));
+        accessControlManager.checkCanViewQueryOwnedBy(new SystemSecurityContext(alice), "any");
+        assertEquals(accessControlManager.filterViewQueryOwnedBy(new SystemSecurityContext(alice), ImmutableSet.of("a", "b")), ImmutableSet.of("a", "b"));
+        assertThrows(AccessDeniedException.class, () -> accessControlManager.checkCanKillQueryOwnedBy(new SystemSecurityContext(alice), "any"));
+
+        assertThrows(AccessDeniedException.class, () -> accessControlManager.checkCanExecuteQuery(new SystemSecurityContext(bob)));
+        assertThrows(AccessDeniedException.class, () -> accessControlManager.checkCanViewQueryOwnedBy(new SystemSecurityContext(bob), "any"));
+        assertEquals(accessControlManager.filterViewQueryOwnedBy(new SystemSecurityContext(bob), ImmutableSet.of("a", "b")), ImmutableSet.of());
+        accessControlManager.checkCanKillQueryOwnedBy(new SystemSecurityContext(bob), "any");
+
+        accessControlManager.checkCanExecuteQuery(new SystemSecurityContext(nonAsciiUser));
+        accessControlManager.checkCanViewQueryOwnedBy(new SystemSecurityContext(nonAsciiUser), "any");
+        assertEquals(accessControlManager.filterViewQueryOwnedBy(new SystemSecurityContext(nonAsciiUser), ImmutableSet.of("a", "b")), ImmutableSet.of("a", "b"));
+        accessControlManager.checkCanKillQueryOwnedBy(new SystemSecurityContext(nonAsciiUser), "any");
+    }
+
+    @Test
+    public void testQueryNotSet()
+    {
+        SystemAccessControl accessControlManager = newFileBasedSystemAccessControl("catalog.json");
+
+        accessControlManager.checkCanExecuteQuery(new SystemSecurityContext(bob));
+        accessControlManager.checkCanViewQueryOwnedBy(new SystemSecurityContext(bob), "any");
+        assertEquals(accessControlManager.filterViewQueryOwnedBy(new SystemSecurityContext(bob), ImmutableSet.of("a", "b")), ImmutableSet.of("a", "b"));
+        accessControlManager.checkCanKillQueryOwnedBy(new SystemSecurityContext(bob), "any");
+    }
+
+    @Test
+    public void testDocsExample()
+    {
+        String rulesFile = new File("../presto-docs/src/main/sphinx/security/query-access.json").getAbsolutePath();
+        SystemAccessControl accessControlManager = newFileBasedSystemAccessControl(ImmutableMap.of("security.config-file", rulesFile));
+
+        accessControlManager.checkCanExecuteQuery(new SystemSecurityContext(admin));
+        accessControlManager.checkCanViewQueryOwnedBy(new SystemSecurityContext(admin), "any");
+        assertEquals(accessControlManager.filterViewQueryOwnedBy(new SystemSecurityContext(admin), ImmutableSet.of("a", "b")), ImmutableSet.of("a", "b"));
+        accessControlManager.checkCanKillQueryOwnedBy(new SystemSecurityContext(admin), "any");
+
+        accessControlManager.checkCanExecuteQuery(new SystemSecurityContext(alice));
+        assertThrows(AccessDeniedException.class, () -> accessControlManager.checkCanViewQueryOwnedBy(new SystemSecurityContext(alice), "any"));
+        assertEquals(accessControlManager.filterViewQueryOwnedBy(new SystemSecurityContext(alice), ImmutableSet.of("a", "b")), ImmutableSet.of());
+        accessControlManager.checkCanKillQueryOwnedBy(new SystemSecurityContext(alice), "any");
+
+        accessControlManager.checkCanExecuteQuery(new SystemSecurityContext(bob));
+        assertThrows(AccessDeniedException.class, () -> accessControlManager.checkCanViewQueryOwnedBy(new SystemSecurityContext(bob), "any"));
+        assertEquals(accessControlManager.filterViewQueryOwnedBy(new SystemSecurityContext(bob), ImmutableSet.of("a", "b")), ImmutableSet.of());
+        assertThrows(AccessDeniedException.class, () -> accessControlManager.checkCanKillQueryOwnedBy(new SystemSecurityContext(bob), "any"));
+    }
+
     @Test
     public void testCatalogOperations()
     {