[网络协议]Websocket 与 SpringBoot 集成使用

Websocket 与 SpringBoot 集成使用

1.依赖

	<dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-websocket</artifactId>
            <exclusions>
            <!--排除 websocket 集成中带有的 tomcat-->
                <exclusion>
                    <groupId>org.springframework.boot</groupId>
                    <artifactId>spring-boot-starter-tomcat</artifactId>
                </exclusion>
            </exclusions>
        </dependency>

2. Websocket 配置

/**
 * @Description: websocket 配置类
 */

@Configuration
@EnableWebSocket  // 开启webscoket 
public class SocketConfig implements WebSocketConfigurer {

    // 自定义 webscoket 业务处理逻辑类，类似于 Controller类，具体实例见下面
    @Autowired
    private SocketAuthHandler socketAuthHandler;
    
     // 自定义 webscoket 握手拦截器，可以对此进行 toke你校验，进行数据赋值等处理，具体实例见下面
    @Autowired
    private SocketInterceptor socketInterceptor;

	// 自定义 webscoket 链接端口，类似于 @PostMapping("tagMail")
    private static final String WEB_SOCKET_PATH = "websocket";

    /**
     * @param webSocketHandlerRegistry
     *  把webscoket信息注册进去
     */
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry webSocketHandlerRegistry) {
        webSocketHandlerRegistry
        		// 配置业务处理逻辑和websocket 链接端口，可以多次调用 .addHandler（） ，配置多个
            	// 前端的链接地址是：ws:// + location.host （ip：port） +  WEB_SOCKET_PATH
            	// 例如 ：ws://localhost:63343/websocket   后面可以带上参数
                .addHandler(socketAuthHandler, WEB_SOCKET_PATH)
                .addInterceptors(socketInterceptor)
            	// 配置支持跨域
                .setAllowedOrigins("*");
    }

}

3. SocketSessionManager SocketSession 会话管理类

import com.google.common.cache.Cache;

/**
 * 
 * @Description: SocketSession 会话管理类
 * 
 */
@Slf4j
public class SocketSessionManager {

    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();


    /**
     * websocket 会话池
     * 如果想长期存在会话池  ，可以使用  Map<String, WebSocketSession>  来存储
     * 
     */
    private final static Cache<String, WebSocketSession> WEB_SOCKET_SESSION_MAP = CacheBuilder.newBuilder()
            .maximumSize(10240) // 缓存的最大大小
            .expireAfterAccess(3, TimeUnit.MINUTES) // 指定过期时间
            .build();

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

    /**
     * 移除 websocket 会话,并将该会话内容返回
     *
     * @param key
     * @return
     */
    public static WebSocketSession remove(String key) {
        WebSocketSession session = WEB_SOCKET_SESSION_MAP.getIfPresent(key);
        WEB_SOCKET_SESSION_MAP.invalidate(key);
        return session;
    }

    /**
     * 删除 websocket,并关闭连接
     *
     * @param key
     */
    public static void removeAndClose(String key) {
        WebSocketSession session = remove(key);
        if (session != null) {
            try {
                session.close();
            } catch (IOException e) {
                log.error("Websocket session close exception ", e);
            }
        }
    }

    /**
     * 获取 websocket 会话
     *
     * @param key
     * @return
     */
    public static WebSocketSession get(String key) {
        return WEB_SOCKET_SESSION_MAP.getIfPresent(key);
    }

    /**
     * 发送消息
     * @param key  自定义 key
     * @param data 数据
     */
    public static <T> void sendMessages(String key, T data) {
        String sendData = "";
        if (data != null) {
            try {
                // json 序列化
                sendData = OBJECT_MAPPER.writeValueAsString(data);
            } catch (JsonProcessingException e) {
                log.error("json序列化异常，{}", e.getMessage());
                return;
            }
        }
        // 根据 key 缓存中获取 WebSocketSession
        WebSocketSession session = WEB_SOCKET_SESSION_MAP.getIfPresent(key);
        try {
            if (session != null && session.isOpen()) {
                // 把json消息包装成 TextMessage 进行发送
                session.sendMessage(new TextMessage(sendData));
            }
        } catch (IOException e) {
            log.error("消息发送异常，{}", e.getMessage());
        }
    }


}

4. SocketInterceptor 握手拦截器

/**
 * @Description: websocket 握手拦截器
 *  可以用于全局赋值，进行 token 校验等
 *  可以从 ServerHttpRequest 中获取信息，来进行全局赋值
 * 
 */
@Slf4j
@Component
public class SocketInterceptor extends HttpSessionHandshakeInterceptor {

    private static final String TOKEN_FIELD = "token";

    @Autowired
    StringRedisTemplate redisTemplate;

