hhyykk/ipms-sjy
1
1
Fork 0
mirror of https://gitee.com/hhyykk/ipms-sjy.git synced 2025-07-13 10:35:07 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

websocket：重新封装 websocket 组件，支持 token 认证，并增加 WebSocketMessageListener 方便处理消息

Browse Source
This commit is contained in:
YunaiV
2023-11-22 23:17:50 +08:00
parent 6a61db8508
commit 1941a7b3e6
27 changed files with 539 additions and 540 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
yudao-framework/yudao-spring-boot-starter-websocket/pom.xml
View File

@ -12,18 +12,22 @@
<packaging>jar</packaging>
<name>${project.artifactId}</name>
<description>WebSocket</description>
<description>WebSocket 框架，支持多节点的广播</description>
<url>https://github.com/YunaiV/ruoyi-vue-pro</url>
<dependencies>
<dependency>
<groupId>cn.iocoder.boot</groupId>
<artifactId>yudao-common</artifactId>
</dependency>
<!-- Web 相关 -->
<dependency>
<!-- 为什么是 websocket 依赖 security 呢？而不是 security 拓展 websocket 呢？
因为 websocket 和 LoginUser 当前登录的用户有一定的相关性，具体可见 WebSocketSessionManagerImpl 逻辑。
如果让 security 拓展 websocket 的话，会导致 websocket 组件的封装很散，进而增大理解成本。
-->
<groupId>cn.iocoder.boot</groupId>
<artifactId>yudao-spring-boot-starter-security</artifactId>
</dependency>

14
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/config/WebSocketHandlerConfig.java
View File

@ -1,14 +0,0 @@
package cn.iocoder.yudao.framework.websocket.config;
import cn.iocoder.yudao.framework.websocket.core.UserHandshakeInterceptor;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.web.socket.server.HandshakeInterceptor;
@EnableConfigurationProperties(WebSocketProperties.class)
public class WebSocketHandlerConfig {
@Bean
public HandshakeInterceptor handshakeInterceptor() {
return new UserHandshakeInterceptor();
}
}

13
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/config/WebSocketProperties.java
View File

@ -15,15 +15,8 @@ import org.springframework.validation.annotation.Validated;
public class WebSocketProperties {
/**
* 路径
* WebSocket 的连接路径
*/
private String path = "";
/**
* 默认最多允许同时在线用户数
*/
private int maxOnlineCount = 0;
/**
* 是否保存session
*/
private boolean sessionMap = true;
private String path = "/ws";
}

45
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/config/YudaoWebSocketAutoConfiguration.java
View File

@ -1,11 +1,17 @@
package cn.iocoder.yudao.framework.websocket.config;
import cn.iocoder.yudao.framework.websocket.core.handler.JsonWebSocketMessageHandler;
import cn.iocoder.yudao.framework.websocket.core.listener.WebSocketMessageListener;
import cn.iocoder.yudao.framework.websocket.core.security.LoginUserHandshakeInterceptor;
import cn.iocoder.yudao.framework.websocket.core.session.WebSocketSessionHandlerDecorator;
import cn.iocoder.yudao.framework.websocket.core.session.WebSocketSessionManager;
import cn.iocoder.yudao.framework.websocket.core.session.WebSocketSessionManagerImpl;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.server.HandshakeInterceptor;
@ -17,18 +23,41 @@ import java.util.List;
* @author xingyu4j
*/
@AutoConfiguration
// 允许使用 yudao.websocket.enable=false 禁用websocket
@ConditionalOnProperty(prefix = "yudao.websocket", value = "enable", matchIfMissing = true)
@EnableWebSocket // 开启 websocket
@ConditionalOnProperty(prefix = "yudao.websocket", value = "enable", matchIfMissing = true) // 允许使用 yudao.websocket.enable=false 禁用 websocket
@EnableConfigurationProperties(WebSocketProperties.class)
public class YudaoWebSocketAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public WebSocketConfigurer webSocketConfigurer(List<HandshakeInterceptor> handshakeInterceptor,
public WebSocketConfigurer webSocketConfigurer(HandshakeInterceptor[] handshakeInterceptors,
WebSocketHandler webSocketHandler,
WebSocketProperties webSocketProperties) {
return registry -> registry
// 添加 WebSocketHandler
.addHandler(webSocketHandler, webSocketProperties.getPath())
.addInterceptors(handshakeInterceptor.toArray(new HandshakeInterceptor[0]));
.addInterceptors(handshakeInterceptors)
// 允许跨域，否则前端连接会直接断开
.setAllowedOriginPatterns("*");
}
}
@Bean
public HandshakeInterceptor handshakeInterceptor() {
return new LoginUserHandshakeInterceptor();
}
@Bean
public WebSocketHandler webSocketHandler(WebSocketSessionManager sessionManager,
List<? extends WebSocketMessageListener<?>> messageListeners) {
// 1. 创建 JsonWebSocketMessageHandler 对象，处理消息
JsonWebSocketMessageHandler messageHandler = new JsonWebSocketMessageHandler(messageListeners);
// 2. 创建 WebSocketSessionHandlerDecorator 对象，处理连接
return new WebSocketSessionHandlerDecorator(messageHandler, sessionManager);
}
@Bean
public WebSocketSessionManager webSocketSessionManager() {
return new WebSocketSessionManagerImpl();
}
}

