Skip to content

Commit c9b441a

Browse files
authored
Refactor TreeAccessCheckVisitor
1 parent 4da1c88 commit c9b441a

File tree

3 files changed

+72
-31
lines changed

3 files changed

+72
-31
lines changed

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java

Lines changed: 62 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -82,9 +82,14 @@
8282
import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowVariablesStatement;
8383
import org.apache.iotdb.db.queryengine.plan.statement.metadata.UnSetTTLStatement;
8484
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.CreateModelStatement;
85+
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.CreateTrainingStatement;
8586
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.DropModelStatement;
87+
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.LoadModelStatement;
88+
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowAIDevicesStatement;
8689
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowAINodesStatement;
90+
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowLoadedModelsStatement;
8791
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowModelsStatement;
92+
import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.UnloadModelStatement;
8893
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.AlterPipeStatement;
8994
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipePluginStatement;
9095
import org.apache.iotdb.db.queryengine.plan.statement.metadata.pipe.CreatePipeStatement;
@@ -123,8 +128,10 @@
123128
import org.apache.iotdb.db.queryengine.plan.statement.sys.AuthorStatement;
124129
import org.apache.iotdb.db.queryengine.plan.statement.sys.ClearCacheStatement;
125130
import org.apache.iotdb.db.queryengine.plan.statement.sys.ExplainAnalyzeStatement;
131+
import org.apache.iotdb.db.queryengine.plan.statement.sys.ExplainStatement;
126132
import org.apache.iotdb.db.queryengine.plan.statement.sys.FlushStatement;
127133
import org.apache.iotdb.db.queryengine.plan.statement.sys.KillQueryStatement;
134+
import org.apache.iotdb.db.queryengine.plan.statement.sys.LoadConfigurationStatement;
128135
import org.apache.iotdb.db.queryengine.plan.statement.sys.SetConfigurationStatement;
129136
import org.apache.iotdb.db.queryengine.plan.statement.sys.SetSqlDialectStatement;
130137
import org.apache.iotdb.db.queryengine.plan.statement.sys.SetSystemStatusStatement;
@@ -161,10 +168,7 @@ public class TreeAccessCheckVisitor extends StatementVisitor<TSStatus, TreeAcces
161168

162169
@Override
163170
public TSStatus visitNode(StatementNode node, TreeAccessCheckContext context) {
164-
if (AuthorityChecker.SUPER_USER.equals(context.getUsername())) {
165-
return SUCCEED;
166-
}
167-
return AuthorityChecker.getTSStatus(false, "Only the admin user can perform this operation");
171+
throw new IllegalStateException("Each operation should have permission check.");
168172
}
169173

