package cn.iocoder.yudao.framework.websocket.core.security;
|
|
import cn.iocoder.yudao.framework.common.enums.UserTypeEnum;
|
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) {
|
return false;
|
}
|
// WebSocket 地址（如 /infra/ws）不带 /admin-api、/app-api 前缀，
|
// 因此 mock 登录时 WebFrameworkUtils#getLoginUserType 取不到用户类型。
|
// 这里兜底为管理员类型，避免 WebSocketSessionManagerImpl 用 null 作为 key 记录会话时抛 NPE 断开连接。
|
if (loginUser.getUserType() == null) {
|
loginUser.setUserType(UserTypeEnum.ADMIN.getValue());
|
}
|
WebSocketFrameworkUtils.setLoginUser(loginUser, attributes);
|
return true;
|
}
|
|
@Override
|
public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
|
WebSocketHandler wsHandler, Exception exception) {
|
// do nothing
|
}
|
|
}
|
