【WebSocket连接异常】前端使用WebSocket子协议传递token时，Java后端的正确打开方式！！！

前言：
本篇文章记录的是使用WebSocket进行双向通信时踩过的坑，希望能够帮助大家找到解决连接异常的正确方法。

1. 背景

本人在使用WebSocket实现“聊天室”的实时双向通信时（发消息、添加好友、处理好友请求等），一开始使用 cookie + session 的方式来管理用户的上下线情况，后来想引入 JWT，使用 token的方式来增强系统的可用性。这时我遇到了一个问题，大部分的接口都是使用 HTTP 协议的方式传输数据，因此我们可以将令牌放在 Header中用于身份校验；而 WebSocket进行双向通信时，前端无法直接在 header添加token。
经过网上查阅资料可知，有其他的方式可以在 HTTP升级为WebSocket时携带 token：（1）在 URL中追加 token（2）使用WebSocket的子协议传递 token。（通过抓包可以知道，token放在Header的 “Sec-WebSocket-Protocol” 中）

2. 代码实现和异常发现

考虑到 token直接暴露在 url 的安全性及优雅性等因素，我最终选择使用 WebSocket子协议来传递 token。以下是个人操作的过程及心路历程，若只想知道解决方法，可直接查看 3.2 从 WebSocket子协议的使用方式入手。

前端代码如下：

var token = localStorage.getItem("token");
let websocket = new WebSocket("ws://" + location.host + "/WebSocketMessage", [token]);

对于后端来说，可以使用自定义拦截器来验证并处理token（存储token信息，以便后续在WebSocketSession中处理消息时使用），具体方法是自定义类继承 HandshakeInterceptor ，并重写它的两个方法。
建立连接前处理token的代码如下：

@Component
public class SaveTokenInterceptor implements HandshakeInterceptor {
    // 握手前的操作，该方法返回 true 代表同意建立 WebSocket连接，false代表拒绝建立连接
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
    	// HTTP协议 未正式升级为 WebSocket时，可以对 HTTP 报文中的信息进行一定的处理
        // 1. 从 header中获取子协议的token，若token为空、过期或非法的，则拒绝建立 WebSocket连接
        String token = request.getHeaders().getFirst("Sec-WebSocket-Protocol");
        Claims claims = JwtUtil.parseToken(token);		// 解析令牌
        if (claims == null) return false;

        // 2. 将 token 中的信息放入到 attributes属性中，后续 WebSoketSession可通过方法获取 attributes，进而获取里面存放的信息
        int id = (int) claims.get(Constant.CLAIM_USERID);
        String username = (String) claims.get(Constant.CLAIM_USERNAME);
        attributes.put(Constant.USER_TOKEN_KEY, new User(id, username, ""));	
        return true;
    }


    // 握手完成后的操作
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
		
    }
}

连接完成后，查看 WebSocketSession 是否能够正确拿到存储到 attributes 中的属性(通过第一个方法查看)

@Component
@Slf4j
public class TestWebSocket extends TextWebSocketHandler {
	// 这个方法会在 WebSocket建立成功后被自动调用
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println("[WebSocketAPI] 连接成功!");
        // session.getAttributes() 得到一个 Map
        // 里面的元素为之前服务器Session存储的Attribute或放进去的其他自定义信息（上述处理token后存储的User对象）
        User user = (User) session.getAttributes().get(Constant.USER_TOKEN_KEY);
        log.info("[WebSocketAPI] afterConnectionEstablished, user: " + user);	// 验证是否将token信息放进去了
        
        if(user == null) {
            System.out.println("用户未登录!");
            return;
        }

        // 往 hash表 中存储对应客户端的WebSocket对象
        onlineUserManager.online(user.getId(), session);
    }


    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 这个方法是在 websocket 收到消息后被自动调用
        System.out.println("[WebSocketAPI] 收到消息! " + message.toString());
    }

    // 这个方法是在连接出现异常时被自动调用
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.out.println("[WebSocketAPI] 连接异常! " + exception.getMessage());
        User user = (User) session.getAttributes().get(Constant.USER_TOKEN_KEY);
        if(user == null) {
            return;
        }
        onlineUserManager.offline(user.getId(), session);
    }


    // 这个方法是在连接正常关闭后被自动调用
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        System.out.println("[WebSocketAPI] 关闭! " + status.toString());
        User user = (User) session.getAttributes().get(Constant.USER_TOKEN_KEY);
        if(user == null) {
            return;
        }
        onlineUserManager.offline(user.getId(), session);
    }
}

通过抓包及后端控制台日志观察上述过程：


可以发现：WebSocketSession 已经能够正确拿到 token里的信息，但是控制台也出现了WebSocket连接异常 和token校验失败两个异常现象。（通过浏览器控制台也可发现连接异常）

3. 解决异常

3.1 从 URL入手

首先，token校验失败原因比较简单，一般是在拦截器拦截 HTTP请求时发生，于是我通过抓包进行分析，但是令人感到奇怪的是所有 HTTP 请求均正常携带了 token，为什么会出现令牌解析不成功的情况呢？

经过一番思考，我决定在拦截器拦截请求时，通过 request获取所有经过拦截器的 HTTP请求的 URL，通过打印每个 HTTP 请求的URL及Header携带的 token 分析是否是前端 WebSocket 使用了子协议而导致被拦截器拦截，从而导致的异常。

@Component
public class LoginInterceptor implements HandlerInterceptor {

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String requestURI = request.getRequestURI();

        // 在方法执行前进行拦截，此处判断哪些方法可以被执行
        // 从 header中的token 判断用户是否登录
        String token = request.getHeader(Constant.USER_TOKEN_HEADER);
        System.out.println(token);
        System.out.println(requestURI);
        if (JwtUtil.parseToken(token) == null) {
            response.setStatus(401);
            return false;
        }
        return true;
    }
}

