WebSocket中间件实现

WebSocket中间件实现

1. 背景

当集成spring-boot-starter-websocket来做websocket逻辑时，我们需要考虑会话存储，单域名多节点服务时，如何找到建立连接的会话上进行通讯等问题。

对此，这里对spring-boot-starter-websocket进行封装，在原来的基础上增加连接鉴权、会话存储、多节点下如何通知原来连接上的节点进行通讯。

这里主要讲思路，代码实现因为时间问题有待完善。

2. 思路图

3. 代码实现

1. SpringBoot的自动配置

引入依赖

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

<!-- springboot配置属性解析 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-configuration-processor</artifactId>
    <!-- 不传递依赖 -->
    <optional>true</optional>
</dependency>

在resources目录下新建META-INF文件夹，并创建spring.factories文件，文件内容为需要读取的配置类

org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
  com.mountain.websocket.autoconfigure.WebSocketAutoConfig

2. 配置类

配置类主要实现WebSocketConfigurer接口的registerWebSocketHandlers()方法为websocket处理器和拦截器注册，并标注@EnableWebSocket注解

WebSocketAutoConfig

/**
 * websocket配置类
 * @author Tarzan写bug
 * @since 2023/02/14
 */
@Configuration
@EnableWebSocket
@ComponentScan(basePackages = "com.mountain.websocket")
public class WebSocketAutoConfig implements WebSocketConfigurer {

    private final WebSocketConfig config;

    public WebSocketAutoConfig(WebSocketConfig config) {
        this.config = config;
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // WebSocket通道
        registry.addHandler(new WebSocketHandler(), config.getWebSocketUrl())
                .addInterceptors(new WebSocketInterceptor())
                .setAllowedOrigins("*");
        // sockJs通道
        registry.addHandler(new WebSocketHandler(), config.getSockJsUrl())
                .addInterceptors(new WebSocketInterceptor())
                .setAllowedOrigins("*");
    }
}

3. 处理器

处理器类继承AbstractWebSocketHandler，实现建立连接、处理消息、异常处理等方法，这里通过策略模式将不同类型的请求使用不同的处理器

WebSocketHandler

/**
 * websocket处理器
 * @author Tarzan写bug
 * @since 2023/02/14
 */
