Java实现Websocket获取protocols中的Token

因为token不希望通过url后缀的方式传输 所以通过Sec-WebSocket-Protocol请求头的方式来进行传输 这样在请求连接的时候就可以进行鉴权 token也不会暴露在url中

1. 前端调用方式

let ws = new WebSocket(url, [token]);

2. WebsocketConfig

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(new WebSocketServer(), "/ws").setAllowedOrigins("*")
                // 设置握手拦截器
                .addInterceptors(handshakeInterceptor());;
    }

    /**
     * 装配 HttpSession 的拦截器，这样就可以在握手阶段
     * 获取 HttpSession 的内容,在使用 WebSocketSession 时
     * 就能直接得到 HttpSession 的数据
     */
    @Bean
    public HandshakeInterceptor handshakeInterceptor(){
        return new WebsocketHandshakeInterceptor();
    }

}

3. WebsocketHandshakeInterceptor

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.basics.pojo.ComSysUser;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;

/**
 * Websocket握手拦截器
 */
@Slf4j
@Component
public class WebsocketHandshakeInterceptor implements HandshakeInterceptor {

    private static final String WEBSOCKET_PROTOCOL = "Sec-WebSocket-Protocol";

    @Resource
    private TokenStore tokenStore;
 
    /**
     * 初次握手访问前 校验token
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler handler, Map<String, Object> map) {
        if (request instanceof ServletServerHttpRequest) {
            HttpServletRequest req = ((ServletServerHttpRequest) request).getServletRequest();
            String authorization = req.getHeader(WEBSOCKET_PROTOCOL);
            log.info("authorization = {}", authorization);
            if (ObjectUtil.isNull(tokenStore)) {
                tokenStore = SpringUtil.getBean("tokenStore", TokenStore.class);
            }
            OAuth2AccessToken accessToken = tokenStore.readAccessToken(authorization);
            if (accessToken == null || StrUtil.isBlank(accessToken.getValue())) {
                response.setStatusCode(HttpStatus.FORBIDDEN);
                log.info("websocket token校验失败. token = {}", authorization);
                return false;
            }
            OAuth2Authentication oAuth2Authentication = tokenStore.readAuthentication(accessToken);
            ComSysUser sysUser = (ComSysUser)oAuth2Authentication.getUserAuthentication().getPrincipal();
            if (ObjectUtil.isNull(sysUser.getUserId())){
                response.setStatusCode(HttpStatus.FORBIDDEN);
                log.info("websocket token校验失败. token = {}", authorization);
                return false;
            }
            map.put("userId", sysUser.getUserId());
            map.put("createTime", System.currentTimeMillis());
            map.put("token", authorization);
            map.put("sessionId", req.getSession().getId());
            log.info("【beforeHandshake】 WEBSOCKET_INFO_MAP: {}", map);
        }
        return true;
    }
 
    /**
     * 初次握手访问后，将前端自定义协议头Sec-WebSocket-Protocol原封不动返回回去，否则会报错
     */
    @Override
    public void afterHandshake(@NonNull ServerHttpRequest request, @NonNull ServerHttpResponse response, @NonNull WebSocketHandler handler, Exception exception) {
        HttpServletRequest httpRequest = ((ServletServerHttpRequest) request).getServletRequest();
        HttpServletResponse httpResponse = ((ServletServerHttpResponse) response).getServletResponse();
        if (StringUtils.isNotEmpty(httpRequest.getHeader("Sec-WebSocket-Protocol"))) {
            httpResponse.addHeader(WEBSOCKET_PROTOCOL, httpRequest.getHeader(WEBSOCKET_PROTOCOL));
        }
    }
}

4.WebSocketServer

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;

@Slf4j
@Service
public class WebSocketServer implements WebSocketHandler {

    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception {
        log.info("afterConnectionEstablished: {}", webSocketSession);
    }

    @Override
    public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> webSocketMessage) throws Exception {
        log.info("handleMessage: {}", webSocketMessage);
    }

    @Override
    public void handleTransportError(WebSocketSession webSocketSession, Throwable throwable) throws Exception {
        log.info("handleTransportError: {}", webSocketSession);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus closeStatus) throws Exception {
        log.info("afterConnectionClosed: {}", webSocketSession);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}