WebSocket (连接前验证token)

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

实战

引入依赖

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

WebSocket配置类

重写modifyHandshake方法，从握手请求中提取token，并尝试从token中获取用户ID，然后将用户ID保存在WebSocket的session属性中。

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

import com.keepc.common.utils.JwtHelper;
import com.keepc.common.utils.uuid.IdUtil;

import java.util.List;
import java.util.Map;

import javax.websocket.HandshakeResponse;
import javax.websocket.server.HandshakeRequest;
import javax.websocket.server.ServerEndpointConfig;

/**
 * WebSocket配置类，用于配置WebSocket的相关设置。
 */
@Configuration
public class WebSocketConfig extends ServerEndpointConfig.Configurator {

    /**
     * 创建并返回一个ServerEndpointExporter的实例。
     * ServerEndpointExporter是Spring提供的用于注册WebSocket端点的组件。
     * 它会扫描并注册所有注解了@ServerEndpoint的类，使得这些WebSocket端点可以被服务器使用。
     *
     * @return ServerEndpointExporter实例
     */
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }

    /**
     * 重写modifyHandshake方法，用于修改握手请求和响应。
     * 此方法会从握手请求中提取token，并尝试从token中获取用户ID，
     * 然后将用户ID保存在WebSocket的session属性中。
     *
     * @param sec      ServerEndpointConfig对象，用于获取用户属性
     * @param request  HandshakeRequest对象，用于获取请求头
     * @param response HandshakeResponse对象，用于设置响应头
     */
    @Override
    public void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {
        // 初始化用户属性Map，可通过session.getUserProperties()获取
        final Map<String, Object> userProperties = sec.getUserProperties();
        // 从请求头中获取token
        Map<String, List<String>> headers = request.getHeaders();
        List<String> tokenHeaders = headers.get("token");
        String token = tokenHeaders.get(0);
        // 尝试从token解析用户id
        String userId = "";
        if (token != null) {
            userId = JwtHelper.getUserId(token);
        }

        // 将解析出的用户id或生成的未知用户id放入userProperties
        if (userId != null) {
            userProperties.put("userId", userId);
        } else {
            userProperties.put("unknownId", "未知用户" + IdUtil.fastSimpleUUID());
        }

    }
}

创建websocket的服务核心类

实现websocket的连接、释放、发送、报错等核心功能

import javax.websocket.OnClose;
import javax.websocket.OnError;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import com.keepc.system.config.WebSocketConfig;

import java.io.IOException;
import java.util.Iterator;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket服务端，提供与客户端的实时通信能力。
 * 地址：ws://127.0.0.1:8800/ws
 */
@Component
@ServerEndpoint(value = "/ws", configurator = WebSocketConfig.class) // 指定WebSocketConfig配置
public class WebSocketServer {

    /**
     * 用于记录日志信息。
     */
    private static final Logger log = LoggerFactory.getLogger(WebSocketServer.class);

    /**
     * 保存所有在线客户端的Session，以支持消息广播和单点发送。
     */
    public static final Map<String, Session> sessionMap = new ConcurrentHashMap<>();