public class WebSocketHandler extends AbstractWebSocketHandler {

    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketHandler.class);
    private final Map<String, SocketHandler> handlers = new HashMap<>();

    private Map<String, SocketHandler> getHandlers() {
        if (handlers.size() > 0) {
            return handlers;
        }

        Map<String, SocketHandler> map = ApplicationContextHelper.getContext()
                .getBeansOfType(SocketHandler.class);
        map.values().forEach(item -> {
            handlers.put(item.name(), item);
        });

        return handlers;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String processor = getProcessor(session.getAttributes());
        if (!StringUtils.isEmpty(processor) && handlers.containsKey(processor)) {
            handlers.get(processor).afterConnectionEstablished(session);
        }

        if (session.isOpen()) {
            handlers.get(WebSocketConstant.DEFAULT_PROCESSOR_NAME).afterConnectionEstablished(session);
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        String sessionId = getSessionId(session);
        try {
            String processor = getProcessor(session.getAttributes());
            if (!StringUtils.isEmpty(processor) && handlers.containsKey(processor)) {
                handlers.get(processor).handleTransportError(session, exception);
            }
        } finally {
            // 清理缓存
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String sessionId = getSessionId(session);
        try {
            String processor = getProcessor(session.getAttributes());
            if (!StringUtils.isEmpty(processor) && handlers.containsKey(processor)) {
                handlers.get(processor).afterConnectionClosed(session, status);
            }
        } finally {
            // 清理缓存
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String processor = getProcessor(session.getAttributes());
        if (!StringUtils.isEmpty(processor) && handlers.containsKey(processor)) {
            handlers.get(processor).handleTextMessage(session, message);
        }
        handlers.get(WebSocketConstant.DEFAULT_PROCESSOR_NAME).handleTextMessage(session, message);
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        String processor = getProcessor(session.getAttributes());
        if (!StringUtils.isEmpty(processor) && handlers.containsKey(processor)) {
            handlers.get(processor).handleBinaryMessage(session, message);
        }
        handlers.get(WebSocketConstant.DEFAULT_PROCESSOR_NAME).handleBinaryMessage(session, message);
    }

    @Override
    protected void handlePongMessage(WebSocketSession session, PongMessage message) throws Exception {
        String processor = getProcessor(session.getAttributes());
        if (!StringUtils.isEmpty(processor) && handlers.containsKey(processor)) {
            handlers.get(processor).handlePongMessage(session, message);
        }
        handlers.get(WebSocketConstant.DEFAULT_PROCESSOR_NAME).handlePongMessage(session, message);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 获取sessionId
     * @param session
     * @return
     */
    private String getSessionId(WebSocketSession session) {
        String sessionId = null;
        try {
            if (session instanceof StandardWebSocketSession) {
                sessionId = session.getId();
            } else if (session instanceof WebSocketServerSockJsSession) {
                sessionId = ((WebSocketSession) FieldUtils.readField(session, "webSocketSession", true)).getId();
            }
            return sessionId;
        } catch (Exception e) {
            LOGGER.error("getSessionId error");
            return null;
        }
    }

    /**
     * 获取WebSocket连接的请求头参数processor
     * @param attributes
     * @return
     */
    private String getProcessor(Map<String, Object> attributes) {
        String processor = null;
        if (attributes != null
                && attributes.containsKey(WebSocketConstant.Attributes.PROCESSOR)) {
            processor = String.valueOf(attributes.get(WebSocketConstant.Attributes.PROCESSOR));
        }
        return processor;
    }
}

4. 拦截器

实现HandshakeInterceptor接口，同样可以用策略模式实现不同请求使用不同的拦截器，拦截器的主要功能可以对请求进行令牌的校验等。

WebSocketInterceptor

/**
 * websocket拦截器
 * @author Tarzan写bug
 * @since 2023/02/14
 */
public class WebSocketInterceptor implements HandshakeInterceptor {

    private final Map<String, SocketIntereptor> interceptors = new HashMap<>();

    private Map<String, SocketIntereptor> getInterceptors() {
        if (interceptors.size() > 0) {
            return interceptors;
        }

        Map<String, SocketIntereptor> map = ApplicationContextHelper.getContext().getBeansOfType(SocketIntereptor.class);
        map.values().forEach(item -> {
            interceptors.put(item.name(), item);
        });
        return interceptors;
    }

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler handler, Map<String, Object> map) throws Exception {
        return interceptors.get(WebSocketConstant.DEFAULT_PROCESSOR_NAME).beforeHandshake(request, response, handler, map);
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler handler, Exception e) {
        interceptors.get(WebSocketConstant.DEFAULT_PROCESSOR_NAME).afterHandshake(request, response, handler, e);
    }
}

5. 会话缓存和多实例接收请求

这两部分功能暂时没在代码中实现，这里说一下思路；

会话缓存在建立连接的时候先缓存到本地内存，比如用Map来保存，再保存在redis中。

多实例接收请求是当服务部署了多个实例后，建立连接的实例A缓存了sessionA，但当请求路由到实例B时，要求实例B向sessionA发送消息。这时候实例B从内存中找不到sessionA，即可以通过redis的订阅发布功能通知其他实例发送消息，当实例A收到redis的消息后找到本地有sessionA，即发送websocket消息。

4. 总结

代码收录于https://gitee.com/ouwenrts/tuyere.git的common-websocket目录中。

好了，到站下车，未来可期…

欢迎同频共振的那一部分人

作者公众号：Tarzan写bug

taobao：Tarzan小店，提供一元解决Java问题和其他方面的解决方案，欢迎咨询