Java SpringBoot集成WebSocket

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

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import java.util.Map;
/**
 * @author wuzhenyong
 * ClassName:WebSocketInterceptor.java
 * date:2023-07-04 13:58
 * Description: WebSocket拦截器配置
 */
@Component
@Slf4j
public class WebSocketInterceptor implements HandshakeInterceptor {
    /**
     * 在握手之前 可用于权限认证
     *
     * @param serverHttpRequest  服务器http请求
     * @param serverHttpResponse 服务器http响应
     * @param webSocketHandler   网络套接字处理程序
     * @param map                地图
     * @return boolean
     * @throws Exception 异常
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Map<String, Object> map) throws Exception {
        log.info("握手之前");
        return true;
    }
    /**
     * 握手之后
     *
     * @param serverHttpRequest  服务器http请求
     * @param serverHttpResponse 服务器http响应
     * @param webSocketHandler   网络套接字处理程序
     * @param e                  e
     */
    @Override
    public void afterHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Exception e) {
        log.info("握手之后");
    }
}

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;
/**
 * @author wuzhenyong
 * ClassName:WebSocketConfig.java
 * date:2023-07-04 15:28
 * Description: WebSocket配置
 */
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    /**
     * 拦截器容器注册
     *
     * @return {@link WebSocketInterceptor}
     */
    @Bean
    public WebSocketInterceptor getWebSocketInterceptor() {
        return new WebSocketInterceptor();
    }
    /**
     * WebSocket服务注册
     *
     * @return {@link WebSocketServer}
     */
    @Bean
    public WebSocketServer myServer() {
        return new WebSocketServer();
    }
    /**
     * 端点注册
     *
     * @return {@link ServerEndpointExporter}
     */
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // 此处定义webSocket的连接地址以及允许跨域
        registry.addHandler(myServer(), "/websocket").addInterceptors(getWebSocketInterceptor()).setAllowedOrigins("*");
        // 同上，同时开启了Sock JS的支持，目的为了支持IE8及以下浏览器
        registry.addHandler(myServer(), "/sockjs/websocket").addInterceptors(getWebSocketInterceptor()).setAllowedOrigins("*").withSockJS();
    }
}

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import java.io.IOException;
import java.util.HashSet;
import java.util.Set;
/**
 * @author wuzhenyong
 * ClassName:WebSocketServer.java
 * date:2023-07-04 15:46
 * Description: WebSocket服务
 */
@Component
@Slf4j
public class WebSocketServer implements WebSocketHandler {
    private static Set<WebSocketSession> webSocketSet = new HashSet<>();
    /**
     * 连接后回调
     *
     * @param session 会话
     * @throws Exception 异常
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        webSocketSet.add(session);
        log.info("有新连接加入.当前在线人数为:{}" , webSocketSet.size());
    }
    /**
     * 收到消息触发
     *
     * @param session 会话
     * @param message 消息
     * @throws Exception 异常
     */
    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        log.info("收到新的消息.内容:{}" ,message.getPayload().toString());
    }
    /**
     * 异常
     *
     * @param session   会话
     * @param throwable throwable
     * @throws Exception 异常
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable throwable) throws Exception {
        webSocketSet.remove(session);
        log.info("websocket发生异常！" , throwable);
    }
    /**
     * 关闭连接
     *
     * @param session     socket会话
     * @param closeStatus 关闭状态
     * @throws Exception 异常
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        webSocketSet.remove(session);
        log.debug("webSocket关闭连接，状态：{}，当前连接数：{}", closeStatus, webSocketSet.size());
    }
    /**
     * 是否支持消息分片
     *
     * @return boolean
     */
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
    /**
     * 发送消息
     *
     * @param message 消息
     * @throws IOException ioexception
     */
    public static void sendString(String message) throws IOException {
        for (WebSocketSession webSocket : webSocketSet) {
            if (webSocket.isOpen()) {
                webSocket.sendMessage(new TextMessage(message));
            }
        }
        log.debug("webSocket发送消息内容：{}，当前连接数：{}", message, webSocketSet.size());
    }
}

import io.swagger.annotations.ApiOperation;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
/**
 * @author wuzhenyong
 * ClassName:WebSocketController.java
 * date:2023-07-04 15:54
 * Description: WebSocket发送消息测试
 */
@RestController
@RequestMapping("/ws")
public class WebSocketController {
    @ApiOperation(value = "发送webSocket消息")
    @PostMapping("/sendMsg")
    public String sendWebSocketMessage(@RequestParam("msg") String msg) {
        // 发送webSocket消息
        WebSocketServer.sendString(msg);
        return "发送成功";
    }
}