使用WebSocket

在Servlet 3.0时代，WebSocket相关的支持被引入Servlet API中，Spring基于此为我们进行了封装，提供了方便的WebSocketHandler组件来快速开发WebSocket服务端程序。这篇笔记我们主要介绍SpringBoot 3.x工程中如何实现WebSocket相关的功能。

引入Maven依赖

对于SpringBoot工程，使用WebSocket需要引入起步依赖。

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

工程配置

工程中我们需要在启动类上标注@EnableWebSocket来开启WebSocket支持。

WsConfig.java

package com.gacfox.demo;

import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.web.socket.config.annotation.EnableWebSocket;

@EnableWebSocket
@SpringBootApplication
public class DemoApplication {

    public static void main(String[] args) {
        SpringApplication.run(DemoApplication.class, args);
    }

}

编写WebSocketHandler

Spring的WebSocket封装提供了WebSocketHandler接口。具体代码中，我们可以继承抽象类AbstractWebSocketHandler，或者更具体的直接继承TextWebSocketHandler或BinaryWebSocketHandler，后两个类分别专用于处理文本消息和二进制消息，足够覆盖绝大部分WebSocket的使用场景了。

这里我们继承TextWebSocketHandler，实现一个最简单的文本消息处理功能，接收客户端发送的输入并打印，然后回复“收到”。为了处理WebSocket消息，我们需要编写我们自己的Handler类。

ChatHandler.java

package com.gacfox.demo.ws;

import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

@Slf4j
public class ChatHandler extends TextWebSocketHandler {
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        log.info(message.getPayload());
        session.sendMessage(new TextMessage("收到"));
    }
}

例子代码比较简单，我们直接继承了TextWebSocketHandler来处理文本消息。其中session参数包含连接的信息，例如需要实现消息广播等功能时，我们就需要保存所有的连接并统一发送消息；message则是具体的文本消息体。具体的Handler逻辑就是直接打印客户端发来的文本消息。

此时我们的Handler类还不能直接使用，我们需要将其注册到框架中，下面代码我们将ChatHandler注册到/ws/chat路径。

package com.gacfox.demo.config;

import com.gacfox.demo.ws.ChatHandler;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
public class WsConfig implements WebSocketConfigurer {
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(chatHandler(), "/ws/chat");
    }

    @Bean
    public ChatHandler chatHandler() {
        return new ChatHandler();
    }
}

跨域设置

虽然WebSocket建立连接后不再是简单的HTTP请求响应了，但它的第一步仍是HTTP Upgrade请求，因此仍存在跨域安全限制。如果我们有跨域需求，需要在注册Handler时进行设置。

@Override
public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
    registry.addHandler(chatHandler(), "/ws/chat")
            .setAllowedOriginPatterns("*");
}

这里我们直接设置了*即允许全部跨域连接，生产环境中这样做并不安全，实际使用时应该根据具体的限制写出明确的跨域域名。

HandshakeInterceptor握手拦截器

除了Handler类，Spring的WebSocket封装还提供了HandshakeInterceptor接口，它用于对WebSocket建立握手连接前和建立握手连接后进行AOP风格的拦截处理。HandshakeInterceptor接口有如下两个重要方法。

package org.springframework.web.socket.server;

import java.util.Map;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.lang.Nullable;
import org.springframework.web.socket.WebSocketHandler;

public interface HandshakeInterceptor {
    boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception;

    void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, @Nullable Exception exception);
}

beforeHandshake：建立握手连接前处理，返回true继续连接，返回false则中断

afterHandshake：建立握手连接后处理

前面说过WebSocket建立连接的第一步是HTTP Upgrade请求，HandshakeInterceptor就是在这一步起作用的，我们可以看到参数中都是ServerHttpRequest、ServerHttpResponse等HTTP协议请求和响应的对象，我们可以借此进行建立WebSocket连接前的认证、鉴权等操作。

下面例子中我们编写了一个拦截器，并进行注册。

ChatInterceptor.java

package com.gacfox.demo.interceptor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

@Slf4j
public class ChatInterceptor implements HandshakeInterceptor {
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        log.info("建立连接前");
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        log.info("建立连接后");
    }
}

WsConfig.java

package com.gacfox.demo.config;

import com.gacfox.demo.interceptor.ChatInterceptor;
import com.gacfox.demo.ws.ChatHandler;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
public class WsConfig implements WebSocketConfigurer {
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(chatHandler(), "/ws/chat")
                .addInterceptors(new ChatInterceptor());
    }

    @Bean
    public ChatHandler chatHandler() {
        return new ChatHandler();
    }
}

代码非常简单，这里就不多解释了。

作者：Gacfox