From 6f64057e3b70af860fa015a68814717abc5704b3 Mon Sep 17 00:00:00 2001
From: shiziyuan9527 <yuhao.li@fit2cloud.com>
Date: Tue, 10 Mar 2020 18:17:52 +0800
Subject: [PATCH] =?UTF-8?q?=E9=A6=96=E6=AC=A1=E7=99=BB=E5=BD=95=E9=9A=8F?=
 =?UTF-8?q?=E6=9C=BA=E9=80=89=E4=B8=AD=E5=B7=A5=E4=BD=9C=E5=8C=BA=E9=97=B4?=
 =?UTF-8?q?=E5=92=8C=E7=BB=84=E7=BB=87?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../controller/LoginController.java           | 24 ++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/backend/src/main/java/io/metersphere/controller/LoginController.java b/backend/src/main/java/io/metersphere/controller/LoginController.java
index 11da340770..48f621dcc9 100644
--- a/backend/src/main/java/io/metersphere/controller/LoginController.java
+++ b/backend/src/main/java/io/metersphere/controller/LoginController.java
@@ -1,19 +1,27 @@
 package io.metersphere.controller;
 
+import io.metersphere.base.domain.UserRole;
 import io.metersphere.controller.request.LoginRequest;
+import io.metersphere.dto.UserDTO;
 import io.metersphere.i18n.Translator;
-import io.metersphere.user.SessionUtils;
+import io.metersphere.service.UserService;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.shiro.SecurityUtils;
 import org.apache.shiro.authc.*;
 import org.apache.shiro.authz.UnauthorizedException;
 import org.apache.shiro.subject.Subject;
 import org.springframework.web.bind.annotation.*;
+import javax.annotation.Resource;
+import java.util.List;
+import java.util.stream.Collectors;
 
 @RestController
 @RequestMapping
 public class LoginController {
 
+    @Resource
+    private UserService userService;
+
     @GetMapping(value = "/isLogin")
     public ResultHolder isLogin() {
         if (SecurityUtils.getSubject().isAuthenticated()) {
@@ -37,6 +45,20 @@ public class LoginController {
         try {
             subject.login(token);
             if (subject.isAuthenticated()) {
+                UserDTO user = (UserDTO) subject.getSession().getAttribute("user");
+                // 自动选中组织，工作空间
+                if (StringUtils.isBlank(user.getLastOrganizationId())) {
+                    List<UserRole> userRoles = user.getUserRoles();
+                    List<UserRole> test = userRoles.stream().filter(ur -> ur.getRoleId().indexOf("test") > -1).collect(Collectors.toList());
+                    List<UserRole> org = userRoles.stream().filter(ur -> ur.getRoleId().indexOf("org") > -1).collect(Collectors.toList());
+                    if (test.size() > 0) {
+                        String wsId = test.get(0).getSourceId();
+                        userService.switchUserRole(user, "workspace", wsId);
+                    } else if (org.size() > 0) {
+                        String orgId = org.get(0).getSourceId();
+                        userService.switchUserRole(user, "organization", orgId);
+                    }
+                }
                 // 返回 userDTO
                 return ResultHolder.success(subject.getSession().getAttribute("user"));
             } else {