Connection limit function. (#16462)

Author: wenyanshi-123

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanType.java

@@ -131,6 +131,8 @@ public enum ConfigPhysicalPlanType {
   @Deprecated
   ListRoleUsers((short) 637),
   CreateUserWithRawPassword((short) 638),
+  UpdateUserMaxSession((short) 639),
+  UpdateUserMinSession((short) 640),
 
   /** Table Author */
   RCreateUser((short) 641),
@@ -164,6 +166,8 @@ public enum ConfigPhysicalPlanType {
   RListRole((short) 670),
   RListUserPrivilege((short) 671),
   RListRolePrivilege((short) 672),
+  RUpdateUserMaxSession((short) 673),
+  RUpdateUserMinSession((short) 674),
 
   /** Function. */
   CreateFunction((short) 700),

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/auth/AuthorPlan.java

@@ -31,6 +31,8 @@ public abstract class AuthorPlan extends ConfigPhysicalReadPlan {
   protected String newPassword;
   protected String userName;
   protected boolean grantOpt;
+  protected int maxSessionPerUser;
+  protected int minSessionPerUser;
 
   public AuthorPlan(final ConfigPhysicalPlanType type) {
     super(type);
@@ -44,13 +46,17 @@ public AuthorPlan(
       String roleName,
       String password,
       String newPassword,
-      boolean grantOpt) {
+      boolean grantOpt,
+      int MaxSessionPerUser,
+      int MinSessionPerUser) {
     super(type);
     this.userName = userName;
     this.roleName = roleName;
     this.password = password;
     this.newPassword = newPassword;
     this.grantOpt = grantOpt;
+    this.maxSessionPerUser = MaxSessionPerUser;
+    this.minSessionPerUser = MinSessionPerUser;
   }
 
   public ConfigPhysicalPlanType getAuthorType() {
@@ -69,6 +75,22 @@ public String getPassword() {
     return password;
   }
 
+  public int getMaxSessionPerUser() {
+    return maxSessionPerUser;
+  }
+
+  public void setMaxSessionPerUser(final int maxSessionPerUser) {
+    this.maxSessionPerUser = maxSessionPerUser;
+  }
+
+  public int getMinSessionPerUser() {
+    return minSessionPerUser;
+  }
+
+  public void setMinSessionPerUser(final int minSessionPerUser) {
+    this.maxSessionPerUser = minSessionPerUser;
+  }
+
   public void setPassword(final String password) {
     this.password = password;
   }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/auth/AuthorRelationalPlan.java

@@ -41,6 +41,32 @@ public AuthorRelationalPlan(final ConfigPhysicalPlanType authorType) {
     super(authorType);
   }
 
+  public AuthorRelationalPlan(
+      final ConfigPhysicalPlanType authorType,
+      final String userName,
+      final String roleName,
+      final String databaseName,
+      final String tableName,
+      final Set<Integer> permissions,
+      final boolean grantOpt,
+      final String password,
+      final int maxSessionPerUser,
+      final int minSessionPerUser) {
+    super(
+        authorType,
+        userName,
+        roleName,
+        password,
+        "",
+        grantOpt,
+        maxSessionPerUser,
+        minSessionPerUser);
+
+    this.databaseName = databaseName;
+    this.tableName = tableName;
+    this.permissions = permissions;
+  }
+
   public AuthorRelationalPlan(
       final ConfigPhysicalPlanType authorType,
       final String userName,
@@ -50,7 +76,8 @@ public AuthorRelationalPlan(
       final Set<Integer> permissions,
       final boolean grantOpt,
       final String password) {
-    super(authorType, userName, roleName, password, "", grantOpt);
+    super(authorType, userName, roleName, password, "", grantOpt, -1, -1);
+
     this.databaseName = databaseName;
     this.tableName = tableName;
     this.permissions = permissions;
@@ -64,7 +91,7 @@ public AuthorRelationalPlan(
       final String tableName,
       final int permission,
       final boolean grantOpt) {
-    super(authorType, userName, roleName, "", "", grantOpt);
+    super(authorType, userName, roleName, "", "", grantOpt, -1, -1);
     this.databaseName = databaseName;
     this.tableName = tableName;
     this.permissions = Collections.singleton(permission);
@@ -140,6 +167,11 @@ protected void serializeImpl(DataOutputStream stream) throws IOException {
     BasicStructureSerDeUtil.write(userName, stream);
     BasicStructureSerDeUtil.write(roleName, stream);
     BasicStructureSerDeUtil.write(password, stream);
+    if (getAuthorType() == ConfigPhysicalPlanType.UpdateUserMaxSession
+        || getAuthorType() == ConfigPhysicalPlanType.UpdateUserMinSession) {
+      BasicStructureSerDeUtil.write(maxSessionPerUser, stream);
+      BasicStructureSerDeUtil.write(minSessionPerUser, stream);
+    }
     BasicStructureSerDeUtil.write(databaseName, stream);
     BasicStructureSerDeUtil.write(tableName, stream);
     stream.writeInt(permissions.size());
@@ -155,6 +187,11 @@ protected void deserializeImpl(ByteBuffer buffer) {
     userName = BasicStructureSerDeUtil.readString(buffer);
     roleName = BasicStructureSerDeUtil.readString(buffer);
     password = BasicStructureSerDeUtil.readString(buffer);
+    if (getAuthorType() == ConfigPhysicalPlanType.UpdateUserMaxSession
+        || getAuthorType() == ConfigPhysicalPlanType.UpdateUserMinSession) {
+      maxSessionPerUser = buffer.getInt();
+      minSessionPerUser = buffer.getInt();
+    }
     databaseName = BasicStructureSerDeUtil.readString(buffer);
     tableName = BasicStructureSerDeUtil.readString(buffer);
     permissions = new HashSet<>();

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/auth/AuthorTreePlan.java

@@ -43,6 +43,42 @@ public AuthorTreePlan(final ConfigPhysicalPlanType type) {
     super(type);
   }
 
+  /**
+   * {@link AuthorTreePlan} Constructor.
+   *
+   * @param authorType author type
+   * @param userName user name
+   * @param roleName role name
+   * @param password password
+   * @param permissions permissions
+   * @param grantOpt with grant option, only grant statement can set grantOpt = true
+   * @param maxSessioPerUser maxSessionPerUser of user
+   * @param minSessionPerUser minSessionPerUser of user
+   */
+  public AuthorTreePlan(
+      final ConfigPhysicalPlanType authorType,
+      final String userName,
+      final String roleName,
+      final String password,
+      final String newPassword,
+      final Set<Integer> permissions,
+      final boolean grantOpt,
+      final List<PartialPath> nodeNameList,
+      final Integer maxSessioPerUser,
+      final Integer minSessionPerUser) {
+    super(
+        authorType,
+        userName,
+        roleName,
+        password,
+        newPassword,
+        grantOpt,
+        maxSessioPerUser,
+        minSessionPerUser);
+    this.permissions = permissions;
+    this.nodeNameList = nodeNameList;
+  }
+
   /**
    * {@link AuthorTreePlan} Constructor.
    *
@@ -63,7 +99,7 @@ public AuthorTreePlan(
       final Set<Integer> permissions,
       final boolean grantOpt,
       final List<PartialPath> nodeNameList) {
-    super(authorType, userName, roleName, password, newPassword, grantOpt);
+    super(authorType, userName, roleName, password, newPassword, grantOpt, -1, -1);
     this.permissions = permissions;
     this.nodeNameList = nodeNameList;
   }
@@ -128,6 +164,11 @@ protected void serializeImpl(DataOutputStream stream) throws IOException {
     BasicStructureSerDeUtil.write(roleName, stream);
     BasicStructureSerDeUtil.write(password, stream);
     BasicStructureSerDeUtil.write(newPassword, stream);
+    if (getAuthorType() == ConfigPhysicalPlanType.UpdateUserMaxSession
+        || getAuthorType() == ConfigPhysicalPlanType.UpdateUserMinSession) {
+      BasicStructureSerDeUtil.write(maxSessionPerUser, stream);
+      BasicStructureSerDeUtil.write(minSessionPerUser, stream);
+    }
     if (permissions == null) {
       stream.write((byte) 0);
     } else {
@@ -150,6 +191,11 @@ protected void deserializeImpl(ByteBuffer buffer) {
     roleName = BasicStructureSerDeUtil.readString(buffer);
     password = BasicStructureSerDeUtil.readString(buffer);
     newPassword = BasicStructureSerDeUtil.readString(buffer);
+    if (getAuthorType() == ConfigPhysicalPlanType.UpdateUserMaxSession
+        || getAuthorType() == ConfigPhysicalPlanType.UpdateUserMinSession) {
+      maxSessionPerUser = buffer.getInt();
+      minSessionPerUser = buffer.getInt();
+    }
     if (buffer.get() == (byte) 0) {
       this.permissions = null;
     } else {

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/AuthorInfo.java

@@ -64,7 +64,6 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -512,7 +511,8 @@ public PermissionInfoResp executeListUsers(final AuthorPlan plan) throws AuthExc
       userList = new ArrayList<>(1);
       userList.add(plan.getUserName());
       User user = authorizer.getUser(plan.getUserName());
-      userInfoList = Collections.singletonList(user.convertToListUserInfo());
+      userInfoList = new ArrayList<>(1);
+      userInfoList.add(user.convertToListUserInfo());
     } else {
       userList = authorizer.listAllUsers();
       userInfoList = authorizer.listAllUsersInfo();
@@ -536,13 +536,6 @@ public PermissionInfoResp executeListUsers(final AuthorPlan plan) throws AuthExc
         }
       }
       userInfoList.removeIf(info -> toRemove.contains(info.username));
-      final Iterator<TListUserInfo> userInfoitr = userInfoList.iterator();
-      while (itr.hasNext()) {
-        User userObj = authorizer.getUser(userInfoitr.next().getUsername());
-        if (userObj == null || !userObj.hasRole(plan.getRoleName())) {
-          itr.remove();
-        }
-      }
     }
     result.setTag(ColumnHeaderConstant.USER);
     result.setMemberInfo(userList);

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java

@@ -467,6 +467,8 @@ public TSStatus executeNonQueryPlan(ConfigPhysicalPlan physicalPlan)
       case RevokeRole:
       case RevokeRoleFromUser:
       case UpdateUser:
+      case UpdateUserMaxSession:
+      case UpdateUserMinSession:
       case CreateUserWithRawPassword:
       case CreateUserDep:
       case CreateRoleDep:
@@ -484,6 +486,8 @@ public TSStatus executeNonQueryPlan(ConfigPhysicalPlan physicalPlan)
       case RDropUser:
       case RDropRole:
       case RUpdateUser:
+      case RUpdateUserMaxSession:
+      case RUpdateUserMinSession:
       case RGrantUserRole:
       case RGrantRoleAny:
       case RGrantUserAny:

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/CNPhysicalPlanGenerator.java

@@ -199,7 +199,6 @@ private void generateUserRolePhysicalPlan(final boolean isUser) {
     try (final DataInputStream dataInputStream =
         new DataInputStream(new BufferedInputStream(inputStream))) {
       int tag = dataInputStream.readInt();
-      boolean fromOldVersion = tag < 0;
       String user;
       if (tag < 0) {
         user = readString(dataInputStream, STRING_ENCODING, strBufferLocal, -1 * tag);
@@ -221,6 +220,23 @@ private void generateUserRolePhysicalPlan(final boolean isUser) {
         createUser.setPermissions(new HashSet<>());
         createUser.setNodeNameList(new ArrayList<>());
         planDeque.add(createUser);
+        if (tag == 2) {
+          final AuthorTreePlan updateUserMaxSession =
+              new AuthorTreePlan(ConfigPhysicalPlanType.UpdateUserMaxSession);
+          updateUserMaxSession.setMaxSessionPerUser(dataInputStream.readInt());
+          updateUserMaxSession.setUserName(user);
+          updateUserMaxSession.setPermissions(new HashSet<>());
+          updateUserMaxSession.setNodeNameList(new ArrayList<>());
+          planDeque.add(updateUserMaxSession);
+          final AuthorTreePlan updateUserMinSession =
+              new AuthorTreePlan(ConfigPhysicalPlanType.UpdateUserMinSession);
+          updateUserMinSession.setMinSessionPerUser(dataInputStream.readInt());
+          updateUserMinSession.setUserName(user);
+          updateUserMinSession.setPermissions(new HashSet<>());
+          updateUserMinSession.setNodeNameList(new ArrayList<>());
+          planDeque.add(updateUserMinSession);
+        }
+
       } else {
         final AuthorTreePlan createRole = new AuthorTreePlan(ConfigPhysicalPlanType.CreateRole);
         createRole.setRoleName(user);

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/persistence/CNPhysicalPlanGeneratorTest.java

@@ -206,6 +206,29 @@ public void userGeneratorTest() throws Exception {
     plan.setNodeNameList(new ArrayList<>());
     answerSet.add(plan.hashCode());
 
+    plan = new AuthorTreePlan(ConfigPhysicalPlanType.UpdateUserMaxSession);
+    plan.setUserName(userName);
+    plan.setPermissions(new HashSet<>());
+    plan.setNodeNameList(new ArrayList<>());
+    plan.setMaxSessionPerUser(-1);
+    authorInfo.authorNonQuery(plan);
+    answerSet.add(plan.hashCode());
+
+    plan = new AuthorTreePlan(ConfigPhysicalPlanType.UpdateUserMinSession);
+    plan.setUserName(userName);
+    plan.setPermissions(new HashSet<>());
+    plan.setNodeNameList(new ArrayList<>());
+    plan.setMinSessionPerUser(-1);
+    authorInfo.authorNonQuery(plan);
+    answerSet.add(plan.hashCode());
+
+    plan = new AuthorTreePlan(ConfigPhysicalPlanType.CreateUserWithRawPassword);
+    plan.setPassword(AuthUtils.encryptPassword("password123456"));
+    plan.setUserName(userName);
+    plan.setPermissions(new HashSet<>());
+    plan.setNodeNameList(new ArrayList<>());
+    answerSet.add(plan.hashCode());
+
     plan = new AuthorTreePlan(ConfigPhysicalPlanType.CreateRole);
     plan.setRoleName("role1");
     plan.setPermissions(new HashSet<>());
@@ -271,7 +294,7 @@ public void userGeneratorTest() throws Exception {
       Assert.assertTrue(answerSet.contains(authPlan.hashCode()));
       count++;
     }
-    Assert.assertEquals(4, count);
+    Assert.assertEquals(6, count);
     final File roleListProfile =
         SystemFileFactory.INSTANCE.getFile(
             snapshotDir

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/auth/role/LocalFileRoleAccessor.java

@@ -150,6 +150,12 @@ protected void loadPrivileges(DataInputStream dataInputStream, Role role)
     role.setObjectPrivilegeMap(objectPrivilegeMap);
   }
 
+  protected void saveSessionPerUser(BufferedOutputStream outputStream, Role role)
+      throws IOException {
+    // Just used in LocalFileUserAccessor.java.
+    // Do nothing.
+  }
+
   protected void saveRoles(Role role) throws IOException {
     // Just used in LocalFileUserAccessor.java.
     // Do nothing.
@@ -191,13 +197,9 @@ public Role loadEntity(String entityName) throws IOException {
     FileInputStream inputStream = new FileInputStream(entityFile);
     try (DataInputStream dataInputStream =
         new DataInputStream(new BufferedInputStream(inputStream))) {
-      boolean fromOldVersion = false;
       int tag = dataInputStream.readInt();
-      if (tag < 0) {
-        fromOldVersion = true;
-      }
 
-      if (fromOldVersion) {
+      if (tag < 0) {
         String name =
             IOUtils.readString(dataInputStream, STRING_ENCODING, strBufferLocal, -1 * tag);
         Role role = new Role(name);
@@ -209,6 +211,11 @@ public Role loadEntity(String entityName) throws IOException {
         }
         role.setPrivilegeList(pathPrivilegeList);
         return role;
+      } else if (tag == 1) {
+        entityName = IOUtils.readString(dataInputStream, STRING_ENCODING, strBufferLocal);
+        Role role = new Role(entityName);
+        loadPrivileges(dataInputStream, role);
+        return role;
       } else {
         assert tag == VERSION;
         entityName = IOUtils.readString(dataInputStream, STRING_ENCODING, strBufferLocal);
@@ -266,6 +273,7 @@ public void saveEntity(Role entity) throws IOException {
         BufferedOutputStream outputStream = new BufferedOutputStream(fileOutputStream)) {
       saveEntityVersion(outputStream);
       saveEntityName(outputStream, entity);
+      saveSessionPerUser(outputStream, entity);
       savePrivileges(outputStream, entity);
       outputStream.flush();
       fileOutputStream.getFD().sync();