Change AccessControl filterCatalogs to have a SecurityContext

AccessControl.filterCatalogs did not accept a SecurityContext like the schema/table/column filters. This caused some places in the code to recreate the SecurityContext and thus eliminate the associated queryId.
trinodb · Jan 21, 2022 · 0228bad · 0228bad
1 parent f81af8f
commit 0228bad
Show file tree

Hide file tree

Showing 10 changed files with 24 additions and 24 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/metadata/MetadataListing.java b/core/trino-main/src/main/java/io/trino/metadata/MetadataListing.java
@@ -65,7 +65,7 @@ public static SortedMap<String, CatalogName> listCatalogs(Session session, Metad
                             Map.Entry::getKey,
                             entry -> entry.getValue().getConnectorCatalogName()));
         }
-        Set<String> allowedCatalogs = accessControl.filterCatalogs(session.getIdentity(), catalogNames.keySet());
+        Set<String> allowedCatalogs = accessControl.filterCatalogs(session.toSecurityContext(), catalogNames.keySet());
 
         ImmutableSortedMap.Builder<String, CatalogName> result = ImmutableSortedMap.naturalOrder();
         for (Map.Entry<String, CatalogName> entry : catalogNames.entrySet()) {
@@ -79,7 +79,7 @@ public static SortedMap<String, CatalogName> listCatalogs(Session session, Metad
     public static SortedMap<String, Catalog> getCatalogs(Session session, Metadata metadata, AccessControl accessControl)
     {
         Map<String, Catalog> catalogs = metadata.getCatalogs(session);
-        Set<String> allowedCatalogs = accessControl.filterCatalogs(session.getIdentity(), catalogs.keySet());
+        Set<String> allowedCatalogs = accessControl.filterCatalogs(session.toSecurityContext(), catalogs.keySet());
 
         ImmutableSortedMap.Builder<String, Catalog> result = ImmutableSortedMap.naturalOrder();
         for (Map.Entry<String, Catalog> entry : catalogs.entrySet()) {

diff --git a/core/trino-main/src/main/java/io/trino/security/AccessControl.java b/core/trino-main/src/main/java/io/trino/security/AccessControl.java
@@ -101,7 +101,7 @@ public interface AccessControl
     /**
      * Filter the list of catalogs to those visible to the identity.
      */
-    Set<String> filterCatalogs(Identity identity, Set<String> catalogs);
+    Set<String> filterCatalogs(SecurityContext context, Set<String> catalogs);
 
     /**
      * Check if identity is allowed to create the specified schema.

diff --git a/core/trino-main/src/main/java/io/trino/security/AccessControlManager.java b/core/trino-main/src/main/java/io/trino/security/AccessControlManager.java
@@ -290,13 +290,13 @@ public void checkCanKillQueryOwnedBy(Identity identity, Identity queryOwner)
     }
 
     @Override
-    public Set<String> filterCatalogs(Identity identity, Set<String> catalogs)
+    public Set<String> filterCatalogs(SecurityContext securityContext, Set<String> catalogs)
     {
-        requireNonNull(identity, "identity is null");
+        requireNonNull(securityContext, "securityContext is null");
         requireNonNull(catalogs, "catalogs is null");
 
         for (SystemAccessControl systemAccessControl : getSystemAccessControls()) {
-            catalogs = systemAccessControl.filterCatalogs(new SystemSecurityContext(identity, Optional.empty()), catalogs);
+            catalogs = systemAccessControl.filterCatalogs(securityContext.toSystemSecurityContext(), catalogs);
         }
         return catalogs;
     }
@@ -373,7 +373,7 @@ public Set<String> filterSchemas(SecurityContext securityContext, String catalog
         requireNonNull(catalogName, "catalogName is null");
         requireNonNull(schemaNames, "schemaNames is null");
 
-        if (filterCatalogs(securityContext.getIdentity(), ImmutableSet.of(catalogName)).isEmpty()) {
+        if (filterCatalogs(securityContext, ImmutableSet.of(catalogName)).isEmpty()) {
             return ImmutableSet.of();
         }
 
@@ -527,7 +527,7 @@ public Set<SchemaTableName> filterTables(SecurityContext securityContext, String
         requireNonNull(catalogName, "catalogName is null");
         requireNonNull(tableNames, "tableNames is null");
 
-        if (filterCatalogs(securityContext.getIdentity(), ImmutableSet.of(catalogName)).isEmpty()) {
+        if (filterCatalogs(securityContext, ImmutableSet.of(catalogName)).isEmpty()) {
             return ImmutableSet.of();
         }
 

diff --git a/core/trino-main/src/main/java/io/trino/security/AllowAllAccessControl.java b/core/trino-main/src/main/java/io/trino/security/AllowAllAccessControl.java
@@ -72,7 +72,7 @@ public void checkCanKillQueryOwnedBy(Identity identity, Identity queryOwner)
     }
 
     @Override
-    public Set<String> filterCatalogs(Identity identity, Set<String> catalogs)
+    public Set<String> filterCatalogs(SecurityContext context, Set<String> catalogs)
     {
         return catalogs;
     }

diff --git a/core/trino-main/src/main/java/io/trino/security/DenyAllAccessControl.java b/core/trino-main/src/main/java/io/trino/security/DenyAllAccessControl.java
@@ -143,7 +143,7 @@ public void checkCanKillQueryOwnedBy(Identity identity, Identity queryOwner)
     }
 
     @Override
-    public Set<String> filterCatalogs(Identity identity, Set<String> catalogs)
+    public Set<String> filterCatalogs(SecurityContext context, Set<String> catalogs)
     {
         return ImmutableSet.of();
     }

diff --git a/core/trino-main/src/main/java/io/trino/security/ForwardingAccessControl.java b/core/trino-main/src/main/java/io/trino/security/ForwardingAccessControl.java
@@ -101,9 +101,9 @@ public void checkCanKillQueryOwnedBy(Identity identity, Identity queryOwner)
     }
 
     @Override
-    public Set<String> filterCatalogs(Identity identity, Set<String> catalogs)
+    public Set<String> filterCatalogs(SecurityContext context, Set<String> catalogs)
     {
-        return delegate().filterCatalogs(identity, catalogs);
+        return delegate().filterCatalogs(context, catalogs);
     }
 
     @Override

diff --git a/core/trino-main/src/main/java/io/trino/testing/AllowAllAccessControlManager.java b/core/trino-main/src/main/java/io/trino/testing/AllowAllAccessControlManager.java
@@ -60,7 +60,7 @@ public Collection<Identity> filterQueriesOwnedBy(Identity identity, Collection<I
     public void checkCanKillQueryOwnedBy(Identity identity, Identity queryOwner) {}
 
     @Override
-    public Set<String> filterCatalogs(Identity identity, Set<String> catalogs)
+    public Set<String> filterCatalogs(SecurityContext context, Set<String> catalogs)
     {
         return catalogs;
     }

diff --git a/core/trino-main/src/main/java/io/trino/testing/TestingAccessControlManager.java b/core/trino-main/src/main/java/io/trino/testing/TestingAccessControlManager.java
@@ -200,10 +200,10 @@ public void denyIdentityTable(BiPredicate<Identity, String> denyIdentityTable)
     }
 
     @Override
-    public Set<String> filterCatalogs(Identity identity, Set<String> catalogs)
+    public Set<String> filterCatalogs(SecurityContext securityContext, Set<String> catalogs)
     {
         return super.filterCatalogs(
-                identity,
+                securityContext,
                 catalogs.stream()
                         .filter(this.deniedCatalogs)
                         .collect(toImmutableSet()));

diff --git a/core/trino-main/src/test/java/io/trino/security/TestAccessControlManager.java b/core/trino-main/src/test/java/io/trino/security/TestAccessControlManager.java
@@ -126,7 +126,7 @@ public void testReadOnlySystemAccessControl()
                     accessControlManager.checkCanGrantExecuteFunctionPrivilege(context, "function", Identity.ofUser("bob"), false);
                     accessControlManager.checkCanGrantExecuteFunctionPrivilege(context, "function", Identity.ofUser("bob"), true);
                     Set<String> catalogs = ImmutableSet.of("catalog");
-                    assertEquals(accessControlManager.filterCatalogs(identity, catalogs), catalogs);
+                    assertEquals(accessControlManager.filterCatalogs(context, catalogs), catalogs);
                     Set<String> schemas = ImmutableSet.of("schema");
                     assertEquals(accessControlManager.filterSchemas(context, "catalog", schemas), schemas);
                     Set<SchemaTableName> tableNames = ImmutableSet.of(new SchemaTableName("schema", "table"));

diff --git a/core/trino-main/src/test/java/io/trino/security/TestFileBasedSystemAccessControl.java b/core/trino-main/src/test/java/io/trino/security/TestFileBasedSystemAccessControl.java
@@ -216,13 +216,13 @@ public void testCatalogOperations()
 
         transaction(transactionManager, accessControlManager)
                 .execute(transactionId -> {
-                    assertEquals(accessControlManager.filterCatalogs(admin, allCatalogs), allCatalogs);
+                    assertEquals(accessControlManager.filterCatalogs(new SecurityContext(transactionId, admin, queryId), allCatalogs), allCatalogs);
                     Set<String> aliceCatalogs = ImmutableSet.of("open-to-all", "alice-catalog", "all-allowed", "staff-catalog");
-                    assertEquals(accessControlManager.filterCatalogs(alice, allCatalogs), aliceCatalogs);
+                    assertEquals(accessControlManager.filterCatalogs(new SecurityContext(transactionId, alice, queryId), allCatalogs), aliceCatalogs);
                     Set<String> bobCatalogs = ImmutableSet.of("open-to-all", "all-allowed", "staff-catalog");
-                    assertEquals(accessControlManager.filterCatalogs(bob, allCatalogs), bobCatalogs);
+                    assertEquals(accessControlManager.filterCatalogs(new SecurityContext(transactionId, bob, queryId), allCatalogs), bobCatalogs);
                     Set<String> nonAsciiUserCatalogs = ImmutableSet.of("open-to-all", "all-allowed", "\u0200\u0200\u0200");
-                    assertEquals(accessControlManager.filterCatalogs(nonAsciiUser, allCatalogs), nonAsciiUserCatalogs);
+                    assertEquals(accessControlManager.filterCatalogs(new SecurityContext(transactionId, nonAsciiUser, queryId), allCatalogs), nonAsciiUserCatalogs);
                 });
     }
 
@@ -234,13 +234,13 @@ public void testCatalogOperationsReadOnly()
 
         transaction(transactionManager, accessControlManager)
                 .execute(transactionId -> {
-                    assertEquals(accessControlManager.filterCatalogs(admin, allCatalogs), allCatalogs);
+                    assertEquals(accessControlManager.filterCatalogs(new SecurityContext(transactionId, admin, queryId), allCatalogs), allCatalogs);
                     Set<String> aliceCatalogs = ImmutableSet.of("open-to-all", "alice-catalog", "all-allowed");
-                    assertEquals(accessControlManager.filterCatalogs(alice, allCatalogs), aliceCatalogs);
+                    assertEquals(accessControlManager.filterCatalogs(new SecurityContext(transactionId, alice, queryId), allCatalogs), aliceCatalogs);
                     Set<String> bobCatalogs = ImmutableSet.of("open-to-all", "all-allowed");
-                    assertEquals(accessControlManager.filterCatalogs(bob, allCatalogs), bobCatalogs);
+                    assertEquals(accessControlManager.filterCatalogs(new SecurityContext(transactionId, bob, queryId), allCatalogs), bobCatalogs);
                     Set<String> nonAsciiUserCatalogs = ImmutableSet.of("open-to-all", "all-allowed", "\u0200\u0200\u0200");
-                    assertEquals(accessControlManager.filterCatalogs(nonAsciiUser, allCatalogs), nonAsciiUserCatalogs);
+                    assertEquals(accessControlManager.filterCatalogs(new SecurityContext(transactionId, nonAsciiUser, queryId), allCatalogs), nonAsciiUserCatalogs);
                 });
     }