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)));
|
}
|
}
|
}
|
|
}
|
