Add query integrity check in AccessControlManager

presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManager.java

@@ -335,6 +335,7 @@ private <C> void createQueryInternal(QueryId queryId, SessionContext sessionCont
 
             // decode session
             session = sessionSupplier.createSession(queryId, sessionContext);
+            accessControl.checkQueryIntegrity(session.getIdentity(), query);
 
             WarningCollector warningCollector = warningCollectorFactory.create();
 

presto-main/src/main/java/com/facebook/presto/security/AccessControl.java

@@ -34,6 +34,12 @@ public interface AccessControl
      */
     void checkCanSetUser(Optional<Principal> principal, String userName);
 
+    /**
+     * Check if the query is unexpectedly modified using the credentials passed in the identity.
+     * @throws com.facebook.presto.spi.security.AccessDeniedException if query is modified.
+     */
+    void checkQueryIntegrity(Identity identity, String query);
+
     /**
      * Filter the list of catalogs to those visible to the identity.
      */

presto-main/src/main/java/com/facebook/presto/security/AccessControlManager.java

@@ -150,6 +150,15 @@ public void checkCanSetUser(Optional<Principal> principal, String userName)
         authenticationCheck(() -> systemAccessControl.get().checkCanSetUser(principal, userName));
     }
 
+    @Override
+    public void checkQueryIntegrity(Identity identity, String query)
+    {
+        requireNonNull(identity, "identity is null");
+        requireNonNull(query, "query is null");
+
+        authenticationCheck(() -> systemAccessControl.get().checkQueryIntegrity(identity, query));
+    }
+
     @Override
     public Set<String> filterCatalogs(Identity identity, Set<String> catalogs)
     {
@@ -753,6 +762,12 @@ public ConnectorTransactionHandle getTransactionHandle(TransactionId transaction
     private static class InitializingSystemAccessControl
             implements SystemAccessControl
     {
+        @Override
+        public void checkQueryIntegrity(Identity identity, String query)
+        {
+            throw new PrestoException(SERVER_STARTING_UP, "Presto server is still initializing");
+        }
+
         @Override
         public void checkCanSetUser(Optional<Principal> principal, String userName)
         {

presto-main/src/main/java/com/facebook/presto/security/AllowAllAccessControl.java

@@ -33,6 +33,11 @@ public void checkCanSetUser(Optional<Principal> principal, String userName)
     {
     }
 
+    @Override
+    public void checkQueryIntegrity(Identity identity, String query)
+    {
+    }
+
     @Override
     public Set<String> filterCatalogs(Identity identity, Set<String> catalogs)
     {

presto-main/src/main/java/com/facebook/presto/security/AllowAllSystemAccessControl.java

@@ -55,6 +55,11 @@ public SystemAccessControl create(Map<String, String> config)
         }
     }
 
+    @Override
+    public void checkQueryIntegrity(Identity identity, String query)
+    {
+    }
+
     @Override
     public void checkCanSetUser(Optional<Principal> principal, String userName)
     {

presto-main/src/main/java/com/facebook/presto/security/DenyAllAccessControl.java

@@ -42,6 +42,7 @@
 import static com.facebook.presto.spi.security.AccessDeniedException.denyGrantRoles;
 import static com.facebook.presto.spi.security.AccessDeniedException.denyGrantTablePrivilege;
 import static com.facebook.presto.spi.security.AccessDeniedException.denyInsertTable;
+import static com.facebook.presto.spi.security.AccessDeniedException.denyQueryIntegrityCheck;
 import static com.facebook.presto.spi.security.AccessDeniedException.denyRenameColumn;
 import static com.facebook.presto.spi.security.AccessDeniedException.denyRenameSchema;
 import static com.facebook.presto.spi.security.AccessDeniedException.denyRenameTable;
@@ -67,6 +68,12 @@ public void checkCanSetUser(Optional<Principal> principal, String userName)
         denySetUser(principal, userName);
     }
 
+    @Override
+    public void checkQueryIntegrity(Identity identity, String query)
+    {
+        denyQueryIntegrityCheck();
+    }
+
     @Override
     public Set<String> filterCatalogs(Identity identity, Set<String> catalogs)
     {

presto-main/src/main/java/com/facebook/presto/security/FileBasedSystemAccessControl.java

@@ -157,6 +157,11 @@ public void checkCanSetUser(Optional<Principal> principal, String userName)
         denySetUser(principal, userName);
     }
 
+    @Override
+    public void checkQueryIntegrity(Identity identity, String query)
+    {
+    }
+
     @Override
     public void checkCanSetSystemSessionProperty(Identity identity, String propertyName)
     {

presto-main/src/main/java/com/facebook/presto/security/ReadOnlySystemAccessControl.java

@@ -58,6 +58,11 @@ public void checkCanSetUser(Optional<Principal> principal, String userName)
     {
     }
 
+    @Override
+    public void checkQueryIntegrity(Identity identity, String query)
+    {
+    }
+
     @Override
     public void checkCanSetSystemSessionProperty(Identity identity, String propertyName)
     {

presto-main/src/test/java/com/facebook/presto/security/TestAccessControlManager.java

@@ -51,18 +51,21 @@
 
 import static com.facebook.presto.spi.ConnectorId.createInformationSchemaConnectorId;
 import static com.facebook.presto.spi.ConnectorId.createSystemTablesConnectorId;
+import static com.facebook.presto.spi.security.AccessDeniedException.denyQueryIntegrityCheck;
 import static com.facebook.presto.spi.security.AccessDeniedException.denySelectColumns;
 import static com.facebook.presto.spi.security.AccessDeniedException.denySelectTable;
 import static com.facebook.presto.transaction.InMemoryTransactionManager.createTestTransactionManager;
 import static com.facebook.presto.transaction.TransactionBuilder.transaction;
 import static java.util.Objects.requireNonNull;
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
 import static org.testng.Assert.fail;
 
 public class TestAccessControlManager
 {
     private static final Principal PRINCIPAL = new BasicPrincipal("principal");
     private static final String USER_NAME = "user_name";
+    private static final String QUERY_TOKEN_FIELD = "query_token";
 
     @Test(expectedExceptions = PrestoException.class, expectedExceptionsMessageRegExp = "Presto server is still initializing")
     public void testInitializing()
@@ -131,6 +134,28 @@ public void testSetAccessControl()
         assertEquals(accessControlFactory.getCheckedPrincipal(), Optional.of(PRINCIPAL));
     }
 
+    @Test
+    public void testCheckQueryIntegrity()
+    {
+        AccessControlManager accessControlManager = new AccessControlManager(createTestTransactionManager());
+
+        TestSystemAccessControlFactory accessControlFactory = new TestSystemAccessControlFactory("test");
+        accessControlManager.addSystemAccessControlFactory(accessControlFactory);
+        accessControlManager.setSystemAccessControl("test", ImmutableMap.of());
+        String testQuery = "test_query";
+
+        accessControlManager.checkQueryIntegrity(new Identity(USER_NAME, Optional.of(PRINCIPAL), ImmutableMap.of(), ImmutableMap.of(QUERY_TOKEN_FIELD, testQuery)), testQuery);
+        assertEquals(accessControlFactory.getCheckedUserName(), USER_NAME);
+        assertEquals(accessControlFactory.getCheckedPrincipal(), Optional.of(PRINCIPAL));
+        assertEquals(accessControlFactory.getCheckedQuery(), testQuery);
+
+        assertThrows(
+                AccessDeniedException.class,
+                () -> accessControlManager.checkQueryIntegrity(
+                        new Identity(USER_NAME, Optional.of(PRINCIPAL), ImmutableMap.of(), ImmutableMap.of(QUERY_TOKEN_FIELD, testQuery + " modified")),
+                        testQuery));
+    }
+
     @Test
     public void testNoCatalogAccessControl()
     {
@@ -219,6 +244,7 @@ private static class TestSystemAccessControlFactory
 
         private Optional<Principal> checkedPrincipal;
         private String checkedUserName;
+        private String checkedQuery;
 
         public TestSystemAccessControlFactory(String name)
         {
@@ -240,6 +266,11 @@ public String getCheckedUserName()
             return checkedUserName;
         }
 
+        public String getCheckedQuery()
+        {
+            return checkedQuery;
+        }
+
         @Override
         public String getName()
         {
@@ -259,6 +290,17 @@ public void checkCanSetUser(Optional<Principal> principal, String userName)
                     checkedUserName = userName;
                 }
 
+                @Override
+                public void checkQueryIntegrity(Identity identity, String query)
+                {
+                    if (!query.equals(identity.getExtraCredentials().get(QUERY_TOKEN_FIELD))) {
+                        denyQueryIntegrityCheck();
+                    }
+                    checkedUserName = identity.getUser();
+                    checkedPrincipal = identity.getPrincipal();
+                    checkedQuery = query;
+                }
+
                 @Override
                 public void checkCanAccessCatalog(Identity identity, String catalogName)
                 {

presto-plugin-toolkit/src/main/java/com/facebook/presto/plugin/base/security/ForwardingSystemAccessControl.java

@@ -52,6 +52,12 @@ public void checkCanSetUser(Optional<Principal> principal, String userName)
         delegate().checkCanSetUser(principal, userName);
     }
 
+    @Override
+    public void checkQueryIntegrity(Identity identity, String query)
+    {
+        delegate().checkQueryIntegrity(identity, query);
+    }
+
     @Override
     public void checkCanSetSystemSessionProperty(Identity identity, String propertyName)
     {

presto-spi/src/main/java/com/facebook/presto/spi/security/AccessDeniedException.java

@@ -36,6 +36,11 @@ public static void denySetUser(Optional<Principal> principal, String userName)
         denySetUser(principal, userName, null);
     }
 
+    public static void denyQueryIntegrityCheck()
+    {
+        throw new AccessDeniedException("Query integrity check failed.");
+    }
+
     public static void denySetUser(Optional<Principal> principal, String userName, String extraInfo)
     {
         throw new AccessDeniedException(format("Principal %s cannot become user %s%s", principal.orElse(null), userName, formatExtraInfo(extraInfo)));

presto-spi/src/main/java/com/facebook/presto/spi/security/SystemAccessControl.java

@@ -53,6 +53,12 @@ public interface SystemAccessControl
      */
     void checkCanSetUser(Optional<Principal> principal, String userName);
 
+    /**
+     * Check if the query is unexpectedly modified using the credentials passed in the identity.
+     * @throws AccessDeniedException if query is modified.
+     */
+    void checkQueryIntegrity(Identity identity, String query);
+
     /**
      * Check if identity is allowed to set the specified system property.
      *