CB-2713 license role management

2026-10-12 23:37:00 +08:00 · 2022-10-16 12:15:11 +02:00
parent db1b06a286
commit f1da5369f2
6 changed files with 137 additions and 26 deletions
@@ -88,4 +88,8 @@ public class WebUser {
    public String toString() {
        return user.getUserId();
    }
+
+    public String getAuthRole() {
+        return user.getAuthRole();
+    }
 }
@@ -42,6 +42,7 @@ type AdminUserInfo {

    linkedAuthProviders: [String!]!
    enabled: Boolean!
+    authRole: String
 }

 type AdminTeamInfo {
@@ -54,6 +54,11 @@ public class AdminUserInfo {
        return user.getUserId();
    }

+    @Property
+    public String getAuthRole() {
+        return user.getAuthRole();
+    }
+
    @Property
    public Map<String, String> getMetaParameters() {
        return user.getMetaParameters();
@@ -27,7 +27,7 @@ import io.cloudbeaver.model.app.WebAuthConfiguration;
 import io.cloudbeaver.model.session.WebAuthInfo;
 import io.cloudbeaver.service.security.db.CBDatabase;
 import io.cloudbeaver.service.security.internal.AuthAttemptSessionInfo;
-import io.cloudbeaver.service.security.internal.RefreshTokenInfo;
+import io.cloudbeaver.service.security.internal.SMTokenInfo;
 import io.cloudbeaver.utils.WebAppUtils;
 import org.jkiss.code.NotNull;
 import org.jkiss.code.Nullable;
@@ -126,7 +126,7 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
            try (JDBCTransaction txn = new JDBCTransaction(dbCon)) {
                createAuthSubject(dbCon, userId, SUBJECT_USER);
                try (PreparedStatement dbStat = dbCon.prepareStatement(
-                    "INSERT INTO CB_USER(USER_ID,IS_ACTIVE,CREATE_TIME,DEFAULT_AUTH_ROLE) VALUES(?,?,?)")
+                    "INSERT INTO CB_USER(USER_ID,IS_ACTIVE,CREATE_TIME,DEFAULT_AUTH_ROLE) VALUES(?,?,?,?)")
                ) {
                    dbStat.setString(1, userId);
                    dbStat.setString(2, enabled ? CHAR_BOOL_TRUE : CHAR_BOOL_FALSE);
@@ -216,13 +216,16 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
    public SMUser getUserById(String userId) throws DBException {
        try (Connection dbCon = database.openConnection()) {
            SMUser user;
-            try (PreparedStatement dbStat = dbCon.prepareStatement("SELECT * FROM CB_USER WHERE USER_ID=?")) {
+            try (PreparedStatement dbStat = dbCon.prepareStatement(
+                "SELECT USER_ID,IS_ACTIVE,DEFAULT_AUTH_ROLE FROM CB_USER WHERE USER_ID=?"
+            )) {
                dbStat.setString(1, userId);
                try (ResultSet dbResult = dbStat.executeQuery()) {
                    if (dbResult.next()) {
                        String userName = dbResult.getString(1);
                        String active = dbResult.getString(2);
-                        user = new SMUser(userName, CHAR_BOOL_TRUE.equals(active));
+                        String authRole = dbResult.getString(3);
+                        user = new SMUser(userName, CHAR_BOOL_TRUE.equals(active), authRole);
                    } else {
                        return null;
                    }
@@ -263,8 +266,9 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
        try (Connection dbCon = database.openConnection()) {
            Map<String, SMUser> result = new LinkedHashMap<>();
            // Read users
-            try (PreparedStatement dbStat = dbCon.prepareStatement("SELECT * FROM CB_USER" +
-                (CommonUtils.isEmpty(userNameMask) ? "\nORDER BY USER_ID" : " WHERE USER_ID=?"))) {
+            try (PreparedStatement dbStat = dbCon.prepareStatement(
+                "SELECT USER_ID,IS_ACTIVE,DEFAULT_AUTH_ROLE FROM CB_USER" +
+                    (CommonUtils.isEmpty(userNameMask) ? "\nORDER BY USER_ID" : " WHERE USER_ID=?"))) {
                if (!CommonUtils.isEmpty(userNameMask)) {
                    dbStat.setString(1, userNameMask);
                }
@@ -272,7 +276,8 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
                    while (dbResult.next()) {
                        String userId = dbResult.getString(1);
                        String active = dbResult.getString(2);
-                        result.put(userId, new SMUser(userId, CHAR_BOOL_TRUE.equals(active)));
+                        String authRole = dbResult.getString(3);
+                        result.put(userId, new SMUser(userId, CHAR_BOOL_TRUE.equals(active), authRole));
                    }
                }
            }
@@ -415,6 +420,19 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
        }
    }

+    @Override
+    public void setUserAuthRole(@NotNull String userId, @Nullable String authRole) throws DBException {
+        try (Connection dbCon = database.openConnection()) {
+            try (PreparedStatement dbStat = dbCon.prepareStatement("UPDATE CB_USER SET DEFAULT_AUTH_ROLE=? WHERE USER_ID=?")) {
+                dbStat.setString(1, authRole);
+                dbStat.setString(2, userId);
+                dbStat.executeUpdate();
+            }
+        } catch (SQLException e) {
+            throw new DBCException("Error while updating user authentication role", e);
+        }
+    }
+
    ///////////////////////////////////////////
    // Credentials

@@ -1172,17 +1190,17 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
    public SMTokens refreshSession(@NotNull String refreshToken) throws DBException {
        var currentUserCreds = getCurrentUserCreds();
        var currentUserAccessToken = currentUserCreds.getSmToken();
-        String currentUserAuthRole = null; // FIXME: read role from auth token
+        String currentUserAuthRole = readTokenAuthRole(currentUserAccessToken);

-        var expectedRefreshTokenInfo = findRefreshToken(currentUserAccessToken);
+        var smTokenInfo = readAccessTokenInfo(currentUserAccessToken);

-        if (!expectedRefreshTokenInfo.getRefreshToken().equals(refreshToken)) {
+        if (!smTokenInfo.getRefreshToken().equals(refreshToken)) {
            throw new SMException("Invalid refresh token");
        }

        try (var dbCon = database.openConnection()) {
            invalidateUserTokens(currentUserAccessToken);
-            return generateNewSessionToken(expectedRefreshTokenInfo.getSessionId(), expectedRefreshTokenInfo.getUserId(), currentUserAuthRole, dbCon);
+            return generateNewSessionToken(smTokenInfo.getSessionId(), smTokenInfo.getUserId(), currentUserAuthRole, dbCon);
        } catch (SQLException e) {
            throw new DBException("Error refreshing sm session", e);
        }
@@ -1220,10 +1238,10 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
        }
    }

-    private RefreshTokenInfo findRefreshToken(String smAccessToken) throws DBException {
+    private SMTokenInfo readAccessTokenInfo(String smAccessToken) throws DBException {
        try (Connection dbCon = database.openConnection();
             PreparedStatement dbStat = dbCon.prepareStatement(
-                 "SELECT REFRESH_TOKEN_ID,SESSION_ID,USER_ID,REFRESH_TOKEN_EXPIRATION_TIME FROM CB_AUTH_TOKEN WHERE TOKEN_ID=?"
+                 "SELECT REFRESH_TOKEN_ID,SESSION_ID,USER_ID,REFRESH_TOKEN_EXPIRATION_TIME,AUTH_ROLE FROM CB_AUTH_TOKEN WHERE TOKEN_ID=?"
             )
        ) {
            dbStat.setString(1, smAccessToken);
@@ -1235,10 +1253,11 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
                var sessionId = dbResult.getString(2);
                var userId = dbResult.getString(3);
                var expiredDate = dbResult.getTimestamp(4);
+                var authRole = dbResult.getString(5);
                if (Timestamp.from(Instant.now()).after(expiredDate)) {
                    throw new SMRefreshTokenExpiredException("Refresh token expired");
                }
-                return new RefreshTokenInfo(refreshToken, sessionId, userId);
+                return new SMTokenInfo(smAccessToken, refreshToken, sessionId, userId, authRole);
            }
        } catch (SQLException e) {
            throw new DBCException("Error reading token info in database", e);
@@ -1287,7 +1306,8 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
                userAuthData,
                finishAuthMonitor,
                activeUserId,
-                activeUserId == null
+                activeUserId == null,
+                authInfo.getAuthRole()
            );

            if (userIdFromCreds == null) {
@@ -1304,6 +1324,8 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
            );
        }

+        String tokenAuthRole = updateUserAuthRoleIfNeeded(activeUserId, authInfo.getAuthRole());
+
        if (smTokens == null && permissions == null) {
            try (Connection dbCon = database.openConnection()) {
                try (JDBCTransaction txn = new JDBCTransaction(dbCon)) {
@@ -1319,9 +1341,10 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
                    } else {
                        smSessionId = authAttemptSessionInfo.getSmSessionId();
                    }
-                    smTokens = generateNewSessionToken(smSessionId, activeUserId, authInfo.getAuthRole(), dbCon);
+
+                    smTokens = generateNewSessionToken(smSessionId, activeUserId, tokenAuthRole, dbCon);
                    permissions = new SMAuthPermissions(
-                        activeUserId, smSessionId, getUserPermissions(activeUserId, authInfo.getAuthRole()));
+                        activeUserId, smSessionId, getUserPermissions(activeUserId, tokenAuthRole));
                    txn.commit();
                }
            } catch (SQLException e) {
@@ -1342,6 +1365,40 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
        );
    }

+    @Nullable
+    private String readUserAuthRole(String userId) throws DBException {
+        try (Connection dbCon = database.openConnection()) {
+            try (PreparedStatement dbStat = dbCon.prepareStatement(
+                "SELECT DEFAULT_AUTH_ROLE FROM CB_USER WHERE USER_ID=?"
+            )) {
+                dbStat.setString(1, userId);
+                try (ResultSet dbResult = dbStat.executeQuery()) {
+                    if (!dbResult.next()) {
+                        throw new SMException("User not found");
+                    }
+                    return dbResult.getString(1);
+
+                }
+            }
+        } catch (SQLException e) {
+            throw new DBException("Failed to update user auth role", e);
+        }
+    }
+
+    private String updateUserAuthRoleIfNeeded(@Nullable String userId, @Nullable String authRole) throws DBException {
+        if (userId == null) {
+            return null;
+        }
+        var currentAuthRole = readUserAuthRole(userId);
+        if (authRole == null) {
+            return currentAuthRole;
+        }
+        if (!Objects.equals(currentAuthRole, authRole)) {
+            setUserAuthRole(userId, authRole);
+        }
+        return authRole;
+    }
+
    private AuthAttemptSessionInfo readAuthAttemptSessionInfo(@NotNull String authId) throws DBException {
        try (Connection dbCon = database.openConnection()) {
            try (PreparedStatement dbStat = dbCon.prepareStatement(
@@ -1373,7 +1430,8 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
        @NotNull Map<String, Object> userCredentials,
        @NotNull DBRProgressMonitor progressMonitor,
        @Nullable String activeUserId,
-        boolean createNewUserIfNotExist
+        boolean createNewUserIfNotExist,
+        String authRole
    ) throws DBException {
        AuthProviderDescriptor authProvider = getAuthProvider(authProviderId);
        SMAuthProvider<?> smAuthProviderInstance = authProvider.getInstance();
@@ -1397,11 +1455,10 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen

            userId = userIdFromCredentials;
            if (!isSubjectExists(userId)) {
-                var newUser = new SMUser(userId, true);
-                createUser(newUser.getUserId(),
-                    newUser.getMetaParameters(),
+                createUser(userId,
+                    Map.of(),
                    true,
-                    resolveUserAuthRole(authProvider, userCredentials)
+                    authRole
                );
                String defaultTeamName = WebAppUtils.getWebApplication().getAppConfiguration().getDefaultUserTeam();
                if (!CommonUtils.isEmpty(defaultTeamName)) {
@@ -1426,7 +1483,7 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
        @NotNull AuthProviderDescriptor authProvider,
        @NotNull Map<String, Object> userAuthData
    ) {
-        //TODO implement logic
+        //TODO read from external auth somehow
        return null;
    }

@@ -1801,6 +1858,25 @@ public class CBEmbeddedSecurityController implements SMAdminController, SMAuthen
    ///////////////////////////////////////////
    // Utils

+    protected String readTokenAuthRole(String smAccessToken) throws DBException {
+        try (Connection dbCon = database.openConnection();
+             PreparedStatement dbStat = dbCon.prepareStatement(
+                 "SELECT AUTH_ROLE FROM CB_AUTH_TOKEN " +
+                     "WHERE TOKEN_ID=?"
+             )
+        ) {
+            dbStat.setString(1, smAccessToken);
+            try (var dbResult = dbStat.executeQuery()) {
+                if (!dbResult.next()) {
+                    throw new SMException("Invalid token");
+                }
+                return dbResult.getString(1);
+            }
+        } catch (SQLException e) {
+            throw new DBCException("Error reading lm role in database", e);
+        }
+    }
+
    public void initializeMetaInformation() throws DBCException {
        try (Connection dbCon = database.openConnection()) {
            try (JDBCTransaction txn = new JDBCTransaction(dbCon)) {
@@ -255,7 +255,7 @@ public class CBDatabase {
        SMUser adminUser = adminSecurityController.getUserById(adminName);

        if (adminUser == null) {
-            adminUser = new SMUser(adminName, true);
+            adminUser = new SMUser(adminName, true, adminAuthRole);
            adminSecurityController.createUser(adminUser.getUserId(), adminUser.getMetaParameters(), true, adminAuthRole);
        }

@@ -18,8 +18,12 @@
 package io.cloudbeaver.service.security.internal;

 import org.jkiss.code.NotNull;
+import org.jkiss.code.Nullable;
+
+public class SMTokenInfo {
+    @NotNull
+    private final String accessToken;

-public class RefreshTokenInfo {
    @NotNull
    private final String refreshToken;

@@ -29,10 +33,21 @@ public class RefreshTokenInfo {
    @NotNull
    private final String userId;

-    public RefreshTokenInfo(@NotNull String refreshToken, @NotNull String sessionId, @NotNull String userId) {
+    @Nullable
+    private final String authRole;
+
+    public SMTokenInfo(
+        @NotNull String accessToken,
+        @NotNull String refreshToken,
+        @NotNull String sessionId,
+        @NotNull String userId,
+        @Nullable String authRole
+    ) {
+        this.accessToken = accessToken;
        this.refreshToken = refreshToken;
        this.sessionId = sessionId;
        this.userId = userId;
+        this.authRole = authRole;
    }

    @NotNull
@@ -49,4 +64,14 @@ public class RefreshTokenInfo {
    public String getUserId() {
        return userId;
    }
+
+    @NotNull
+    public String getAccessToken() {
+        return accessToken;
+    }
+
+    @Nullable
+    public String getAuthRole() {
+        return authRole;
+    }
 }