SessionUser

backend/src/main/java/io/metersphere/commons/constants/SessionConstants.java

@@ -0,0 +1,5 @@
+package io.metersphere.commons.constants;
+
+public class SessionConstants {
+    public static final String ATTR_USER = "user";
+}

backend/src/main/java/io/metersphere/dto/UserDTO.java

@@ -1,6 +1,7 @@
 package io.metersphere.dto;
 
 import io.metersphere.base.domain.Role;
+import io.metersphere.base.domain.UserRole;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -22,6 +23,10 @@ public class UserDTO {
 
     private String lastSourceId;
 
+    private List<Role> roles = new ArrayList<>();
+
+    private List<UserRole> userRoles = new ArrayList<>();
+
     public String getLastSourceId() {
         return lastSourceId;
     }
@@ -30,8 +35,6 @@ public class UserDTO {
         this.lastSourceId = lastSourceId;
     }
 
-    private List<Role> roles = new ArrayList<>();
-
     public List<Role> getRoles() {
         return roles;
     }
@@ -97,4 +100,12 @@ public class UserDTO {
     public void setUpdateTime(Long updateTime) {
         this.updateTime = updateTime;
     }
+
+    public List<UserRole> getUserRoles() {
+        return userRoles;
+    }
+
+    public void setUserRoles(List<UserRole> userRoles) {
+        this.userRoles = userRoles;
+    }
 }

backend/src/main/java/io/metersphere/security/ShiroDBRealm.java

@@ -5,7 +5,7 @@ import io.metersphere.base.domain.Role;
 import io.metersphere.dto.UserDTO;
 import io.metersphere.service.UserService;
 import io.metersphere.user.SessionUser;
-import org.apache.shiro.SecurityUtils;
+import io.metersphere.user.SessionUtils;
 import org.apache.shiro.authc.*;
 import org.apache.shiro.authz.AuthorizationInfo;
 import org.apache.shiro.authz.SimpleAuthorizationInfo;
@@ -68,7 +68,7 @@ public class ShiroDBRealm extends AuthorizingRealm {
         // TODO 密码验证
 
         SessionUser sessionUser = SessionUser.fromUser(user);
-        SecurityUtils.getSubject().getSession().setAttribute("user", sessionUser);
+        SessionUtils.putUser(sessionUser);
         return new SimpleAuthenticationInfo(userId, password, getName());
     }
 

backend/src/main/java/io/metersphere/service/UserService.java

@@ -89,6 +89,8 @@ public class UserService {
         if (CollectionUtils.isEmpty(userRoleList)) {
             return userDTO;
         }
+        // 设置 user_role
+        userDTO.setUserRoles(userRoleList);
 
         List<String> roleIds = userRoleList.stream().map(UserRole::getRoleId).collect(Collectors.toList());
 

backend/src/main/java/io/metersphere/user/SessionUser.java

@@ -1,17 +1,44 @@
 package io.metersphere.user;
 
 import io.metersphere.dto.UserDTO;
+import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.BeanUtils;
 
 import java.io.Serializable;
 
+import static io.metersphere.commons.constants.RoleConstants.*;
+
 public class SessionUser extends UserDTO implements Serializable {
 
     private static final long serialVersionUID = -7149638440406959033L;
 
+    private String workspaceId;
+    private String organizationId;
+
+    public String getWorkspaceId() {
+        return workspaceId;
+    }
+
+    public String getOrganizationId() {
+        return organizationId;
+    }
+
     public static SessionUser fromUser(UserDTO user) {
         SessionUser sessionUser = new SessionUser();
         BeanUtils.copyProperties(user, sessionUser);
+        String lastSourceId = sessionUser.getLastSourceId();
+        user.getUserRoles().forEach(ur -> {
+            if (StringUtils.equals(ur.getSourceId(), lastSourceId)) {
+                if (StringUtils.equals(ur.getRoleId(), ORG_ADMIN)) {
+                    sessionUser.organizationId = lastSourceId;
+                    return;
+                }
+                if (StringUtils.equalsAny(ur.getRoleId(), TEST_MANAGER, TEST_USER, TEST_VIEWER)) {
+                    sessionUser.workspaceId = lastSourceId;
+                }
+            }
+        });
+
         return sessionUser;
     }
 

backend/src/main/java/io/metersphere/user/SessionUtils.java

@@ -4,15 +4,32 @@ import org.apache.shiro.SecurityUtils;
 import org.apache.shiro.session.Session;
 import org.apache.shiro.subject.Subject;
 
+import java.util.Optional;
+
+import static io.metersphere.commons.constants.SessionConstants.ATTR_USER;
+
 public class SessionUtils {
 
     public static SessionUser getUser() {
         try {
             Subject subject = SecurityUtils.getSubject();
             Session session = subject.getSession();
-            return (SessionUser) session.getAttribute("user");
+            return (SessionUser) session.getAttribute(ATTR_USER);
         } catch (Exception e) {
             return null;
         }
     }
+
+    //
+    public static void putUser(SessionUser sessionUser) {
+        SecurityUtils.getSubject().getSession().setAttribute(ATTR_USER, sessionUser);
+    }
+
+    public static String getCurrentWorkspaceId() {
+        return Optional.ofNullable(getUser()).orElse(new SessionUser()).getWorkspaceId();
+    }
+
+    public static String getCurrentOrganizationId() {
+        return Optional.ofNullable(getUser()).orElse(new SessionUser()).getOrganizationId();
+    }
 }