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
     */
    @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();

        // 获取token
        Map<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所标注的对象
     */
    @Override
    public <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<>();

    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public 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 客户端发送过来的消息
     */
    @OnMessage
    public void onMessage(String message, Session session) {

        // 获取ID
        String 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("发送对象集合不能为空！");
        }
    }

    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public 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
     */
    @OnError
    public 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) {
        // 判断是否有合法的id
        final 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"]
}