    /**
     * websocket 握手之前
     *
     * @param serverHttpRequest
     * @param serverHttpResponse
     * @param webSocketHandler
     * @param map   websocket 中的 attributes
     * @return  true 继续执行，false 终端
     * @throws Exception
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest serverHttpRequest, 		                                            ServerHttpResponse serverHttpResponse,                                                    WebSocketHandler webSocketHandler, 
                                   Map<String, Object> map) throws Exception {
        // 获取请求参数 token
        String token = serverHttpRequest.getHeaders().getFirst(TOKEN_FIELD);
        String authorization = HttpUtil.decodeParamMap(serverHttpRequest.getURI().getQuery(), Charsets.UTF_8).get(TOKEN_FIELD);
        String auth = StrUtil.blankToDefault(authorization, token);
        log.debug("websocket 开始握手,token：{}",auth);
        // 业务处理逻辑
        xxxxxxxxxxxxxxxxxxxx
        
        return true;
    }

    

    /**
     * websocket 握手之后
     *
     * @param serverHttpRequest
     * @param serverHttpResponse
     * @param webSocketHandler
     * @param e
     */
    @Override
    public void afterHandshake(ServerHttpRequest serverHttpRequest, 
                               ServerHttpResponse serverHttpResponse,
                               WebSocketHandler webSocketHandler, 
                               Exception e) {
        log.debug("握手完成!");
    }
}

5. SocketAuthHandler 业务逻辑处理类

/**
 * @Description: websocket 业务逻辑处理类
 * Spring提供了TextWebSocketHandler和BinaryWebSocketHandler分别处理文本消息和二进制消息
 */
@Slf4j
@Component
public class SocketAuthHandler extends TextWebSocketHandler {

    /**
     * 握手成功之后 回调方法
     *  往 session 中的 Attribute 添加全局使用的数据
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
//       可以从 session 中的 Attributes 获取数据，添加数据
        UserDto moblieUser = (UserDto) session.getAttributes().get(SystemTool.MOBLIE_USER);
       
//        return;   正常处理
//       throw new BusinessException("用户登录已失效");  阻断流程
    }

    @Autowired
    private ApplicationContext applicationContext;

    /**
     * 接收客户端消息
     *
     * @param session
     * @param message
     * @throws Exception
     *  
     * 前端 websocket 会发送 send Beat 的事件到后端
     * 所以对于前端信息序列化失败后 不能报错直接返回就好，如下
     * 
     * try {
            exchangeData = JSONObject.parseObject(message.getPayload(), ExchangeData.class);
        } catch (Exception e) {
            log.error("序列化失败:{}", e.getMessage());
            return;
        }
     * 
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, 
                                     TextMessage message) throws Exception {
     	// 从 session 中的 Attributes 获取数据
       UserDto pcUser = (UserDto) session.getAttributes().get(SystemTool.PC_USER);
        // message.getPayload() 获取客户端信息 进行事件校验
        // 客户端 发送 ping  服务端回复 pong 保持连接
        if (pcUser != null && "ping".equalsIgnoreCase(message.getPayload())) {
            SocketSessionManager.get(SystemTool.PC_USER + pcUser.getId() + "_" + pcUser.getTenantId());
            
            // 发送信息给客户端
            session.sendMessage(new TextMessage("pong"));
            return;
        }
       
        //接收服务  自定义接收类
        ExchangeData exchangeData;
        try {
            // 反序列化为 实体类
            exchangeData = JSONObject.parseObject(message.getPayload(), ExchangeData.class);
        } catch (Exception e) {
            log.error("序列化失败:{}", e.getMessage());
            return;
        }
        
        // 根据客户端发送的信息进行具体的业务逻辑处理
        xxxxxxxxxxxxxxxxxxxxxxx
        
        
    }

   
	/**
     *  任何原因导致WebSocket连接中断时，Spring会自动调用afterConnectionClosed方法
     * @param session
     * @param status
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        
        UserDto pcUser = (UserDto) session.getAttributes().get(SystemTool.PC_USER);
        if (Objects.nonNull(pcUser)) {
            log.debug("pc用户 [{}] 断开连接", pcUser.getId() + "_" + pcUser.getTenantId());
            // 删除  SocketSession 缓存
            SocketSessionManager.removeAndClose(SystemTool.MOBLIE_USER + pcUser.getId() + "_" + pcUser.getTenantId());
        }
    }

    /**
    * 消息传输出错时调用
    */
    @Override
    public void handleTransportError(WebSocketSession session,
                                     Throwable exception) throws Exception {
        // 消息传输出错时调用
        log.error("socket session{}", session);
    }
}

6. 前端开发

// 创建WebSocket连接:
var ws = new WebSocket('ws://' + location.host + '/chat');
// 连接成功时:
ws.addEventListener('open', function (event) {
    console.log('websocket connected.');
});
// 收到消息时:
ws.addEventListener('message', function (event) {
    console.log('message: ' + event.data);
    var msgs = JSON.parse(event.data);
    // TODO:
});
// 连接关闭时:
ws.addEventListener('close', function () {
    console.log('websocket closed.');
});
// 绑定到全局变量:
window.chatWs = ws;

// websocket 链接成功后，客户端在任务地方都可以调用
var inputText = 'Hello, WebSocket.';
window.chatWs.send(JSON.stringify({text: inputText}));