SpringSecurity整合WebSocket并携带token

目的

导入SpringSecurity的SpringBoot项目，在连接WebSocket时进行token校验

实现

SpringBoot整合SpringBoot的相关知识就不过多赘述，本文主要介绍WebSocket权限校验相关

1. 前端

WebSocket连接

 var windowTag = `${user.id}-${Math.random().toString(36).substr(2)}`;
 var token = user.token;

websocket = new WebSocket(`ws://localhost:9001/ws/chat/${windowTag}`,[token]);

windowTag是生成的随机窗口唯一标识符，token是用户登录后生成的令牌token
当前端发起WebSocket连接请求时，请求头在通信子协议Sec-WebSocket-Protocol里携带token

2. 后端

前端通过WebSocket的通信子协议携带token发送给后端，现在我们只需要获取到该token就能获取用户信息

/**
  WebSocket配置
*/
@Configuration
public class WebSocketConfig extends ServerEndpointConfig.Configurator {

    @Bean
    public ServerEndpointExporter serverEndpointExporter(){
        return new ServerEndpointExporter();
    }

    /**
     * 建立握手时，连接前的操作
     */
    @Override
    public void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {
        // 这个userProperties 可以通过 session.getUserProperties()获取
        final Map<String, Object> userProperties = sec.getUserProperties();
        Map<String, List<String>> headers = request.getHeaders();
        List<String> protocol = headers.get(WEBSOCKET_PROTOCOL);
        // 存放自己想要的header信息
        if(protocol != null){
            userProperties.put(WEBSOCKET_PROTOCOL, protocol.get(0));
        }
    }

    /**
     * 初始化端点对象,也就是被@ServerEndpoint所标注的对象
     */
    @Override
    public <T> T getEndpointInstance(Class<T> clazz) throws InstantiationException {
        return super.getEndpointInstance(clazz);
    }

}

将请求头中Sec-WebSocket-Protocol携带的token放入session的userProperties中，方便连接时获取token

@Slf4j
@Component
@ServerEndpoint(value = "/ws/chat/{windowTag}",configurator = WebSocketConfig.class)
public class ChatEndPoint {

    //用线程安全的map来保存当前用户
    private static Map<String, ChatEndPoint> onlineUsers = new ConcurrentHashMap<>();
    //声明一个session对象，通过该对象可以发送消息给指定用户，不能设置为静态，每个ChatEndPoint有一个session才能区分.(websocket的session)
    private Session session;

    //建立连接
    @OnOpen
    public void onOpen(Session session, @PathParam("windowTag") String windowTag){
        this.session = session;
        String username = getUserName(session);
        log.info("上线用户名称: {}", username);
        onlineUsers.put(username + "-" + windowTag, this);
        log.info("在线用户数: {}", onlineUsers.size());
    }

	......

    // 获取用户名
    private String getUserName(Session session){
        String token = getHeader(session, WEBSOCKET_PROTOCOL);
        return new TokenManager().getUserInfoFromToken(token);
    }

	public String getHeader(Session session, String headerName) {
        String header = (String) session.getUserProperties().get(headerName);
        if (StringUtils.isBlank(header)) {
            try {
                session.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return header;
    }
}

通过子协议中携带的token获取用户名，并与窗口标识符拼接成连接标识符，之后将连接的session存放进线程安全的Map中

/**
	SpringSecurity的权限校验器
*/
public class TokenAuthFilter extends BasicAuthenticationFilter {

    private TokenManager tokenManager;

    public TokenAuthFilter(AuthenticationManager authenticationManager, TokenManager tokenManager) {
        super(authenticationManager);
        this.tokenManager = tokenManager;
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain chain) throws IOException, ServletException {
        System.out.println("根据token获取用户权限并放入security上下文...");
        // websocket需要验证Sec-WebSocket-Protocol中的token
        String token = request.getHeader(WEBSOCKET_PROTOCOL);
        if(token == null){
            token = request.getHeader("token");
        }
        //获取当前认证成功用户权限信息
        UsernamePasswordAuthenticationToken authRequest = getAuthentication(token);
        //判断如果有权限信息，放到权限上下文中
        if(authRequest != null) {
            SecurityContextHolder.getContext().setAuthentication(authRequest);
        }
        chain.doFilter(request,response);
    }

    private UsernamePasswordAuthenticationToken getAuthentication(String token) {
        System.out.println("请求头中的token: " + token);
        if(!token.equals("null")) {
            //从token获取用户名
            String username = tokenManager.getUserInfoFromToken(token);
            System.out.println("token获取用户名:"+username);
            //从token获取对应权限列表
            List<String> permissionValueList = tokenManager.getUserPermissionList(token);
            Collection<GrantedAuthority> authority = new ArrayList<>();
            if(permissionValueList != null){
                for(String permissionValue : permissionValueList) {
                    SimpleGrantedAuthority auth = new SimpleGrantedAuthority(permissionValue);
                    authority.add(auth);
                }
            }
            return new UsernamePasswordAuthenticationToken(username,token,authority);
        }
        return null;
    }
}

在SpringSecurity权限校验时先获取Sec-WebSocket-Protocol携带的token
当我们以为一切准备就绪时，运行时发现报错了

WebSocket握手阶段出错：发送了非空“Sec-WebSocket-Protocol”请求头但是响应中没有此字段。在后端握手时设置一下请求头（Sec-WebSocket-Protocol）即可，前端发来什么值，这里就写什么值

@Order(1)
@Component
@WebFilter(filterName = "WebsocketFilter", urlPatterns = "/ws/**")
public class WebSocketFilter implements Filter {

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
        HttpServletResponse response = (HttpServletResponse) servletResponse;

        String token = ((HttpServletRequest) servletRequest).getHeader(WEBSOCKET_PROTOCOL);
        // 解决 Sent non-empty 'Sec-WebSocket-Protocol' header but no response was received
        response.setHeader(WEBSOCKET_PROTOCOL,token);

        filterChain.doFilter(servletRequest, servletResponse);
    }

    @Override
    public void destroy() {}
}

连接成功！！


功能测试