优化websocket获取user信息

zhou-hao · zhou-hao · commit d6f95eecf5cd · 2017-05-05T09:53:48.000+08:00
diff --git a/hsweb-web-websocket/src/main/java/org/hsweb/web/socket/message/SimpleWebSocketMessageManager.java b/hsweb-web-websocket/src/main/java/org/hsweb/web/socket/message/SimpleWebSocketMessageManager.java
@@ -4,6 +4,7 @@
 import org.hsweb.web.core.exception.AuthorizeException;
 import org.hsweb.web.core.session.HttpSessionManager;
 import org.hsweb.web.socket.utils.SessionUtils;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketSession;
 
@@ -26,17 +27,13 @@ public class SimpleWebSocketMessageManager implements WebSocketMessageManager {
 
     private HttpSessionManager httpSessionManager;
 
+    @Autowired
     public void setHttpSessionManager(HttpSessionManager httpSessionManager) {
         this.httpSessionManager = httpSessionManager;
     }
 
     public Map<String, WebSocketSession> getSessionMap(String userId) {
-        Map<String, WebSocketSession> map = session_map.get(userId);
-        if (map == null) {
-            map = Collections.synchronizedMap(new HashMap<>());
-            session_map.put(userId, map);
-        }
-        return map;
+        return session_map.computeIfAbsent(userId, k -> Collections.synchronizedMap(new HashMap<>()));
     }
 
     @Override
diff --git a/hsweb-web-websocket/src/main/java/org/hsweb/web/socket/utils/SessionUtils.java b/hsweb-web-websocket/src/main/java/org/hsweb/web/socket/utils/SessionUtils.java
@@ -5,15 +5,15 @@
 import org.springframework.http.HttpHeaders;
 import org.springframework.web.socket.WebSocketSession;
 
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
+import java.util.function.Function;
 
 /**
  * Created by zhouhao on 16-5-30.
  */
 public class SessionUtils {
     public static User getUser(WebSocketSession session, HttpSessionManager sessionManager) {
+        if (sessionManager == null) return null;
         User user = ((User) session.getAttributes().get("user"));
         if (user != null) return user;
         HttpHeaders headers = session.getHandshakeHeaders();
@@ -22,16 +22,21 @@ public static User getUser(WebSocketSession session, HttpSessionManager sessionM
             return null;
         }
         String[] cookie = cookies.get(0).split("[;]");
-        Map<String, Object> sessionId = new HashMap<>();
-        for (int i = 0; i < cookie.length; i++) {
-            String[] tmp = cookie[i].split("[=]");
+        Map<String, Set<String>> sessionId = new HashMap<>();
+        for (String aCookie : cookie) {
+            String[] tmp = aCookie.split("[=]");
             if (tmp.length == 2)
-                sessionId.put(tmp[0].trim(), tmp[1].trim());
+                sessionId.computeIfAbsent(tmp[0].trim(), k -> new HashSet<>())
+                        .add(tmp[1].trim());
         }
-        user = sessionManager.getUserBySessionId((String) sessionId.get("SESSION"));
-        if (user == null) {
-            user = sessionManager.getUserBySessionId((String) sessionId.get("JSESSIONID"));
-        }
-        return user;
+
+        Function<Set<String>, Optional<User>> userGetter = set ->
+                set == null ? Optional.empty() : set.stream()
+                        .map(sessionManager::getUserBySessionId)
+                        .filter(Objects::nonNull).findFirst();
+
+        return userGetter.apply(sessionId.get("SESSION"))
+                .orElseGet(() -> userGetter.apply(sessionId.get("JSESSIONID")).orElse(null));
+
     }
 }
