Springboot实现websocket（连接前jwt验证token）

背景

用户连接服务器weksocket前，需经过jwt的token验证（token中包含账号信息），验证合法后，才可以于服务器正常交互。

实现

一、配置依赖（pom.xml）

        <!-- websocket --><dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-websocket</artifactId></dependency>

二、因为springboot的websocket连接时不会显示header信息，也就无法拿到cookie中的token信息，需要在连接前处理，新建一个WebSocketConfig.class，在连接前做一个jwt的token验证，并获取用户的账号信息添加到session中。

(关于jwt的token验证工具类我这里就不详细讲了，可以去看我的另一篇文章Springboot实现jwt的token验证（超简单）-CSDN博客)

package com.example.springboot.common.config;import com.example.springboot.utils.CharUtil;
import com.example.springboot.utils.CookieUtil;
import com.example.springboot.utils.TokenUtils;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;
import javax.websocket.HandshakeResponse;
import javax.websocket.server.HandshakeRequest;
import javax.websocket.server.ServerEndpointConfig;
import java.util.List;
import java.util.Map;/*** websocket开启支持*/
@Configuration
public class WebSocketConfig extends ServerEndpointConfig.Configurator {/*** 这个bean会自动注册使用了@ServerEndpoint注解声明的对象* @return*/@Beanpublic ServerEndpointExporter serverEndpointExporter() {return new ServerEndpointExporter();}/*** 建立握手时，连接前的操作*/@Overridepublic void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {// 这个userProperties 可以通过 session.getUserProperties()获取final Map<String, Object> userProperties = sec.getUserProperties();// 获取tokenMap<String, List<String>> headers = request.getHeaders();List<String> cookie = headers.get("cookie");String token = "";if (cookie != null) {token = CookieUtil.getCookie("token", cookie.get(0));}// 判断用户token是否合法，并获取用户id，入果非法，生成一个临时用户用于识别String id = "";try {TokenUtils.verify(token);id = TokenUtils.getTokenInfo(token).getClaim("user").asString();userProperties.put("id", id);} catch (Exception err) {id = "未知用户" + CharUtil.randomVerify();userProperties.put("unknownId", id);}}/*** 初始化端点对象,也就是被@ServerEndpoint所标注的对象*/@Overridepublic <T> T getEndpointInstance(Class<T> clazz) throws InstantiationException {return super.getEndpointInstance(clazz);}
}

 三、再新建一个websocket的服务核心类WebSocketServer.class，实现websocket的连接、释放、发送、报错等4大核心功能。

package com.example.springboot.service;import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.example.springboot.common.config.WebSocketConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;/*** @author websocket服务*/
@ServerEndpoint(value = "/ws",configurator = WebSocketConfig.class)
@Component
public class WebSocketServer {/*** 打印日志*/private static final Logger log = LoggerFactory.getLogger(WebSocketServer.class);/*** 记录当前在线连接数*/public static final Map<String, Session> sessionMap = new ConcurrentHashMap<>();/*** 连接建立成功调用的方法*/@OnOpenpublic void onOpen(Session session) {// 判断是否合法连接，如果不是，直接执行session.close()关闭连接final boolean isverify = openVerify(session);if (isverify) {// 获取用户账号String id = (String) session.getUserProperties().get("id");// 根据账号添加到session列表中sessionMap.put(id, session);log.info("用户ID为={}加入连接, 当前在线人数为：{}", id, sessionMap.size());} else {// 非法连接，进行释放try {session.close();} catch (IOException e) {e.printStackTrace();}}}/*** 后台收到客户端发送过来的消息* @param message 客户端发送过来的消息*/@OnMessagepublic void onMessage(String message, Session session) {// 获取IDString id = (String) session.getUserProperties().get("id");log.info("服务端收到来自用户ID为={}的消息:{}", id, message);// 将JSON数据转换为对象，方便操作JSONObject obj = JSONUtil.parseObj(message);// 储存需要发送的对象JSONArray users = new JSONArray();// 判断是否为群聊，1为否，2为是if (obj.getStr("type").equals("1")) {users.add(obj.getStr("send_id"));users.add(obj.getStr("accept_id"));} else if (obj.getStr("type").equals("2")) {users = obj.getJSONArray("accept_group");}// 判断是否存在发送对象if (users.size() > 0) {for (int i=0;i<users.size();i++){Session toSession = sessionMap.get(users.get(i));if (toSession != null) {this.sendMessage(obj.toString(), toSession);log.info("服务器发送消息给用户ID为={}，消息：{}", users.get(i), obj.toString());} else {log.info("发送失败，用户ID为={}未在线", users.get(i));}}} else {log.info("发送对象集合不能为空！");}}/*** 连接关闭调用的方法*/@OnClosepublic void onClose(Session session) {// 判断断开的连接是否是合法的String id = (String) session.getUserProperties().get("id");if (id == "" & id == null) {sessionMap.remove(id);log.info("有一连接正常关闭，移除username={}的用户session, 当前在线人数为：{}", id, sessionMap.size());} else {id = (String) session.getUserProperties().get("unknownId");sessionMap.remove(id);log.info("token验证不通过，移除username={}的用户session, 当前在线人数为：{}", id, sessionMap.size());}}/*** 异常报错* @param session* @param error*/@OnErrorpublic void onError(Session session, Throwable error) {log.error("websocket发生异常错误:");error.printStackTrace();}/*** 服务端发送消息给客户端*/private void sendMessage(String message, Session toSession) {try {log.info("服务端给客户端[{}]发送消息{}", toSession.getId(), message);toSession.getBasicRemote().sendText(message);} catch (Exception e) {log.error("服务端发送消息给客户端失败", e);}}/*** 判断是否是合法用户。是就返回true，反之返回false* @param session* @return*/public static boolean openVerify (Session session) {// 判断是否有合法的idfinal String id = (String) session.getUserProperties().get("id");if (id == "" | id == null) {return false;} else {return true;}}
}

 四、连接路径如下

http://localhost:8899/ws

五、前端发送的JSON数据格式

// 单聊
{"type":1,"content":"你好啊","send_id":"001","accept_id":"002","accept_group":[]
}// 群聊
{"type":2,"content":"大家好啊","send_id":"001","accept_id":"","accept_group":["001","002","003"]
}