在springboot中使用websocket

1、添加如下依赖

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

2、定义一个类，继承TextWebSocketHandler，并重写如下方法

@Component
public class MessageHandler extends TextWebSocketHandler {

    private Map<Long,WebSocketSession> SESSIONS=new HashMap<>();

    private ObjectMapper objectMapper=new ObjectMapper();

    @Autowired
    private MessageDAO messageDAO;

    //客户端与服务端建立websocket连接后执行
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Long uid = (Long) session.getAttributes().get("uid");
        SESSIONS.put(uid,session);
    }

    //服务端接收到消息后执行
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage textMessage) throws Exception {
        Long uid = (Long) session.getAttributes().get("uid");
        JsonNode jsonNode = objectMapper.readTree(textMessage.getPayload());
        long toId = jsonNode.get("toId").asLong();
        String msg = jsonNode.get("msg").asText();
        Message message = Message.builder()
                .from(UserData.USER_MAP.get(uid))
                .to(UserData.USER_MAP.get(toId))
                .msg(msg)
                .build();
        message = messageDAO.saveMessage(message);
        WebSocketSession toSession = SESSIONS.get(toId);
        if (toSession!=null&&toSession.isOpen()){
            toSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(message)));
            messageDAO.updateMessageState(message.getId(),2);
        }
    }

    //客户端与服务端关闭websocket连接后执行
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Long uid = (Long) session.getAttributes().get("uid");
        SESSIONS.remove(uid);
    }
}

其中，afterConnectionEstablished方法是在客户端与服务端建立websocket连接后执行，handleTextMessage方法是服务端接收到消息后执行，通过WebSocketSession 类的对象调用sendMessage方法可以向客户端发送消息，afterConnectionClosed方法是客户端与服务端关闭websocket连接后执行。

3、定义一个类继承HandshakeInterceptor接口，作为拦截器，并重写如下方法

@Component
public class MessageHandShakeInterceptor implements HandshakeInterceptor {
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        String path = request.getURI().getPath();
        String[] split = StringUtils.split(path, '/');
        if (split.length!=2){
            return false;
        }
        if (!StringUtils.isNumeric(split[1])){
            return false;
        }
        attributes.put("uid",Long.valueOf(split[1]));
        return true;
    }

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

    }
}

其中，beforeHandshake方法是在握手前执行，也就是在websocket连接建立之前执行，其中的attributes参数可以用来设置WebSocketSession 对象中的属性，afterHandshake方法是在握手成功后执行。

4、定义一个类继承WebSocketConfigurer接口，并重写如下方法，添加如下注解

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Autowired
    private MessageHandler messageHandler;

    @Autowired
    private MessageHandShakeInterceptor messageHandShakeInterceptor;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(messageHandler,"/ws/{uid}")
                .setAllowedOriginPatterns("*")
                .addInterceptors(messageHandShakeInterceptor);
    }
}

其中，registerWebSocketHandlers方法里面可以设置websocket服务的访问路径，设置跨域，并添加拦截器。