Spring Boot 整合 Websocket

最近有个需求是关于扫码的，用户A提供码让用户B来扫，用户B扫了之后就要给用户A下达一个成功的提醒，但是传统的HTTP无法在服务端推送信息给特定用户，于是在调研了一些方法后选择了Websocket实现这个功能，简要记录本次的实现。

为什么使用Websocket

传统的HTTP只能由客户端发起连接，服务端响应请求返回数据。如果需要服务端主动向客户端推送内容，那就比较麻烦了，但是在实际当中这样的应用场景有很多，比如说新的微博、好友动态更新，以及本次需求的被扫码成功通知等。传统的做法是通过轮询实现，即客户端每隔一段时间发送请求获取是否存在更新，通过短时间间隔不断请求模拟出实时收到服务端更新。这个方法非常简单，但是存在一定的性能问题，毕竟短时间内频繁请求服务端必定会造成压力，而时间间隔过长又影响了实时性，于是又出现了长轮询（Comet），服务端在接收到请求后，会阻塞请求直到有数据或者超时才返回，客户端处理返回信息后再次发出请求，但是长轮询仍然存在资源浪费。于是选择具备客户端/服务端双向通信的WebSocket。

引入依赖

以Maven为例。

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

配置

引入依赖后少不了的就是配置了，新建一个类并实现WebSocketConfigurer接口，类上添加@Configuration和@EnableWebSocket两个注解。

@Configuration
@EnableWebSocket
public class WebSocketConfiguration implements WebSocketConfigurer {

    @Autowired
    private SignInWebSocketHandler signInWebSocketHandler;
    @Autowired
    private SignInHandshakeInterceptor signInHandshakeInterceptor;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(signInWebSocketHandler, "/url")
                .addInterceptors(signInHandshakeInterceptor).setAllowedOrigins("*");
    }
}

拦截器

和普通的HTTP请求的拦截器类似，我们可以在Websocket连接建立的时候，通过配置拦截器进行鉴权，避免他人恶意连接占用服务器资源。实现HandshakeInterceptor接口的类即可作为拦截器，在beforeHandshake方法中进行鉴权，此外还可以将鉴定完毕的用户信息保存在attributes这个Map当中，供之后的方法使用。

public class SignInHandshakeInterceptor implements HandshakeInterceptor {

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        if (request instanceof ServletServerHttpRequest) {
            // 重点：首先将request强制转换为ServletServerHttpRequest
            ServletServerHttpRequest req = (ServletServerHttpRequest) request;
            String session = req.getServletRequest().getParameter("session");
            if(session == null){
                // 拒绝
                return false;
            }
            // ... 其他逻辑
            attributes.put("Session",session);
            return true;
        }
        return false;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {

    }
}

处理Websocket业务逻辑

当请求通过拦截器进入Handler之后，就可以开始处理业务逻辑了，实现WebSocketHandler接口（或继承他的实现类，如TextWebSocketHandler），之后可根据自己的需要填充业务逻辑。

public class SignInWebSocketHandler extends TextWebSocketHandler {
    private final ConcurrentHashMap<String,WebSocketSession> userSessionMap;

    public SignInWebSocketHandler(){
        userSessionMap = new ConcurrentHashMap<>();
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // ... 连接建立成功
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // ... 处理对方发来的文本消息
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        // ... 连接关闭
    }

}

向特定客户端发送信息

下面来到重点，即向特定的客户端主动发送信息，这也是传统HTTP协议难以实现的一点。这里，我使用WebSocketSession的其中一个方法getAttributes()以获取到之前拦截器设置的session信息，将这个信息和WebSocketSession存入Map中即可在之后的逻辑中取出并发送信息，这一操作的核心方法是sendMessage。

    private void sendMsg(WebSocketSession session, SignInMessageModel signInMessageModel){
        try {
            if(session.isOpen()){
                session.sendMessage(new TextMessage(new Gson().toJson(signInMessageModel)));
            }
        } catch (IOException e) {
            log.error("WebSocket IO Error",e);
        }
    }

与Nginx配合使用的一个坑

在本地测试没有问题，但是线上测试环境部署后无法建立Websocket连接，日志提示错误：

Handshake failed due to invalid Upgrade header: null

查找资料后发现是nginx反向代理的时候，忽略了headers中的：Upgrade:websocket，因此需要进行以下配置，重启nginx。

# WebSocket support (nginx 1.4)
proxy_http_version 1.1;
proxy_set_header Upgrade $http_upgrade;
proxy_set_header Connection "upgrade";