新增字段

wujianwei

2025-08-15 8776c82342ede39fe894dec9a5e0ace216d1c99d

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
package com.ruoyi.framework.socket;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
 
import com.alibaba.fastjson2.JSON;
import com.ruoyi.common.core.domain.entity.SysDept;
import com.ruoyi.common.core.domain.model.LoginUser;
import com.ruoyi.framework.web.domain.Message;
import com.ruoyi.framework.web.service.TokenService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;
 
import java.io.IOException;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
 
@Component
public class MyWebSocketHandler extends TextWebSocketHandler {
    private static final ObjectMapper objectMapper = new ObjectMapper();
 
    // 存储已连接的用户
    private static final ConcurrentHashMap<Long, WebSocketSession> clients = new ConcurrentHashMap<>();
    @Autowired
    TokenService tokenService;
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String token = getToken(session);
        LoginUser loginUser = tokenService.getLoginUser(token);
        if (loginUser != null) {
            clients.put(loginUser.getUserId(), session);
            sendToUser(loginUser.getUserId(), new Message("Connection Successful"));
        }else{
            session.close(CloseStatus.SERVER_ERROR);
        }
 
 
    }
 
 
    // **解析 Token**
    private String getToken(WebSocketSession session) {
        String query = Objects.requireNonNull(session.getUri()).getQuery();
        if (query == null || !query.contains("token=")) return null;
        return query.split("token=")[1];
    }
 
    /**
     * 获取用户
     * @param session
     * @return
     */
    private Long getUserId(WebSocketSession session) {
        for (Map.Entry<Long, WebSocketSession> entry :  clients.entrySet()) {
            if (entry.getValue().equals(session)) return entry.getKey();
        }
        return null;
    }
 
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        if (message.getPayload().equals("{\"type\":\"ping\"}")) {
            sendToUser(getUserId(session), new Message("pong"));
        }
 
        super.handleTextMessage(session, message);
    }
 
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        clients.values().removeIf(c -> c == session);
    }
 
 
    /**
     * 发送给指定用户
     * @param userId 用户ID
     * @param message 消息内容
     * @throws IOException
     */
    private void sendToUser(Long userId, Message message) throws IOException {
        WebSocketSession webSocketSession = clients.get(userId);
        if (webSocketSession != null && webSocketSession.isOpen()) {
            webSocketSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(message)));
        }
    }
 
    /**
     * 发送给所有人
     * @param message 消息内容
     * @throws IOException
     */
    private void sendToAll(Message message) throws IOException {
        for (WebSocketSession webSocketSession : clients.values()) {
            if (webSocketSession.isOpen()) {
                webSocketSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(message)));
            }
        }
    }
 
}