24
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/UserHandshakeInterceptor.java
View File

@ -1,24 +0,0 @@
package cn.iocoder.yudao.framework.websocket.core;
import cn.iocoder.yudao.framework.security.core.LoginUser;
import cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import java.util.Map;
public class UserHandshakeInterceptor implements HandshakeInterceptor {
@Override
public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
LoginUser loginUser = SecurityFrameworkUtils.getLoginUser();
attributes.put(WebSocketKeyDefine.LOGIN_USER, loginUser);
return true;
}
@Override
public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
}
}

9
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/WebSocketKeyDefine.java
View File

@ -1,9 +0,0 @@
package cn.iocoder.yudao.framework.websocket.core;
import lombok.Data;
@Data
public class WebSocketKeyDefine {
public static final String LOGIN_USER ="LOGIN_USER";
}

24
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/WebSocketMessageDO.java
View File

@ -1,24 +0,0 @@
package cn.iocoder.yudao.framework.websocket.core;
import lombok.Data;
import lombok.experimental.Accessors;
import java.util.List;
@Data
@Accessors(chain = true)
public class WebSocketMessageDO {
/**
* 接收消息的seesion
*/
private List<Object> seesionKeyList;
/**
* 发送消息
*/
private String msgText;
public static WebSocketMessageDO build(List<Object> seesionKeyList, String msgText) {
return new WebSocketMessageDO().setMsgText(msgText).setSeesionKeyList(seesionKeyList);
}
}

36
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/WebSocketSessionHandler.java
View File

@ -1,36 +0,0 @@
package cn.iocoder.yudao.framework.websocket.core;
import org.springframework.web.socket.WebSocketSession;
import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
public final class WebSocketSessionHandler {
private WebSocketSessionHandler() {
}
private static final Map<String, WebSocketSession> SESSION_MAP = new ConcurrentHashMap<>();
public static void addSession(Object sessionKey, WebSocketSession session) {
SESSION_MAP.put(sessionKey.toString(), session);
}
public static void removeSession(Object sessionKey) {
SESSION_MAP.remove(sessionKey.toString());
}
public static WebSocketSession getSession(Object sessionKey) {
return SESSION_MAP.get(sessionKey.toString());
}
public static Collection<WebSocketSession> getSessions() {
return SESSION_MAP.values();
}
public static Set<String> getSessionKeys() {
return SESSION_MAP.keySet();
}
}

31
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/WebSocketUtils.java
View File

@ -1,31 +0,0 @@
package cn.iocoder.yudao.framework.websocket.core;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import java.io.IOException;
@Slf4j
public class WebSocketUtils {
public static boolean sendMessage(WebSocketSession seesion, String message) {
if (seesion == null) {
log.error("seesion 不存在");
return false;
}
if (seesion.isOpen()) {
try {
seesion.sendMessage(new TextMessage(message));
} catch (IOException e) {
log.error("WebSocket 消息发送异常 Session={} | msg= {} | exception={}", seesion, message, e);
return false;
}
}
return true;
}
public static boolean sendMessage(Object sessionKey, String message) {
WebSocketSession session = WebSocketSessionHandler.getSession(sessionKey);
return sendMessage(session, message);
}
}

