spring+springMVC+webSocket實現給指定用戶推送消息（一）

本片文章實現的功能是後端程序給指定的用戶通過webSocket推送消息。根據前端用戶登錄時給後臺傳遞用戶ID，然後在後臺將用戶ID和對應的webSocketSession存儲起來，用於推送消息時使用。

一、pom.xml文件中加入maven依賴

<dependency>
	<groupId>org.springframework</groupId>
	<artifactId>spring-websocket</artifactId>
	<version>4.3.9.RELEASE</version>
</dependency>
<dependency>
	<groupId>org.springframework</groupId>
	<artifactId>spring-webmvc</artifactId>
	<version>4.3.9.RELEASE</version>
</dependency>
<dependency>
	<groupId>org.springframework</groupId>
	<artifactId>spring-context</artifactId>
	<version>4.3.9.RELEASE</version>
</dependency>

二、websocket配置類

package framework.websocket;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurerAdapter;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
@EnableWebSocket
public class WebSocketConfig extends WebMvcConfigurerAdapter implements WebSocketConfigurer {

    /**
     * 注入websocket處理類
     */
    @Autowired
    private WebSocketHandler webSocketHandler;

    /**
     * 註冊SocketHandler
     */
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // 原生支持websocket的使用地址/webSocketServer
        registry.addHandler(webSocketHandler, "/webSocketService").addInterceptors(
                new WebsocketHandshakeInterceptor()).setAllowedOrigins("*");

        // 不支持websocket的使用sockjs，地址/webSocketServer/sockjs
        registry.addHandler(webSocketHandler, "/webSocketService/sockjs").setAllowedOrigins("*")
                .addInterceptors(new WebsocketHandshakeInterceptor()).withSockJS();
    }
}

三、websocket握手攔截器（握手請求和響應，並將屬性傳遞給目標）

package framework.websocket;

import java.util.Map;
import java.util.logging.Logger;

import javax.servlet.http.HttpServletRequest;

import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

public class WebsocketHandshakeInterceptor implements HandshakeInterceptor {

    /**
     * 獲取日誌對象（全類名）
     */
    private static final Logger logger = Logger.getLogger(WebsocketHandshakeInterceptor.class.getName());

    /** 
     * 在處理握手之前調用，這裏是把獲取的請求數據綁定到session的map對象中（attributes）
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
            WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {

        // 將ServerHttpRequest轉成ServletServerHttpRequest
        ServletServerHttpRequest servletServerHttpRequest = (ServletServerHttpRequest) request;

        // 獲取請求參數，先獲取HttpServletRequest對象才能獲取請求參數
        HttpServletRequest httpServletrequest = servletServerHttpRequest.getServletRequest();
        logger.info(String.format("Websocket Handshake Interceptor， sessionID：%s",
                httpServletrequest.getSession().getId()));

        // 獲取請求的數據
        String clientID = httpServletrequest.getParameter(WebSocketConstant.CLIENT_ID);
        if (null == clientID || "".equals(clientID)) {
            return false;
        }
        
        // 把獲取的請求數據綁定到session的map對象中（attributes）
        attributes.put(WebSocketConstant.CLIENT_ID, clientID);
        return true;
    }

    /* 
     * 握手完成後調用。響應狀態和頭指示握手的結果，即握手是否成功。
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
            WebSocketHandler wsHandler, Exception exception) {
        // TODO Auto-generated method stub
    }
}

四、webSocket處理類

package framework.websocket;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.logging.Logger;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

@Component
public class WebSocketHandler extends TextWebSocketHandler {

    /**
     * 獲取日誌對象（全類名）
     */
    private static final Logger logger = Logger.getLogger(WebSocketHandler.class.getName());

    /**
     * 在線用戶列表
     */
    private static final Map<String, List<WebSocketSession>> onlineUsers;

