From 1c27ad91a31fffb8dd8b0ceefbd7ad57ffcff05e Mon Sep 17 00:00:00 2001 From: HangyuanLiu <460660596@qq.com> Date: Fri, 20 Oct 2023 10:31:33 +0800 Subject: [PATCH] [Enhancement] Support is_role_in_session function (#32984) Signed-off-by: HangyuanLiu <460660596@qq.com> (cherry picked from commit b6954d3b4262c392296f4ce3f7d72b5b59f7b3a7) Signed-off-by: HangyuanLiu <460660596@qq.com> # Conflicts: # fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorFunctions.java --- .../com/starrocks/catalog/FunctionSet.java | 3 + .../starrocks/privilege/AuthorizationMgr.java | 13 ++- .../sql/analyzer/ExpressionAnalyzer.java | 11 +++ .../rewrite/ScalarOperatorFunctions.java | 15 +++ .../AuthenticationManagerTest.java | 91 +++++++++++++++++++ gensrc/script/functions.py | 3 + test/sql/test_rbac/R/test_is_role_in_session | 70 ++++++++++++++ test/sql/test_rbac/T/test_is_role_in_session | 25 +++++ 8 files changed, 230 insertions(+), 1 deletion(-) create mode 100644 test/sql/test_rbac/R/test_is_role_in_session create mode 100644 test/sql/test_rbac/T/test_is_role_in_session diff --git a/fe/fe-core/src/main/java/com/starrocks/catalog/FunctionSet.java b/fe/fe-core/src/main/java/com/starrocks/catalog/FunctionSet.java index 7a1a4a5e23bd1..6957137ed1b8e 100644 --- a/fe/fe-core/src/main/java/com/starrocks/catalog/FunctionSet.java +++ b/fe/fe-core/src/main/java/com/starrocks/catalog/FunctionSet.java @@ -456,6 +456,9 @@ public class FunctionSet { // dict query function public static final String DICT_MAPPING = "dict_mapping"; + //user and role function + public static final String IS_ROLE_IN_SESSION = "is_role_in_session"; + public static final String QUARTERS_ADD = "quarters_add"; public static final String QUARTERS_SUB = "quarters_sub"; public static final String WEEKS_ADD = "weeks_add"; diff --git a/fe/fe-core/src/main/java/com/starrocks/privilege/AuthorizationMgr.java b/fe/fe-core/src/main/java/com/starrocks/privilege/AuthorizationMgr.java index 8c2b992d66190..6d3fbd244ff88 100644 --- a/fe/fe-core/src/main/java/com/starrocks/privilege/AuthorizationMgr.java +++ b/fe/fe-core/src/main/java/com/starrocks/privilege/AuthorizationMgr.java @@ -701,7 +701,7 @@ public boolean canExecuteAs(UserIdentity currentUser, Set roleIds, UserIde } } - public boolean allowGrant(UserIdentity currentUser, Set roleIds, ObjectType type, + public boolean allowGrant(UserIdentity currentUser, Set roleIds, ObjectType type, List wants, List objects) { try { PrivilegeCollectionV2 collection = mergePrivilegeCollection(currentUser, roleIds); @@ -974,6 +974,17 @@ public RolePrivilegeCollectionV2 getRolePrivilegeCollection(long roleId) { } } + public void getRecursiveRole(Set roleNames, Long roleId) { + RolePrivilegeCollectionV2 rolePrivilegeCollection = getRolePrivilegeCollection(roleId); + if (rolePrivilegeCollection != null) { + roleNames.add(rolePrivilegeCollection.getName()); + + for (Long parentId : rolePrivilegeCollection.getParentRoleIds()) { + getRecursiveRole(roleNames, parentId); + } + } + } + public RolePrivilegeCollectionV2 getRolePrivilegeCollectionUnlocked(long roleId, boolean exceptionIfNotExists) throws PrivilegeException { RolePrivilegeCollectionV2 collection = roleIdToPrivilegeCollection.get(roleId); diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/analyzer/ExpressionAnalyzer.java b/fe/fe-core/src/main/java/com/starrocks/sql/analyzer/ExpressionAnalyzer.java index 6a1cfef08d49c..335832caf4a2f 100644 --- a/fe/fe-core/src/main/java/com/starrocks/sql/analyzer/ExpressionAnalyzer.java +++ b/fe/fe-core/src/main/java/com/starrocks/sql/analyzer/ExpressionAnalyzer.java @@ -1400,6 +1400,17 @@ private void checkFunction(String fnName, FunctionCallExpr node, Type[] argument } break; } + + case FunctionSet.IS_ROLE_IN_SESSION: { + if (node.getChildren().size() != 1) { + throw new SemanticException("IS_ROLE_IN_SESSION currently only supports a single parameter"); + } + + if (!(node.getChild(0) instanceof StringLiteral)) { + throw new SemanticException("IS_ROLE_IN_SESSION currently only supports constant parameters"); + } + break; + } } } diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorFunctions.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorFunctions.java index 40173521e80f3..b00f4cbb8129b 100644 --- a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorFunctions.java +++ b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorFunctions.java @@ -61,6 +61,7 @@ import com.starrocks.connector.PartitionUtil; import com.starrocks.connector.hive.Partition; import com.starrocks.privilege.AccessDeniedException; +import com.starrocks.privilege.AuthorizationMgr; import com.starrocks.privilege.ObjectType; import com.starrocks.privilege.PrivilegeType; import com.starrocks.qe.ConnectContext; @@ -90,6 +91,7 @@ import java.time.temporal.IsoFields; import java.time.temporal.TemporalAdjusters; import java.time.temporal.TemporalUnit; +import java.util.HashSet; import java.util.Map; import java.util.Optional; import java.util.Set; @@ -1303,4 +1305,17 @@ public static ConstantOperator coalesce(ConstantOperator... values) { } return values[values.length - 1]; } + + @ConstantFunction(name = "is_role_in_session", argTypes = {VARCHAR}, returnType = BOOLEAN) + public static ConstantOperator isRoleInSession(ConstantOperator role) { + AuthorizationMgr manager = GlobalStateMgr.getCurrentState().getAuthorizationMgr(); + Set roleNames = new HashSet<>(); + ConnectContext connectContext = ConnectContext.get(); + + for (Long roleId : connectContext.getCurrentRoleIds()) { + manager.getRecursiveRole(roleNames, roleId); + } + + return ConstantOperator.createBoolean(roleNames.contains(role.getVarchar())); + } } diff --git a/fe/fe-core/src/test/java/com/starrocks/authentication/AuthenticationManagerTest.java b/fe/fe-core/src/test/java/com/starrocks/authentication/AuthenticationManagerTest.java index 340103b16cecd..e5be7e4c181b9 100644 --- a/fe/fe-core/src/test/java/com/starrocks/authentication/AuthenticationManagerTest.java +++ b/fe/fe-core/src/test/java/com/starrocks/authentication/AuthenticationManagerTest.java @@ -28,10 +28,13 @@ import com.starrocks.sql.ast.CreateRoleStmt; import com.starrocks.sql.ast.CreateUserStmt; import com.starrocks.sql.ast.DropUserStmt; +import com.starrocks.sql.ast.GrantRoleStmt; import com.starrocks.sql.ast.SetDefaultRoleStmt; import com.starrocks.sql.ast.SetUserPropertyStmt; import com.starrocks.sql.ast.StatementBase; import com.starrocks.sql.ast.UserIdentity; +import com.starrocks.sql.optimizer.operator.scalar.ConstantOperator; +import com.starrocks.sql.optimizer.rewrite.ScalarOperatorFunctions; import com.starrocks.utframe.UtFrameUtils; import org.junit.AfterClass; import org.junit.Assert; @@ -480,4 +483,92 @@ public void testSortUserIdentity() throws Exception { Assert.assertEquals(Arrays.asList( "'sort_user'@'10.1.1.1'", "'sort_user'@'10.1.1.2'", "'sort_user'@['host01']", "'sort_user'@'%'"), l); } + + @Test + public void testIsRoleInSession() throws Exception { + AuthenticationMgr masterManager = ctx.getGlobalStateMgr().getAuthenticationMgr(); + AuthorizationMgr authorizationManager = ctx.getGlobalStateMgr().getAuthorizationMgr(); + + String sql = "create role test_in_role_r1"; + CreateRoleStmt createStmt = + (CreateRoleStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + authorizationManager.createRole(createStmt); + + sql = "create role test_in_role_r2"; + createStmt = (CreateRoleStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + authorizationManager.createRole(createStmt); + + sql = "create role test_in_role_r3"; + createStmt = (CreateRoleStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + authorizationManager.createRole(createStmt); + + sql = "create role test_in_role_r4"; + createStmt = (CreateRoleStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + authorizationManager.createRole(createStmt); + + sql = "grant test_in_role_r3 to role test_in_role_r2"; + GrantRoleStmt grantRoleStmt = (GrantRoleStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + authorizationManager.grantRole(grantRoleStmt); + + sql = "grant test_in_role_r2 to role test_in_role_r1"; + grantRoleStmt = (GrantRoleStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + authorizationManager.grantRole(grantRoleStmt); + + sql = "create user test_in_role_u1 default role test_in_role_r1"; + CreateUserStmt stmt = (CreateUserStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + masterManager.createUser(stmt); + + ctx.setCurrentUserIdentity(new UserIdentity("test_in_role_u1", "%")); + ctx.setCurrentRoleIds(new UserIdentity("test_in_role_u1", "%")); + + Assert.assertTrue(ScalarOperatorFunctions.isRoleInSession( + ConstantOperator.createVarchar("test_in_role_r1")).getBoolean()); + Assert.assertTrue(ScalarOperatorFunctions.isRoleInSession( + ConstantOperator.createVarchar("test_in_role_r2")).getBoolean()); + Assert.assertTrue(ScalarOperatorFunctions.isRoleInSession( + ConstantOperator.createVarchar("test_in_role_r3")).getBoolean()); + + Assert.assertFalse(ScalarOperatorFunctions.isRoleInSession( + ConstantOperator.createVarchar("test_in_role_r4")).getBoolean()); + + sql = "create user test_in_role_u2 default role test_in_role_r2"; + stmt = (CreateUserStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + masterManager.createUser(stmt); + + ctx.setCurrentUserIdentity(new UserIdentity("test_in_role_u2", "%")); + ctx.setCurrentRoleIds(new UserIdentity("test_in_role_u2", "%")); + + Assert.assertFalse(ScalarOperatorFunctions.isRoleInSession( + ConstantOperator.createVarchar("test_in_role_r1")).getBoolean()); + Assert.assertTrue(ScalarOperatorFunctions.isRoleInSession( + ConstantOperator.createVarchar("test_in_role_r2")).getBoolean()); + Assert.assertTrue(ScalarOperatorFunctions.isRoleInSession( + ConstantOperator.createVarchar("test_in_role_r3")).getBoolean()); + + ctx.setCurrentRoleIds(new HashSet<>()); + + Assert.assertFalse(ScalarOperatorFunctions.isRoleInSession( + ConstantOperator.createVarchar("test_in_role_r1")).getBoolean()); + Assert.assertFalse(ScalarOperatorFunctions.isRoleInSession( + ConstantOperator.createVarchar("test_in_role_r2")).getBoolean()); + Assert.assertFalse(ScalarOperatorFunctions.isRoleInSession( + ConstantOperator.createVarchar("test_in_role_r3")).getBoolean()); + + + sql = "select is_role_in_session(v1) from (select 1 as v1) t"; + try { + stmt = (CreateUserStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + Assert.fail(); + } catch (AnalysisException e) { + Assert.assertTrue(e.getMessage().contains("IS_ROLE_IN_SESSION currently only supports constant parameters")); + } + + sql = "select is_role_in_session(\"a\", \"b\") from (select 1 as v1) t"; + try { + stmt = (CreateUserStmt) UtFrameUtils.parseStmtWithNewParser(sql, ctx); + Assert.fail(); + } catch (AnalysisException e) { + Assert.assertTrue(e.getMessage().contains("IS_ROLE_IN_SESSION currently only supports a single parameter")); + } + } } diff --git a/gensrc/script/functions.py b/gensrc/script/functions.py index 689495287ed33..c5a30470a2f27 100644 --- a/gensrc/script/functions.py +++ b/gensrc/script/functions.py @@ -1011,4 +1011,7 @@ # struct functions [170500, 'row', 'ANY_STRUCT', ['ANY_ELEMENT', "..."], 'StructFunctions::new_struct'], [170501, 'named_struct', 'ANY_STRUCT', ['ANY_ELEMENT', "..."], 'StructFunctions::named_struct'], + + # user function + [180000, 'is_role_in_session', 'BOOLEAN', ['VARCHAR'], 'nullptr'] ] diff --git a/test/sql/test_rbac/R/test_is_role_in_session b/test/sql/test_rbac/R/test_is_role_in_session new file mode 100644 index 0000000000000..7d6fdfca954bd --- /dev/null +++ b/test/sql/test_rbac/R/test_is_role_in_session @@ -0,0 +1,70 @@ +-- name: test_is_role_in_session +drop role if exists r1; +-- result: +-- !result +create role r1; +-- result: +-- !result +drop role if exists r2; +-- result: +-- !result +create role r2; +-- result: +-- !result +drop role if exists r3; +-- result: +-- !result +create role r3; +-- result: +-- !result +drop user if exists u1; +-- result: +-- !result +create user u1; +-- result: +-- !result +grant impersonate on user root to u1; +-- result: +-- !result +grant r3 to role r2; +-- result: +-- !result +grant r2 to role r1; +-- result: +-- !result +grant r1 to u1; +-- result: +-- !result +execute as u1 with no revert; +-- result: +-- !result +select is_role_in_session("r1"); +-- result: +0 +-- !result +select is_role_in_session("r2"); +-- result: +0 +-- !result +select is_role_in_session("r3"); +-- result: +0 +-- !result +set role all; +-- result: +-- !result +select is_role_in_session("r1"); +-- result: +1 +-- !result +select is_role_in_session("r2"); +-- result: +1 +-- !result +select is_role_in_session("r3"); +-- result: +1 +-- !result +execute as root with no revert; +-- result: +-- !result \ No newline at end of file diff --git a/test/sql/test_rbac/T/test_is_role_in_session b/test/sql/test_rbac/T/test_is_role_in_session new file mode 100644 index 0000000000000..a59e2c92152d8 --- /dev/null +++ b/test/sql/test_rbac/T/test_is_role_in_session @@ -0,0 +1,25 @@ +-- name: test_is_role_in_session +drop role if exists r1; +create role r1; +drop role if exists r2; +create role r2; +drop role if exists r3; +create role r3; +drop user if exists u1; +create user u1; +grant impersonate on user root to u1; + +grant r3 to role r2; +grant r2 to role r1; +grant r1 to u1; + +execute as u1 with no revert; +select is_role_in_session("r1"); +select is_role_in_session("r2"); +select is_role_in_session("r3"); +set role all; +select is_role_in_session("r1"); +select is_role_in_session("r2"); +select is_role_in_session("r3"); + +execute as root with no revert; \ No newline at end of file