49
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/YudaoWebSocketHandlerDecorator.java
View File

@ -1,49 +0,0 @@
package cn.iocoder.yudao.framework.websocket.core;
import cn.iocoder.yudao.framework.security.core.LoginUser;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
public class YudaoWebSocketHandlerDecorator extends WebSocketHandlerDecorator {
public YudaoWebSocketHandlerDecorator(WebSocketHandler delegate) {
super(delegate);
}
/**
* websocket 连接时执行的动作
* @param session websocket session 对象
* @throws Exception 异常对象
*/
@Override
public void afterConnectionEstablished(final WebSocketSession session) throws Exception {
Object sessionKey = sessionKeyGen(session);
WebSocketSessionHandler.addSession(sessionKey, session);
}
/**
* websocket 关闭连接时执行的动作
* @param session websocket session 对象
* @param closeStatus 关闭状态对象
* @throws Exception 异常对象
*/
@Override
public void afterConnectionClosed(final WebSocketSession session, CloseStatus closeStatus) throws Exception {
Object sessionKey = sessionKeyGen(session);
WebSocketSessionHandler.removeSession(sessionKey);
}
public Object sessionKeyGen(WebSocketSession webSocketSession) {
Object obj = webSocketSession.getAttributes().get(WebSocketKeyDefine.LOGIN_USER);
if (obj instanceof LoginUser) {
LoginUser loginUser = (LoginUser) obj;
// userId 作为唯一区分
return String.valueOf(loginUser.getId());
}
return null;
}
}

80
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/handler/JsonWebSocketMessageHandler.java Normal file
View File

@ -0,0 +1,80 @@
package cn.iocoder.yudao.framework.websocket.core.handler;
import cn.hutool.core.util.StrUtil;
import cn.hutool.core.util.TypeUtil;
import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import cn.iocoder.yudao.framework.websocket.core.listener.WebSocketMessageListener;
import cn.iocoder.yudao.framework.websocket.core.message.JsonWebSocketMessage;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import java.lang.reflect.Type;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Consumer;
/**
* JSON 格式 {@link WebSocketHandler} 实现类
*
* 基于 {@link JsonWebSocketMessage#getType()} 消息类型，调度到对应的 {@link WebSocketMessageListener} 监听器。
*
* @author 芋道源码
*/
@Slf4j
public class JsonWebSocketMessageHandler extends TextWebSocketHandler {
/**
* type 与 WebSocketMessageListener 的映射
*/
private final Map<String, WebSocketMessageListener<Object>> listeners = new HashMap<>();
@SuppressWarnings({"rawtypes", "unchecked"})
public JsonWebSocketMessageHandler(List<? extends WebSocketMessageListener> listenersList) {
listenersList.forEach((Consumer<WebSocketMessageListener>)
listener -> listeners.put(listener.getType(), listener));
}
@Override
protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
// 1.1 空消息，跳过
if (message.getPayloadLength() == 0) {
return;
}
// 1.2 ping 心跳消息，直接返回 pong 消息。
if (message.getPayloadLength() == 4 && Objects.equals(message.getPayload(), "ping")) {
session.sendMessage(new TextMessage("pong"));
return;
}
// 2.1 解析消息
try {
JsonWebSocketMessage jsonMessage = JsonUtils.parseObject(message.getPayload(), JsonWebSocketMessage.class);
if (jsonMessage == null) {
log.error("[handleTextMessage][session({}) message({}) 解析为空]", session.getId(), message.getPayload());
return;
}
if (StrUtil.isEmpty(jsonMessage.getType())) {
log.error("[handleTextMessage][session({}) message({}) 类型为空]", session.getId(), message.getPayload());
return;
}
// 2.2 获得对应的 WebSocketMessageListener
WebSocketMessageListener<Object> messageListener = listeners.get(jsonMessage.getType());
if (messageListener == null) {
log.error("[handleTextMessage][session({}) message({}) 监听器为空]", session.getId(), message.getPayload());
return;
}
// 2.3 处理消息
Type type = TypeUtil.getTypeArgument(messageListener.getClass(), 0);
Object messageObj = JsonUtils.parseObject(jsonMessage.getMessage(), type);
messageListener.onMessage(session, messageObj);
} catch (Throwable ex) {
log.error("[handleTextMessage][session({}) message({}) 处理异常]", session.getId(), message.getPayload());
}
}
}