    static {
        onlineUsers = Collections.synchronizedMap(new HashMap<>());
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {

        // 獲取WebSocket客戶端ID
        String clientID = getWebSocketClientID(session);

        List<WebSocketSession> webSocketSessionList = onlineUsers.get(clientID);
        if (null != webSocketSessionList) {
            webSocketSessionList.add(session);
        } else {
            webSocketSessionList = new ArrayList<>();
            webSocketSessionList.add(session);

            onlineUsers.put(clientID, webSocketSessionList);
        }
        recordLog("Websocket connection successful， wsClientID：%s，wsSessionID：%s。",
                clientID, session.getId());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status)
            throws Exception {

        // 獲取WebSocket客戶端ID
        String clientID = getWebSocketClientID(session);

        List<WebSocketSession> webSocketSessionList = onlineUsers.get(clientID);
        if (null != webSocketSessionList) {
            webSocketSessionList.remove(session);
        }
        recordLog("Websocket disconnected successfully， wsClientID：%s，wsSessionID：%s。",
                clientID, session.getId());
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) {

        // 獲取WebSocket客戶端ID
        String clientID = getWebSocketClientID(session);
        recordLog("Websocket received the message successfully， wsClientID：%s，wsSessionID：%s。",
                clientID, session.getId());

        // 將前端接收的消息回發給客戶端
        sendMessage(clientID, session.getId(), message);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {

        // 獲取WebSocket客戶端ID
        String clientID = getWebSocketClientID(session);

        List<WebSocketSession> webSocketSessionList = onlineUsers.get(clientID);
        if (null != webSocketSessionList) {
            webSocketSessionList.remove(session);
        }
        recordLog("Websocket handles connection exceptions， wsClientID：%s，wsSessionID：%s。",
                clientID, session.getId());
    }

    /**
     * 獲取WebSocket客戶端ID
     * @param session WebSocketSession
     * @return WebSocket客戶端ID
     */
    private String getWebSocketClientID(WebSocketSession session) {
        return (String) session.getAttributes().get(WebSocketConstant.CLIENT_ID);
    }

    /**
     * 發送消息
     * @param clientID 客戶端ID
     * @param textMessage 消息對象
     */
    private static void sendMessage(String clientID, String sessionID, TextMessage textMessage) {
        // 根據客戶端ID查詢webSocketSession集合
        List<WebSocketSession> webSocketSessionList = onlineUsers.get(clientID);

        if (null == webSocketSessionList) {
//            recordLog("Not find webSocketSession by clientID， wsClientID：%s。", clientID);
            return;
        }

        for (WebSocketSession webSocketSession : webSocketSessionList) {
            // 得到webSocketSessionID
            String webSocketSessionID = webSocketSession.getId();

            boolean flag = false;
            // 當sessionID爲null時，表示給所有根據clientID查詢出來的客戶端都發送；當sessionID不爲null時，表示只給和sessionID相同的客戶端單獨發送；
            if (null == sessionID || sessionID.equals(webSocketSessionID)) {
                flag = true;
            }

            // 判斷連接是否仍然處於連接狀態
            if (webSocketSession.isOpen()) {
                if(flag) {
                    try {
                        webSocketSession.sendMessage(textMessage);
                        recordLog("Send webSocket message success， wsClientID：%s，wsSessionID：%s。",
                                clientID, webSocketSessionID);
                    } catch (IOException e) {
                        recordLog(
                                "Send webSocket message exception，wsClientID：%s，wsSessionID：%s，exceptionMessage：%s。",
                                clientID, webSocketSessionID, e.getMessage());
                        e.printStackTrace();
                    }
                }
            } else {
                recordLog(
                        "Send webSocket message faile，client state is disconnected，wsClientID：%s，wsSessionID：%s。",
                        clientID, webSocketSessionID);
            }
        }
    }

    /**
     * 發送信息
     * @param clientID 客戶端ID
     * @param message 消息
     */
    public static void sendMessage(String clientID, String message) {
        sendMessage(clientID, null, new TextMessage(message));
    }
    
    /**
     * 記錄日誌
     * @param message 日誌信息
     * @param args 日誌信息中的參數值
     */
    private static void recordLog(String message, Object... args) {
        logger.info(String.format(message, args));
    }
}

五、websocket常量類

package framework.websocket;

public final class WebSocketConstant {

    /**
     * websocket客戶端ID常量
     */
    public static final String CLIENT_ID = "clientID";

}