170174
@Override
@@ -229,7 +233,7 @@ private TSStatus checkTemplateShowRelated(
229233
return SUCCEED;
230234
}
231235
// own SYSTEM can see all, otherwise can only see PATHS that user has READ_SCHEMA auth
232-
if (!AuthorityChecker.checkSystemPermission(context.getUsername(), PrivilegeType.SYSTEM)) {
236+
if (!checkHasGlobalAuth(context.getUsername(), PrivilegeType.SYSTEM)) {
233237
statement.setCanSeeAll(false);
234238
return visitAuthorityInformation(statement, context);
235239
} else {
@@ -289,11 +293,7 @@ public TSStatus visitAlterSchemaTemplate(
289293
if (AuthorityChecker.SUPER_USER.equals(context.getUsername())) {
290294
return SUCCEED;
291295
}
292-
return AuthorityChecker.getTSStatus(
293-
AuthorityChecker.checkSystemPermission(context.getUsername(), PrivilegeType.SYSTEM)
294-
|| AuthorityChecker.checkSystemPermission(
295-
context.getUsername(), PrivilegeType.EXTEND_TEMPLATE),
296-
PrivilegeType.SYSTEM);
296+
return checkGlobalAuth(context.getUsername(), PrivilegeType.EXTEND_TEMPLATE);
297297
}
298298

299299
// ============================= timeseries view related ===============
@@ -455,8 +455,7 @@ public TSStatus visitAuthor(AuthorStatement statement, TreeAccessCheckContext co
455455
}
456456

457457
case LIST_ROLE:
458-
if (AuthorityChecker.checkSystemPermission(
459-
context.getUsername(), PrivilegeType.MANAGE_ROLE)) {
458+
if (checkHasGlobalAuth(context.getUsername(), PrivilegeType.MANAGE_ROLE)) {
460459
return SUCCEED;
461460
}
462461
// list roles of other user is not allowed
@@ -572,6 +571,36 @@ public TSStatus visitDropModel(DropModelStatement statement, TreeAccessCheckCont
572571
return checkModelManagement(context.getUsername());
573572
}
574573

574+
@Override
575+
public TSStatus visitCreateTraining(
576+
CreateTrainingStatement createTrainingStatement, TreeAccessCheckContext context) {
577+
return checkModelManagement(context.getUsername());
578+
}
579+
580+
@Override
581+
public TSStatus visitUnloadModel(
582+
UnloadModelStatement unloadModelStatement, TreeAccessCheckContext context) {
583+
return checkModelManagement(context.getUsername());
584+
}
585+
586+
@Override
587+
public TSStatus visitLoadModel(
588+
LoadModelStatement loadModelStatement, TreeAccessCheckContext context) {
589+
return checkModelManagement(context.getUsername());
590+
}
591+
592+
@Override
593+
public TSStatus visitShowAIDevices(
594+
ShowAIDevicesStatement showAIDevicesStatement, TreeAccessCheckContext context) {
595+
return checkModelManagement(context.getUsername());
596+
}
597+
598+
@Override
599+
public TSStatus visitShowLoadedModels(
600+
ShowLoadedModelsStatement showLoadedModelsStatement, TreeAccessCheckContext context) {
601+
return SUCCEED;
602+
}
603+
575604
@Override
576605
public TSStatus visitShowModels(ShowModelsStatement statement, TreeAccessCheckContext context) {
577606
return SUCCEED;
@@ -691,10 +720,7 @@ private TSStatus checkTriggerManagement(String userName) {
691720
if (AuthorityChecker.SUPER_USER.equals(userName)) {
692721
return SUCCEED;
693722
}
694-
return AuthorityChecker.getTSStatus(
695-
AuthorityChecker.checkSystemPermission(userName, PrivilegeType.SYSTEM)
696-
|| AuthorityChecker.checkSystemPermission(userName, PrivilegeType.USE_TRIGGER),
697-
PrivilegeType.SYSTEM);
723+
return checkGlobalAuth(userName, PrivilegeType.USE_TRIGGER);
698724
}
699725

700726
// ============================== database related ===========================
@@ -744,11 +770,7 @@ public TSStatus visitDeleteStorageGroup(
744770
if (AuthorityChecker.SUPER_USER.equals(context.getUsername())) {
745771
return SUCCEED;
746772
}
747-
return AuthorityChecker.getTSStatus(
748-
AuthorityChecker.checkSystemPermission(context.getUsername(), PrivilegeType.SYSTEM)
749-
|| AuthorityChecker.checkSystemPermission(
750-
context.getUsername(), PrivilegeType.MANAGE_DATABASE),
751-
PrivilegeType.SYSTEM);
773+
return checkGlobalAuth(context.getUsername(), PrivilegeType.MANAGE_DATABASE);
752774
}
753775

754776
private TSStatus checkCreateOrAlterDatabasePermission(String userName, PartialPath databaseName) {
@@ -762,10 +784,7 @@ private TSStatus checkCreateOrAlterDatabasePermission(String userName, PartialPa
762784
return SUCCEED;
763785
}
764786

765-
return AuthorityChecker.getTSStatus(
766-
AuthorityChecker.checkSystemPermission(userName, PrivilegeType.SYSTEM)
767-
|| AuthorityChecker.checkSystemPermission(userName, PrivilegeType.MANAGE_DATABASE),
768-
PrivilegeType.SYSTEM);
787+
return checkGlobalAuth(userName, PrivilegeType.MANAGE_DATABASE);
769788
}
770789

771790
private TSStatus checkShowOrCountDatabasePermission(
@@ -855,6 +874,11 @@ public TSStatus visitExplainAnalyze(
855874
return statement.getQueryStatement().accept(this, context);
856875
}
857876

877+
@Override
878+
public TSStatus visitExplain(ExplainStatement explainStatement, TreeAccessCheckContext context) {
879+
return explainStatement.getQueryStatement().accept(this, context);
880+
}
881+
858882
// ============================= timeseries related =================================
859883
public static TSStatus checkTimeSeriesPermission(
860884
String userName, List<? extends PartialPath> checkedPaths, PrivilegeType permission) {
@@ -1280,6 +1304,12 @@ public TSStatus visitShowCurrentTimestamp(
12801304
return SUCCEED;
12811305
}
12821306

1307+
@Override
1308+
public TSStatus visitLoadConfiguration(
1309+
LoadConfigurationStatement loadConfigurationStatement, TreeAccessCheckContext context) {
1310+
return checkOnlySuperUser(context.getUsername());
1311+
}
1312+
12831313
// ======================== TTL related ===========================
12841314
@Override
12851315
public TSStatus visitSetTTL(SetTTLStatement statement, TreeAccessCheckContext context) {
@@ -1429,4 +1459,11 @@ protected void setCanSeeAuditDB(AuthorityInformationStatement statement, String
14291459
statement.setCanSeeAuditDB(false);
14301460
}
14311461
}
1462+
1463+
private TSStatus checkOnlySuperUser(String userName) {
1464+
if (AuthorityChecker.SUPER_USER.equals(userName)) {
1465+
return SUCCEED;
1466+
}
1467+
return AuthorityChecker.getTSStatus(false, "Only the admin user can perform this operation");
1468+
}
14321469
}

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/auth/entity/User.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,12 @@
3636
/** This class contains all information of a User. */
3737
public class User extends Role {
3838

39+
public static final long INTERNAL_USER_END_ID = 9999;
40+
41+
public static final long INTERNAL_SYSTEM_ADMIN = 1;
42+
public static final long INTERNAL_SECURITY_ADMIN = 2;
43+
public static final long INTERNAL_AUDIT_ADMIN = 3;
44+
3945
private long userId = -1;
4046

4147
private String password;

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/auth/user/BasicUserManager.java

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@
4040
import java.io.IOException;
4141
import java.util.Map;
4242

43+
import static org.apache.iotdb.commons.auth.entity.User.INTERNAL_USER_END_ID;
44+
4345
/** This class stores information of each user. */
4446
public abstract class BasicUserManager extends BasicRoleManager {
4547

@@ -55,7 +57,7 @@ protected String getNoSuchEntityError() {
5557
return "No such user %s";
5658
}
5759

58-
protected long nextUserId = 9999;
60+
protected long nextUserId = INTERNAL_USER_END_ID;
5961

6062
/**
6163
* BasicUserManager Constructor.
@@ -151,11 +153,7 @@ private void initInternalAuditorWhenNecessary() throws AuthException {
151153
private void initUserId() {
152154
try {
153155
long maxUserId = this.accessor.loadUserId();
154-
if (maxUserId < 9999) {
155-
nextUserId = 9999;
156-
} else {
157-
nextUserId = maxUserId;
158-
}
156+
nextUserId = Math.max(maxUserId, INTERNAL_USER_END_ID);
159157

160158
for (Map.Entry<String, Role> userEntry : entityMap.entrySet()) {
161159
User user = (User) userEntry.getValue();

0 commit comments

Comments
 (0)