31
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/listener/WebSocketMessageListener.java Normal file
View File

@ -0,0 +1,31 @@
package cn.iocoder.yudao.framework.websocket.core.listener;
import cn.iocoder.yudao.framework.websocket.core.message.JsonWebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
/**
* WebSocket 消息监听器接口
*
* 目的：前端发送消息给后端后，处理对应 {@link #getType()} 类型的消息
*
* @param <T> 泛型，消息类型
*/
public interface WebSocketMessageListener<T> {
/**
* 处理消息
*
* @param session Session
* @param message 消息
*/
void onMessage(WebSocketSession session, T message);
/**
* 获得消息类型
*
* @see JsonWebSocketMessage#getType()
* @return 消息类型
*/
String getType();
}

27
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/message/JsonWebSocketMessage.java Normal file
View File

@ -0,0 +1,27 @@
package cn.iocoder.yudao.framework.websocket.core.message;
import cn.iocoder.yudao.framework.websocket.core.listener.WebSocketMessageListener;
import lombok.Data;
/**
* JSON 格式的 WebSocket 消息帧
*
* @author 芋道源码
*/
@Data
public class JsonWebSocketMessage {
/**
* 消息类型
*
* 目的：用于分发到对应的 {@link WebSocketMessageListener} 实现类
*/
private String type;
/**
* 消息内容
*
* 要求 JSON 对象
*/
private String message;
}

42
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/security/LoginUserHandshakeInterceptor.java Normal file
View File

@ -0,0 +1,42 @@
package cn.iocoder.yudao.framework.websocket.core.security;
import cn.iocoder.yudao.framework.security.core.LoginUser;
import cn.iocoder.yudao.framework.security.core.filter.TokenAuthenticationFilter;
import cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils;
import cn.iocoder.yudao.framework.websocket.core.util.WebSocketFrameworkUtils;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.server.HandshakeInterceptor;
import java.util.Map;
/**
* 登录用户的 {@link HandshakeInterceptor} 实现类
*
* 流程如下：
* 1. 前端连接 websocket 时，会通过拼接 ?token={token} 到 ws:// 连接后，这样它可以被 {@link TokenAuthenticationFilter} 所认证通过
* 2. {@link LoginUserHandshakeInterceptor} 负责把 {@link LoginUser} 添加到 {@link WebSocketSession} 中
*
* @author 芋道源码
*/
public class LoginUserHandshakeInterceptor implements HandshakeInterceptor {
@Override
public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
WebSocketHandler wsHandler, Map<String, Object> attributes) {
LoginUser loginUser = SecurityFrameworkUtils.getLoginUser();
if (loginUser != null) {
WebSocketFrameworkUtils.setLoginUser(loginUser, attributes);
}
return true;
}
@Override
public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
WebSocketHandler wsHandler, Exception exception) {
// do nothing
}
}

24
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/security/WebSocketAuthorizeRequestsCustomizer.java Normal file
View File

@ -0,0 +1,24 @@
package cn.iocoder.yudao.framework.websocket.core.security;
import cn.iocoder.yudao.framework.security.config.AuthorizeRequestsCustomizer;
import cn.iocoder.yudao.framework.websocket.config.WebSocketProperties;
import lombok.RequiredArgsConstructor;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.config.annotation.web.configurers.ExpressionUrlAuthorizationConfigurer;
/**
* WebSocket 的权限自定义
*
* @author 芋道源码
*/
@RequiredArgsConstructor
public class WebSocketAuthorizeRequestsCustomizer extends AuthorizeRequestsCustomizer {
private final WebSocketProperties webSocketProperties;
@Override
public void customize(ExpressionUrlAuthorizationConfigurer<HttpSecurity>.ExpressionInterceptUrlRegistry registry) {
registry.antMatchers(webSocketProperties.getPath()).permitAll();
}
}

49
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/session/WebSocketSessionHandlerDecorator.java Normal file
View File

