简介

Maven依赖

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

方案一：Spring原生注解

1.添加配置类

@Configuration
public class WebSocketConfiguration {
    /**
     * 	注入ServerEndpointExporter，
     * 	这个bean会自动注册使用了@ServerEndpoint注解声明的Websocket endpoint
     */
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

2.创建EndPoint类

@ServerEndpoint("/ws/{userPk}")
@Slf4j
public class WSServerEndpoint {
    private static int onlineCount = 0;
    // 保存在线会话,有需要可以存在redis中,同时考虑用sessionId作为会话标识，适配同一个账号打开多个websocket连接
    private static ConcurrentHashMap<String, Session> webSocketMap = new ConcurrentHashMap<>();

    /**
     * 与某个客户端的连接会话，需要通过它来给客户端发送数据
     */
    private Session session;
    /**
     * 接收userName
     */
    private String userName = "";

    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("userPk") String userPk) {
        if (!webSocketMap.containsKey(userPk)) {
            addOnlineCount(); // 在线数 +1
        }
        this.session = session;
        this.userName = userPk;
        webSocketMap.put(userName, session);
        log.debug("用户连接:" + userName + ",当前在线人数为:" + getOnlineCount());
    }

    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public void onClose() {
        if (webSocketMap.containsKey(userName)) {
            webSocketMap.remove(userName);
            if (webSocketMap.size() > 0) {
                subOnlineCount();
            }
        }
        log.debug(userName + "用户退出,当前在线人数为:" + getOnlineCount());
    }

    /**
     * 收到客户端消息后调用的方法
     *
     * @param message 客户端发送过来的消息
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        log.debug("收到用户消息:" + userName + ",报文:" + message);
        JSONObject json = JSONObject.parseObject(message);
        String msgHandler = json.getString("msgHandler");
        if (StringUtils.isEmpty(msgHandler)){
            log.warn("message未指定handler类，无法执行业务逻辑。");  
            return;
        }
        Map<String, IWSMessageHandler> beans = SpringContextUtils.getBeans(IWSMessageHandler.class);
        IWSMessageHandler handler = beans.get(json.getString("msgHandler"));
        if (handler == null){
            log.warn("message指定handler类不存在，请确认是否已注册到Spring容器。");
            return;
        }
        handler.handlerMessage(json.get("body"));
    }

    /**
     * @param session
     * @param error
     */
    @OnError
    public void onError(Session session, Throwable error) {
        log.error("用户错误:" + this.userName + ",原因:" + error.getMessage());
    }
    
    public static void sendMessage(Session session, String message) {
		try {
			System.out.println("向sid为：" + session.getId() + "，发送：" + message);
			session.getBasicRemote().sendText(message);
		} catch (IOException e) {
			throw new RuntimeException(e);
		}
	}
	
	// 向指定用户推送消息 
	public static void sendMessage(long userId, String message) {
		Session session = webSocketMap.get(USER_ID + userId);
		if(session != null) {
			sendMessage(session, message);
		}
	}

    public static synchronized int getOnlineCount() {
        return onlineCount;
    }

    public static synchronized void addOnlineCount() {
        WSServerEndpoint.onlineCount++;
    }

    public static synchronized void subOnlineCount() {
        WSServerEndpoint.onlineCount--;
    }

    public static synchronized ConcurrentHashMap<String, Session> getWebSocketMap(){
        return webSocketMap;
    }

}

3. 身份认证

在onOpen方法中进行token信息进行鉴权，若认证失败可以直接抛出异常，拒绝连接。onOpen是在握手成功后创建连接时调用的，在这里鉴权稍显滞后，可以考虑用方案二自行封装WebSocket实现在握手时进行鉴权。

方案二：Spring封装WebSocket

参考质料：

Spring Boot 集成 WebSocket（原生注解与Spring封装

使用 Sa-Token 解决 WebSocket 握手身份认证

认证框架sa-token