可以发现：上面的 HTTP 请求都符合预期，出现异常是使用 token 验证用户身份时，由浏览器默认发起的 favicon/ico（GET请求）并不会像其他 HTTP 请求一样，在其 Header 上携带 token，因此出现了令牌校验失败的情况。

通过上述偶然出现的异常也可以发现该程序上的一个问题，若一个 HTTP 的 Header 没有携带 token（即 token == null）就不需要进行令牌解析了，直接拦截即可。

因此只需将上述拦截器的拦截规则多做一个判断即可解决令牌解析的异常。（代码如下）

@Component
public class LoginInterceptor implements HandlerInterceptor {

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String requestURI = request.getRequestURI();

        // 在方法执行前进行拦截，此处判断哪些方法可以被执行
        // 从 header中的token 判断用户是否登录
        String token = request.getHeader(Constant.USER_TOKEN_HEADER);
        if (token == null || JwtUtil.parseToken(token) == null) {
            response.setStatus(401);
            return false;
        }
        return true;
    }
}

3.2 从 WebSocket子协议的使用方式入手（真正原因）

由于抓包并不能找到问题出现的原因，因此我查阅了 WebSocket 子协议的相关使用方式发现：如果前端使用了子协议携带了 token，在 WebSocket连接完成后，返回的响应报文应该携带相同的子协议内容。

因此我立马通过抓包查看了响应报文：

可以发现，响应报文确实没有携带对应的 Header，为了验证 WebSocket连接异常的原因导致及上述说法的正确性，我对代码作出了如下修改：

@Component
public class SaveTokenInterceptor implements HandshakeInterceptor {
    // 握手前的操作
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        // 1. 从 header中获取子协议的token，若token为空、过期或非法的，则拒绝建立 WebSocket连接
        String token = request.getHeaders().getFirst("Sec-WebSocket-Protocol");
        System.out.println("[SaveTokenInterceptor] beforeHandshake方法，token: " + token);
        Claims claims = JwtUtil.parseToken(token);
        if (claims == null) return false;

        // 2. 将 id 和 username 存入WebSocket的 attributes中
        int id = (int) claims.get(Constant.CLAIM_USERID);
        String username = (String) claims.get(Constant.CLAIM_USERNAME);
        attributes.put(Constant.USER_TOKEN_KEY, new User(id, username, ""));
        return true;
    }


    // 握手完成后的操作
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        // 获取 Servlet 的 HttpServletRequest 和 HttpServletResponse 对象
        // httpRequest 可以获取 HTTP协议升级前 请求报文的信息，如 header中的键值对等
        // httpResponse 可以设置 HTTP响应 的相关信息，如状态码、ContentType、header信息等
        HttpServletRequest httpRequest = ((ServletServerHttpRequest) request).getServletRequest();
        HttpServletResponse httpResponse = ((ServletServerHttpResponse) response).getServletResponse();
        if (httpRequest.getHeader("Sec-WebSocket-Protocol") != null) {
            httpResponse.addHeader("Sec-WebSocket-Protocol", httpRequest.getHeader("Sec-WebSocket-Protocol"));
        }
    }
}

上述代码即在 WebSocket 连接完成后，针对响应增加了一个子协议的 header。

注意：无法直接通过 afterHandshake 方法参数的 ServerHttpResponse 修改响应内容，因为该接口并没有提供修改响应的方法。由于ServerHttpResponse是一个接口，通过源码我们可以发现：
ServletServerHttpResponse类 实现了该接口，且在Spring中 ServletServerHttpResponse 对 Servlet的 HttpServletResponse 类进行了封装，因此我们可以将 方法参数中的 response 强转为底层的实现类ServletServerHttpResponse，再通过 ServletServerHttpResponse 类中的方法获取封装的 HttpServletResponse 类，然后就可以使用该类设置响应报文的内容。

对代码作出上述修改后，运行程序的结果如下：

4. 总结（仍然存在的问题）

通过上述修改后，已经能够使用 token 验证用户身份，管理用户上下线情况，但仍然存在问题：

1. 在使用 cookie-session 验证用户登录状态和上下线状态时，服务器重启重启会导致存储在内存的 session 消失，因此用户后续的任何请求都可能触发拦截器的拦截操作，需重新进行登录才能正常进行后续的操作。

而对于使用 token 来代替 cookie-session，虽然触发 HTTP 请求的操作能够做到 “用户无感知”，即服务器因某种原因重启后，用户不用二次登录依然可以完成操作；但对于使用 WebSocket 进行实时通信的消息转发、好友请求转发等功能来说，该程序使用 ConcurrentHashMap 来存储 WebSocketSession，服务器一旦重启，哈希表保存的登录信息就没了，这部分功能也因此直接“失效”了。

要想解决这个问题，可能需要引入 Redis 这样的中间件或使用其他的机制来实现 WebSocket 重连，以保证用户的使用体验。

2. 当令牌达到过期时间，而用户没有触发发送 HTTP 请求的操作，而是进行发送消息这种操作，那么上述存储用户信息的方式则是错误的，因为这种做法虽然可以让接口代码只有小幅度修改，但会出现用户令牌虽然过期了但 ConcurrentHashMap 存储的 WebSocketSession 并不会被立即移除的情况，仍然能够进行消息发送（上一次操作停留在对话框界面）。

以上就是本篇文章的全部内容了，如果这篇文章对你有些许帮助，你的点赞、收藏和评论就是对我最大的支持。
另外，文章可能存在许多不足之处，也希望你可以给我一点小小的建议，我会努力检查并改进。