@ -0,0 +1,49 @@
package cn.iocoder.yudao.framework.websocket.core.session;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.ConcurrentWebSocketSessionDecorator;
import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
/**
* {@link WebSocketHandler} 的装饰类，实现了以下功能：
*
* 1. {@link WebSocketSession} 连接或关闭时，使用 {@link #sessionManager} 进行管理
* 2. 封装 {@link WebSocketSession} 支持并发操作
*
* @author 芋道源码
*/
public class WebSocketSessionHandlerDecorator extends WebSocketHandlerDecorator {
/**
* 发送时间的限制，单位：毫秒
*/
private static final Integer SEND_TIME_LIMIT = 1000 * 5;
/**
* 发送消息缓冲上线，单位：bytes
*/
private static final Integer BUFFER_SIZE_LIMIT = 1024 * 100;
private final WebSocketSessionManager sessionManager;
public WebSocketSessionHandlerDecorator(WebSocketHandler delegate,
WebSocketSessionManager sessionManager) {
super(delegate);
this.sessionManager = sessionManager;
}
@Override
public void afterConnectionEstablished(WebSocketSession session) {
// 实现 session 支持并发，可参考 https://blog.csdn.net/abu935009066/article/details/131218149
session = new ConcurrentWebSocketSessionDecorator(session, SEND_TIME_LIMIT, BUFFER_SIZE_LIMIT);
// 添加到 WebSocketSessionManager 中
sessionManager.addSession(session);
}
@Override
public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) {
sessionManager.removeSession(session);
}
}

53
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/session/WebSocketSessionManager.java Normal file
View File

@ -0,0 +1,53 @@
package cn.iocoder.yudao.framework.websocket.core.session;
import org.springframework.web.socket.WebSocketSession;
import java.util.Collection;
/**
* {@link WebSocketSession} 管理器的接口
*
* @author 芋道源码
*/
public interface WebSocketSessionManager {
/**
* 添加 Session
*
* @param session Session
*/
void addSession(WebSocketSession session);
/**
* 移除 Session
*
* @param session Session
*/
void removeSession(WebSocketSession session);
/**
* 获得指定编号的 Session
*
* @param id Session 编号
* @return Session
*/
WebSocketSession getSession(String id);
/**
* 获得指定用户类型的 Session 列表
*
* @param userType 用户类型
* @return Session 列表
*/
Collection<WebSocketSession> getSessionList(Integer userType);
/**
* 获得指定用户编号的 Session 列表
*
* @param userType 用户类型
* @param userId 用户编号
* @return Session 列表
*/
Collection<WebSocketSession> getSessionList(Integer userType, Long userId);
}

113
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/session/WebSocketSessionManagerImpl.java Normal file
View File

