Enumerate roles until enabled role is found

Previously when SqlStandardAccessControl was checking if given role is enabled, it listed all role grants and check if that role is is among all listed role grants. Now it list all role grants until it finds that role. Extracted-From: prestodb/presto#10904
trinodb · Jan 29, 2019 · 6910a27 · 6910a27
1 parent 7427e3a
commit 6910a27
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 4 deletions.
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/thrift/ThriftMetastoreUtil.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/thrift/ThriftMetastoreUtil.java
@@ -236,6 +236,28 @@ public static Set<RoleGrant> listApplicableRoles(PrestoPrincipal principal, Func
         return ImmutableSet.copyOf(result);
     }
 
+    public static boolean isRoleApplicable(SemiTransactionalHiveMetastore metastore, PrestoPrincipal principal, String role)
+    {
+        if (principal.getType() == ROLE && principal.getName().equals(role)) {
+            return true;
+        }
+        Set<String> seenRoles = new HashSet<>();
+        Queue<PrestoPrincipal> queue = new ArrayDeque<>();
+        queue.add(principal);
+        while (!queue.isEmpty()) {
+            Set<RoleGrant> grants = metastore.listRoleGrants(queue.poll());
+            for (RoleGrant grant : grants) {
+                if (grant.getRoleName().equals(role)) {
+                    return true;
+                }
+                if (seenRoles.add(grant.getRoleName())) {
+                    queue.add(new PrestoPrincipal(ROLE, grant.getRoleName()));
+                }
+            }
+        }
+        return false;
+    }
+
     public static Set<String> listApplicableRoles(SemiTransactionalHiveMetastore metastore, PrestoPrincipal principal)
     {
         return listApplicableRoles(principal, metastore::listRoleGrants)
@@ -276,6 +298,50 @@ private static Set<HivePrivilegeInfo> listTablePrivileges(SemiTransactionalHiveM
         return result.build();
     }
 
+    public static boolean isRoleEnabled(ConnectorIdentity identity, Function<PrestoPrincipal, Set<RoleGrant>> listRoleGrants, String role)
+    {
+        if (role.equals(PUBLIC_ROLE_NAME)) {
+            return true;
+        }
+
+        if (identity.getRole().isPresent() && identity.getRole().get().getType() == SelectedRole.Type.NONE) {
+            return false;
+        }
+
+        PrestoPrincipal principal;
+        if (!identity.getRole().isPresent() || identity.getRole().get().getType() == SelectedRole.Type.ALL) {
+            principal = new PrestoPrincipal(USER, identity.getUser());
+        }
+        else {
+            principal = new PrestoPrincipal(ROLE, identity.getRole().get().getRole().get());
+        }
+
+        if (principal.getType() == ROLE && principal.getName().equals(role)) {
+            return true;
+        }
+
+        if (role.equals(ADMIN_ROLE_NAME)) {
+            // The admin role must be enabled explicitly, and so it should checked above
+            return false;
+        }
+
+        Set<String> seenRoles = new HashSet<>();
+        Queue<PrestoPrincipal> queue = new ArrayDeque<>();
+        queue.add(principal);
+        while (!queue.isEmpty()) {
+            Set<RoleGrant> grants = listRoleGrants.apply(queue.poll());
+            for (RoleGrant grant : grants) {
+                if (grant.getRoleName().equals(role)) {
+                    return true;
+                }
+                if (seenRoles.add(grant.getRoleName())) {
+                    queue.add(new PrestoPrincipal(ROLE, grant.getRoleName()));
+                }
+            }
+        }
+        return false;
+    }
+
     public static Set<String> listEnabledRoles(ConnectorIdentity identity, Function<PrestoPrincipal, Set<RoleGrant>> listRoleGrants)
     {
         Optional<SelectedRole> role = identity.getRole();

diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/security/SqlStandardAccessControl.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/security/SqlStandardAccessControl.java
@@ -39,9 +39,10 @@
 import static io.prestosql.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege.OWNERSHIP;
 import static io.prestosql.plugin.hive.metastore.HivePrivilegeInfo.HivePrivilege.SELECT;
 import static io.prestosql.plugin.hive.metastore.HivePrivilegeInfo.toHivePrivilege;
+import static io.prestosql.plugin.hive.metastore.thrift.ThriftMetastoreUtil.isRoleApplicable;
+import static io.prestosql.plugin.hive.metastore.thrift.ThriftMetastoreUtil.isRoleEnabled;
 import static io.prestosql.plugin.hive.metastore.thrift.ThriftMetastoreUtil.listApplicableRoles;
 import static io.prestosql.plugin.hive.metastore.thrift.ThriftMetastoreUtil.listApplicableTablePrivileges;
-import static io.prestosql.plugin.hive.metastore.thrift.ThriftMetastoreUtil.listEnabledRoles;
 import static io.prestosql.plugin.hive.metastore.thrift.ThriftMetastoreUtil.listEnabledTablePrivileges;
 import static io.prestosql.spi.security.AccessDeniedException.denyAddColumn;
 import static io.prestosql.spi.security.AccessDeniedException.denyCreateRole;
@@ -317,7 +318,7 @@ public void checkCanRevokeRoles(ConnectorTransactionHandle transactionHandle, Co
     public void checkCanSetRole(ConnectorTransactionHandle transaction, ConnectorIdentity identity, String role, String catalogName)
     {
         SemiTransactionalHiveMetastore metastore = metastoreProvider.apply(((HiveTransactionHandle) transaction));
-        if (!listApplicableRoles(metastore, new PrestoPrincipal(USER, identity.getUser())).contains(role)) {
+        if (!isRoleApplicable(metastore, new PrestoPrincipal(USER, identity.getUser()), role)) {
             denySetRole(role);
         }
     }
@@ -343,7 +344,7 @@ public void checkCanShowRoleGrants(ConnectorTransactionHandle transactionHandle,
     private boolean isAdmin(ConnectorTransactionHandle transaction, ConnectorIdentity identity)
     {
         SemiTransactionalHiveMetastore metastore = metastoreProvider.apply(((HiveTransactionHandle) transaction));
-        return listEnabledRoles(identity, metastore::listRoleGrants).contains(ADMIN_ROLE_NAME);
+        return isRoleEnabled(identity, metastore::listRoleGrants, ADMIN_ROLE_NAME);
     }
 
     private boolean isDatabaseOwner(ConnectorTransactionHandle transaction, ConnectorIdentity identity, String databaseName)
@@ -369,7 +370,7 @@ private boolean isDatabaseOwner(ConnectorTransactionHandle transaction, Connecto
         if (database.getOwnerType() == USER && identity.getUser().equals(database.getOwnerName())) {
             return true;
         }
-        if (database.getOwnerType() == ROLE && listEnabledRoles(identity, metastore::listRoleGrants).contains(database.getOwnerName())) {
+        if (database.getOwnerType() == ROLE && isRoleEnabled(identity, metastore::listRoleGrants, database.getOwnerName())) {
             return true;
         }
         return false;