WebSocket实现广播及一对一发送

请求示例

• ws-协议，如果是https对应要改成wss
• 第二个ws，WebSocketConfig配置文件中自定义的名称
• Connection：Upgrade，协议升级
• Upgrade：升级成websocket

/**
 * @description: WebSocket配置类
 * @date: 2023-09-19 15:52
 * @version: 1.0
 */
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {


    @Autowired
    private HttpAuthHandler httpAuthHandler;
    @Autowired
    private MyInterceptor myInterceptor;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry
                .addHandler(httpAuthHandler, "ws")
                .addInterceptors(myInterceptor)
                .setAllowedOrigins("*")
//                .withSockJS()
                ;
    }

}

注：启用.withSockJS()的话，会在浏览器不支持websocket协议的情况下自动切换为最优的其他协议，因为目前大部分的浏览器是支持WebSocket的，所以不启用该项；启用的话连接方式需要发生变化

/**
 * @description: websocket拦截器
 * @date: 2023-09-25 16:59
 * @version: 1.0
 */
@Component
public class MyInterceptor implements HandshakeInterceptor {


    /**
     * 握手前
     *
     * @param request
     * @param response
     * @param wsHandler
     * @param attributes
     * @return
     * @throws Exception
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        System.out.println("握手开始");
        // 获得请求参数
        // 获取url传递的参数，通过attributes在Interceptor处理结束后传递给WebSocketHandler
        // WebSocketHandler可以通过WebSocketSession的getAttributes()方法获取参数
        ServletServerHttpRequest serverRequest = (ServletServerHttpRequest) request;
        String uid = serverRequest.getServletRequest().getParameter("userId");

        if (StrUtil.isNotBlank(uid)) {
            // 放入属性域
            attributes.put("userId", uid);
            System.out.println("用户 " + uid + " 握手成功！");
            return true;
        }
        System.out.println("用户登录已失效");
        return false;
    }

    /**
     * 握手后
     *
     * @param request
     * @param response
     * @param wsHandler
     * @param exception
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler  wsHandler, Exception exception) {
        System.out.println("握手完成");
    }
}

/**
 * @description: 对 websocket 的事件进行处理
 * @date: 2023-09-25 16:52
 * @version: 1.0
 */
@Component
public class HttpAuthHandler extends TextWebSocketHandler {

    /**
     * socket 建立成功事件
     *
     * @param session
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        Object userId = session.getAttributes().get("userId");
        if (userId != null) {
            // 用户连接成功，放入在线用户缓存
            WsSessionManager.add(userId.toString(), session);
        } else {
            throw new RuntimeException("用户登录已经失效!");
        }
    }

    /**
     * 接收消息事件
     *
     * @param session
     * @param message
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 获得客户端传来的消息
        String payload = message.getPayload();
        Object userId = session.getAttributes().get("userId");
        Map<String, Object> stringObjectMap = JsonUtil.json2map(payload);
        if(Objects.nonNull(stringObjectMap.get("flag"))){
            if (stringObjectMap.get("flag").equals(1)) {
                //正常消息
                //session.sendMessage(new TextMessage("pong"));
                System.out.println("server 接收到 " + userId + " 发送的 " + payload);
            } else {
                //心跳
                System.out.println("server 接收到-心跳 " + userId + " 发送的 " + payload);
            }
        }
    }

    /**
     * socket 断开连接时
     *
     * @param session
     * @param status
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        Object userId = session.getAttributes().get("userId");
        if (userId != null) {
            // 用户退出，移除缓存
            WsSessionManager.remove(userId.toString());
        }
    }
}

/**
 * @description: 实现 session 池
 * @date: 2023-09-25 16:55
 * @version: 1.0
 */
@Slf4j
public class WsSessionManager {

    /**
     * 保存连接 session 的地方
     */
    private static ConcurrentHashMap<String, WebSocketSession> SESSION_POOL = new ConcurrentHashMap<>();

    /**
     * 添加 session
     *
     * @param key
     */
    public static void add(String key, WebSocketSession session) {
        // 添加 session
        SESSION_POOL.put(key, session);
    }

    /**
     * 删除 session,会返回删除的 session
     *
     * @param key
     * @return
     */
    public static WebSocketSession remove(String key) {
        // 删除 session
        return SESSION_POOL.remove(key);
    }

    /**
     * 删除并同步关闭连接
     *
     * @param key
     */
    public static void removeAndClose(String key) {
        WebSocketSession session = remove(key);
        if (session != null) {
            try {
                // 关闭连接
                session.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 获得 session
     *
     * @param key
     * @return
     */
    public static WebSocketSession get(String key) {
        // 获得 session
        return SESSION_POOL.get(key);
    }


}

//查询各个设备编码对应的userid，并推送数据
//一对一传输
if(!CollectionUtils.isEmpty(userIdList)){
    for(String userId : userIdList){
        resultList.forEach(t -> t.setUserId(userId));
        WebSocketSession session = WsSessionManager.get(userId);
        if(Objects.nonNull(session)){
            try {
                session.sendMessage(new TextMessage(JsonUtil.obj2Json(resultList)));
            } catch (IOException e) {
                logger.error("websocket推送失败", e);
            }
        }
    }
}

这里是使用userId作为区分每个客户端的依据，可以按需调整