@ -0,0 +1,113 @@
package cn.iocoder.yudao.framework.websocket.core.session;
import cn.hutool.core.collection.CollUtil;
import cn.iocoder.yudao.framework.security.core.LoginUser;
import cn.iocoder.yudao.framework.websocket.core.util.WebSocketFrameworkUtils;
import org.springframework.web.socket.WebSocketSession;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CopyOnWriteArrayList;
/**
* 默认的 {@link WebSocketSessionManager} 实现类
*
* @author 芋道源码
*/
public class WebSocketSessionManagerImpl implements WebSocketSessionManager {
/**
* id 与 WebSocketSession 映射
*
* key：Session 编号
*/
private final ConcurrentMap<String, WebSocketSession> idSessions = new ConcurrentHashMap<>();
/**
* user 与 WebSocketSession 映射
*
* key1：用户类型
* key2：用户编号
*/
private final ConcurrentMap<Integer, ConcurrentMap<Long, CopyOnWriteArrayList<WebSocketSession>>> userSessions
= new ConcurrentHashMap<>();
@Override
public void addSession(WebSocketSession session) {
// 添加到 idSessions 中
idSessions.put(session.getId(), session);
// 添加到 userSessions 中
LoginUser user = WebSocketFrameworkUtils.getLoginUser(session);
if (user == null) {
return;
}
ConcurrentMap<Long, CopyOnWriteArrayList<WebSocketSession>> userSessionsMap = userSessions.get(user.getUserType());
if (userSessionsMap == null) {
userSessionsMap = new ConcurrentHashMap<>();
if (userSessions.putIfAbsent(user.getUserType(), userSessionsMap) != null) {
userSessionsMap = userSessions.get(user.getUserType());
}
}
CopyOnWriteArrayList<WebSocketSession> sessions = userSessionsMap.get(user.getId());
if (sessions == null) {
sessions = new CopyOnWriteArrayList<>();
if (userSessionsMap.putIfAbsent(user.getId(), sessions) != null) {
sessions = userSessionsMap.get(user.getId());
}
}
sessions.add(session);
}
@Override
public void removeSession(WebSocketSession session) {
// 移除从 idSessions 中
idSessions.remove(session.getId(), session);
// 移除从 idSessions 中
LoginUser user = WebSocketFrameworkUtils.getLoginUser(session);
if (user == null) {
return;
}
ConcurrentMap<Long, CopyOnWriteArrayList<WebSocketSession>> userSessionsMap = userSessions.get(user.getUserType());
if (userSessionsMap == null) {
return;
}
CopyOnWriteArrayList<WebSocketSession> sessions = userSessionsMap.get(user.getId());
sessions.removeIf(session0 -> session0.getId().equals(session.getId()));
if (CollUtil.isEmpty(sessions)) {
userSessionsMap.remove(user.getId(), sessions);
}
}
@Override
public WebSocketSession getSession(String id) {
return idSessions.get(id);
}
@Override
public Collection<WebSocketSession> getSessionList(Integer userType) {
ConcurrentMap<Long, CopyOnWriteArrayList<WebSocketSession>> userSessionsMap = userSessions.get(userType);
if (CollUtil.isEmpty(userSessionsMap)) {
return new ArrayList<>();
}
LinkedList<WebSocketSession> result = new LinkedList<>(); // 避免扩容
for (List<WebSocketSession> sessions : userSessionsMap.values()) {
if (CollUtil.isNotEmpty(sessions)) {
continue;
}
result.addAll(sessions);
}
return result;
}
@Override
public Collection<WebSocketSession> getSessionList(Integer userType, Long userId) {
ConcurrentMap<Long, CopyOnWriteArrayList<WebSocketSession>> userSessionsMap = userSessions.get(userType);
if (CollUtil.isEmpty(userSessionsMap)) {
return new ArrayList<>();
}
CopyOnWriteArrayList<WebSocketSession> sessions = userSessionsMap.get(userId);
return CollUtil.isNotEmpty(sessions) ? new ArrayList<>(sessions) : new ArrayList<>();
}
}

59
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/core/util/WebSocketFrameworkUtils.java Normal file
View File

@ -0,0 +1,59 @@
package cn.iocoder.yudao.framework.websocket.core.util;
import cn.iocoder.yudao.framework.security.core.LoginUser;
import org.springframework.lang.Nullable;
import org.springframework.web.socket.WebSocketSession;
import java.util.Map;
/**
* 专属于 web 包的工具类
*
* @author 芋道源码
*/
public class WebSocketFrameworkUtils {
public static final String ATTRIBUTE_LOGIN_USER = "LOGIN_USER";
/**
* 设置当前用户
*
* @param loginUser 登录用户
* @param attributes Session
*/
public static void setLoginUser(LoginUser loginUser, Map<String, Object> attributes) {
attributes.put(ATTRIBUTE_LOGIN_USER, loginUser);
}
/**
* 获取当前用户
*
* @return 当前用户
*/
public static LoginUser getLoginUser(WebSocketSession session) {
return (LoginUser) session.getAttributes().get(ATTRIBUTE_LOGIN_USER);
}
/**
* 获得当前用户的编号
*
* @return 用户编号
*/
@Nullable
public static Long getLoginUserId(WebSocketSession session) {
LoginUser loginUser = getLoginUser(session);
return loginUser != null ? loginUser.getId() : null;
}
/**
* 获得当前用户的类型
*
* @return 用户编号
*/
@Nullable
public static Integer getLoginUserType(WebSocketSession session) {
LoginUser loginUser = getLoginUser(session);
return loginUser != null ? loginUser.getUserType() : null;
}
}

3
yudao-framework/yudao-spring-boot-starter-websocket/src/main/java/cn/iocoder/yudao/framework/websocket/package-info.java
View File

@ -1 +1,4 @@
/**
* WebSocket 框架，支持多节点的广播
*/
package cn.iocoder.yudao.framework.websocket;