    /**
     * 当WebSocket连接打开时的处理逻辑。
     * 
     * @param session WebSocket会话对象，用于与客户端进行通信。
     * @param userId  通过URL路径参数传递的用户ID，用于标识用户。
     */
    @OnOpen
    public void onOpen(Session session) {
        // 进行用户连接验证
        final boolean isverify = openVerify(session);
        if (isverify) {
            // 验证通过，添加用户到在线列表
            String id = (String) session.getUserProperties().get("userId");
            sessionMap.put(id, session);
            log.info("用户ID为={}加入连接, 当前在线人数为：{}", id, sessionMap.size());
        } else {
            // 验证不通过，关闭连接
            try {
                session.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 当WebSocket连接关闭时的处理逻辑。
     * 
     * @param session WebSocket会话对象。
     */
    @OnClose
    public void onClose(Session session) {
        // 从在线列表中移除断开连接的用户
        String id = (String) session.getUserProperties().get("userId");
        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 message 客户端发送的消息。
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        // 向服务端发送消息，并进行日志记录
        String id = (String) session.getUserProperties().get("userId");
        log.info("服务端收到来自用户ID为={}的消息:{}", id, message);
        sendOneMessage(id, "服务端收到消息:" + message);
    }

    /**
     * 当WebSocket发生错误时的处理逻辑。
     * 
     * @param session WebSocket会话对象。
     * @param error   异常错误。
     */
    @OnError
    public void onError(Session session, Throwable error) {
        // 记录异常错误日志
        log.error("websocket发生异常错误:");
        error.printStackTrace();
    }

    /**
     * 广播消息到所有连接的客户端。
     * 
     * @param message 要广播的消息内容。
     */
    public void sendAllMessage(String message) {
        // 向所有在线用户广播消息
        log.info("【WebSocket消息】广播消息：" + message);
        Iterator<Entry<String, Session>> entries = sessionMap.entrySet().iterator();
        while (entries.hasNext()) {
            Entry<String, Session> entry = entries.next();
            Session toSession = entry.getValue();
            if (toSession.isOpen()) {
                try {
                    log.info("服务端给客户端[{}]，用户{}，发送消息{}", toSession.getId(), entry.getKey(), message);
                    toSession.getAsyncRemote().sendText(message);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 向指定用户发送单点消息。
     * 
     * @param userId  目标用户的ID。
     * @param message 要发送的消息内容。
     */
    public void sendOneMessage(String userId, String message) {
        // 向指定用户发送消息
        Session toSession = sessionMap.get(userId);
        if (toSession != null && toSession.isOpen()) {
            try {
                synchronized (toSession) {
                    log.info("【WebSocket消息】单点消息：" + message);
                    toSession.getAsyncRemote().sendText(message);
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 向多个指定用户发送单点消息。
     * 
     * @param userIds 目标用户的ID列表。
     * @param message 要发送的消息内容。
     */
    public void sendMoreMessage(String[] userIds, String message) {
        // 向多个指定用户发送消息
        for (String userId : userIds) {
            sendOneMessage(userId, message);
        }
    }

    /**
     * 判断是否是合法用户。根据用户ID进行验证。
     * 
     * @param session WebSocket会话对象，用于获取用户ID。
     * @return 如果用户ID合法返回true，否则返回false。
     */
    public static boolean openVerify(Session session) {
        // 验证用户ID是否合法
        final String id = (String) session.getUserProperties().get("userId");
        if (id == "" || id == null) {
            return false;
        } else {
            return true;
        }
    }
}

Controller调用

import com.keepc.common.result.Result;
import com.keepc.system.webscoket.WebSocketServer;

import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

@RestController
@RequestMapping("/admin/system/ws")
@Api(tags = "WebSocket管理")
public class WebSocketController {

    @Autowired
    private WebSocketServer webSocketServer;

    // @PreAuthorize("hasAuthority('btn.ws.broadcas')")
    @ApiOperation(value = "广播信息")
    @PostMapping("/broadcas")
    public Result broadcasMessage(@RequestBody String message) {
        webSocketServer.sendAllMessage(message);
        return Result.ok();
    }

}

测试

使用postman连接

调用api发送广播消息

控制台输出日志

遇到的问题

websocket一直无法连接？

检查WebSocket配置类和核心服务类代码是否正确；是否使用过滤器、拦截器等组件拦截了请求；是否使用SpringSecurity等框架。修改配置放行 /ws请求。

过滤器放行参考

// 如果是websocket接口，直接放行
if ("/ws".equals(request.getRequestURI())) {
    chain.doFilter(request, response);
    return;
}

SpringSecurity放行参考

http.antMatchers("/ws").permitAll();

/**
 * 配置哪些请求不拦截
 * 排除swagger相关请求
 *
 * @param web WebSecurity对象
 * @throws Exception 异常情况
 */
@Override
public void configure(WebSecurity web) throws Exception {
    web.ignoring().antMatchers("/favicon.ico", "/swagger-resources/**", "/webjars/**", "/v2/**",
            "/swagger-ui.html/**", "/doc.html", "/ws/**");
}

测试类无法启动？

测试类报错： Caused by: java.lang.IllegalStateException: jakarta.websocket.server.ServerContainer not available

使用随机端口启动Spring Boot的Web应用程序进行测试。

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
/**
 * 使用随机端口启动Spring Boot的Web应用程序进行测试。
 * 这个注解配置使得测试时Spring Boot应用会随机选择一个端口来启动，避免了端口冲突的问题。
 * 适用于需要进行Web层测试的场景，例如RESTful服务的测试。
 */

扩展

案例消息以String类型为例，可自定义JSON消息规则。接受消息时根据规则判断即可。

{
    "type": 1, //消息类型
    "content": "Hello", //消息内容
    "send_id": "001", //消息发送人
    "accept_id": "", //消息接收人
    //消息接收组
    "accept_group": [
        "001",
        "002",
        "003"
    ] 
}