WebSocket的故事（五）—— Springboot中，實現網頁聊天室之自定義訊息代理

概述

WebSocket的故事系列計劃分五大篇六章，旨在由淺入深的介紹WebSocket以及在Springboot中如何快速構建和使用WebSocket提供的能力。本系列計劃包含如下幾篇文章：

第一篇，什麼是WebSocket以及它的用途
第二篇，Spring中如何利用STOMP快速構建WebSocket廣播式訊息模式
第三篇，Springboot中，如何利用WebSocket和STOMP快速構建點對點的訊息模式(1)
第四篇，Springboot中，如何利用WebSocket和STOMP快速構建點對點的訊息模式(2)
第五篇，Springboot中，實現網頁聊天室之自定義WebSocket訊息代理
第六篇，Springboot中，實現更靈活的WebSocket

本篇的主線

本篇將通過一個接近真實的網頁聊天室Demo，來詳細講述如何利用WebSocket來實現一些具體的產品功能。本篇將只採用WebSocket本身，不再使用STOMP等這些封裝。親自動手實現訊息的接收、處理、傳送以及WebSocket的會話管理。這也是本系列的最重要的一篇，不管你們激不激動，反正我是激動了。下面我們就開始。

本篇適合的讀者

想了解如何在Springboot上自定義實現更為複雜的WebSocket產品邏輯的同學以及各路有志青年。

小小網頁聊天室的需求

為了能夠目標明確的表達本文中所要講述的技術要點，我設計了一個小小聊天室產品，先列出需求，這樣大家在看後面的實現時能夠知其所以然。

以上就是我們本篇要實現的需求。簡單說，就是：

使用者可加入，退出某房間，加入後可向房間內所有人傳送訊息，也可向某個人傳送悄悄話訊息。

需求分析和設計

設計使用者儲存

很容易想到我們設計的主體就是使用者、會話和房間，那麼在使用者管理上，我們就可以用下面這個圖來表示他們之間的關係：

這樣我們就可以用一個簡單的Map來儲存房間<->使用者組這樣的對映關係，在使用者組內我們再使用一個Map來儲存使用者名稱<->會話Session這樣的對映關係（假設沒有重名）。這樣，我們就解決了房間和使用者組、使用者和會話，這些關係的儲存和維護。

設計使用者行為與使用者的關係

有兄弟看到這說了，“你講這麼半天了，跟之前幾篇講的什麼STOMP，什麼訊息代理，有毛線的關係？”大兄弟你先消消氣，我們學STOMP，學訊息代理，學點對點訊息，重要的是學思想，你說對不？下面我們就用上了。

當使用者加入到某房間之後，房間裡有任何風吹草動，即有人加入、退出或者發公屏訊息，都會“通知”給該使用者。到此，我們就可以將建立房間理解成“建立訊息代理”，將使用者加入房間，看成是對房間這個“訊息代理”的一個“訂閱”，將使用者退出房間，看成是對房間這個“訊息代理”的一個“解除訂閱”。

那麼，第一個加入房間的人，我們定義為“建立房間”，即建立了一個訊息代理。為了好理解，上圖：

其中紅色的小人表示第一個加入房間的使用者，即建立房間的人。當某使用者傳送訊息時，如果選擇將訊息傳送給聊天室的所有人，即相當於在房間裡傳送了一個廣播，所有訂閱這個房間的使用者，都會收到這個廣播訊息；如果選擇傳送悄悄話，則只將訊息傳送給特定使用者名稱的使用者，即點對點訊息。

總結一下我們要實現的要點：

• 使用者儲存，即使用者，房間，會話之間的關係和物件訪問方式。
• 動態建立訊息代理（房間），並實現使用者對房間的繫結（訂閱）。
• 單獨傳送給某個使用者訊息的能力。

大體設計就到此為止，還有一些細節，我們先來看一下演示效果，再來看通過程式碼來講解實現。

聊天室效果展示

用瀏覽器開啟客戶端頁面後，展示輸入框和加入按鈕。輸入房間號1和使用者名稱小銘，點選進入房間。

進入房間成功後，展示當前房間人數和歡迎語。

當有其他人加入或退出房間時，展示通知資訊。可以傳送公屏訊息和私聊訊息。

下面就讓我們看一下這些主要功能如何來實現吧。

程式碼實現

按照我們上述的設計，我會著重介紹重點部分的程式碼設計和技術要點。

服務端實現

1. 配置WebSocket

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(new MyHandler(), "/webSocket/{INFO}").setAllowedOrigins("*")
                .addInterceptors(new WebSocketInterceptor());
    }
}
複製程式碼

要點解析：

• 註冊WebSocketHandler（MyHandler），這是用來處理WebSocket建立以及訊息處理的類，後面會詳細介紹。
• 註冊WebSocketInterceptor攔截器，此攔截器用來在客戶端向服務端發起初次連線時，記錄客戶端攔截資訊。
• 註冊WebSocket地址，並附帶了{INFO}引數，用來註冊的時候攜帶使用者資訊。

以上都會在後續程式碼中詳細介紹。

2. 實現握手攔截器

public class WebSocketInterceptor implements HandshakeInterceptor {
    @Override
    public boolean beforeHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Map<String, Object> map) throws Exception {
        if (serverHttpRequest instanceof ServletServerHttpRequest) {
            String INFO = serverHttpRequest.getURI().getPath().split("INFO=")[1];
            if (INFO != null && INFO.length() > 0) {
                JSONObject jsonObject = new JSONObject(INFO);
                String command = jsonObject.getString("command");
                if (command != null && MessageKey.ENTER_COMMAND.equals(command)) {
                    System.out.println("當前session的ID="+ jsonObject.getString("name"));
                    ServletServerHttpRequest request = (ServletServerHttpRequest) serverHttpRequest;
                    HttpSession session = request.getServletRequest().getSession();
                    map.put(MessageKey.KEY_WEBSOCKET_USERNAME, jsonObject.getString("name"));
                    map.put(MessageKey.KEY_ROOM_ID, jsonObject.getString("roomId"));
                }
            }
        }
        return true;
    }
}

複製程式碼

要點解析：

• HandshakeInterceptor用來攔截客戶端第一次連線服務端時的請求，即客戶端連線/webSocket/{INFO}時，我們可以獲取到對應INFO的資訊。
• 實現beforeHandshake方法，進行使用者資訊儲存，這裡我們將使用者名稱和房間號儲存到Session上。

3. 實現訊息處理器WebSocketHandler

public class MyHandler implements WebSocketHandler {

    //用來儲存使用者、房間、會話三者。使用雙層Map實現對應關係。
    private static final Map<String, Map<String, WebSocketSession>> sUserMap = new HashMap<>(3);

    //使用者加入房間後，會呼叫此方法，我們在這個節點，向其他使用者傳送有使用者加入的通知訊息。
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println("成功建立連線");
        String INFO = session.getUri().getPath().split("INFO=")[1];
        System.out.println(INFO);
        if (INFO != null && INFO.length() > 0) {
            JSONObject jsonObject = new JSONObject(INFO);
            String command = jsonObject.getString("command");
            String roomId = jsonObject.getString("roomId");
            if (command != null && MessageKey.ENTER_COMMAND.equals(command)) {
                Map<String, WebSocketSession> mapSession = sUserMap.get(roomId);
                if (mapSession == null) {
                    mapSession = new HashMap<>(3);
                    sUserMap.put(roomId, mapSession);
                }
                mapSession.put(jsonObject.getString("name"), session);
                session.sendMessage(new TextMessage("當前房間線上人數" + mapSession.size() + "人"));
                System.out.println(session);
            }
        }
        System.out.println("當前線上人數：" + sUserMap.size());
    }

    //訊息處理方法
    @Override
    public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> webSocketMessage) {
        try {
            JSONObject jsonobject = new JSONObject(webSocketMessage.getPayload().toString());
            Message message = new Message(jsonobject.toString());
            System.out.println(jsonobject.toString());
            System.out.println(message + ":來自" + webSocketSession.getAttributes().get(MessageKey.KEY_WEBSOCKET_USERNAME) + "的訊息");
            if (message.getName() != null && message.getCommand() != null) {
                switch (message.getCommand()) {
                        //有新人加入房間資訊
                    case MessageKey.ENTER_COMMAND:
                        sendMessageToRoomUsers(message.getRoomId(), new TextMessage("【" + getNameFromSession(webSocketSession) + "】加入了房間，歡迎！"));
                        break;
                        //聊天資訊
                    case MessageKey.MESSAGE_COMMAND:
                        if (message.getName().equals("all")) {
                            sendMessageToRoomUsers(message.getRoomId(), new TextMessage(getNameFromSession(webSocketSession) +
                                    "說：" + message.getInfo()
                            ));
                        } else {
                            sendMessageToUser(message.getRoomId(), message.getName(), new TextMessage(getNameFromSession(webSocketSession) +
                                    "悄悄對你說：" + message.getInfo()));
                        }
                        break;
                        //有人離開房間資訊
                    case MessageKey.LEAVE_COMMAND:
                        sendMessageToRoomUsers(message.getRoomId(), new TextMessage("【" + getNameFromSession(webSocketSession) + "】離開了房間，歡迎下次再來"));
                        break;
                        default:
                            break;
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 傳送資訊給指定使用者
     */
    public boolean sendMessageToUser(String roomId, String name, TextMessage message) {
        if (roomId == null || name == null) return false;
        if (sUserMap.get(roomId) == null) return false;
        WebSocketSession session = sUserMap.get(roomId).get(name);
        if (!session.isOpen()) return false;
        try {
            session.sendMessage(message);
        } catch (IOException e) {
            e.printStackTrace();
            return false;
        }
        return true;
    }

    /**
     * 廣播資訊給某房間內的所有使用者
     */
    public boolean sendMessageToRoomUsers(String roomId, TextMessage message) {
        if (roomId == null) return false;
        if (sUserMap.get(roomId) == null) return false;
        boolean allSendSuccess = true;
        Collection<WebSocketSession> sessions = sUserMap.get(roomId).values();
        for (WebSocketSession session : sessions) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(message);
                }
            } catch (IOException e) {
                e.printStackTrace();
                allSendSuccess = false;
            }
        }

        return allSendSuccess;
    }

    //退出房間時的處理
    @Override
    public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus closeStatus) {
        System.out.println("連線已關閉：" + closeStatus);
        Map<String, WebSocketSession> map = sUserMap.get(getRoomIdFromSession(webSocketSession));
        if (map != null) {
            map.remove(getNameFromSession(webSocketSession));
        }
    }
}
複製程式碼

要點解析：

• 使用sUserMap這個靜態變數來儲存使用者資訊。對應我們上述的關係圖。
• 實現afterConnectionEstablished方法，當使用者進入房間成功後，儲存使用者資訊到Map，並呼叫sendMessageToRoomUsers廣播新人加入資訊。
• 實現handleMessage方法，處理使用者加入，離開和傳送資訊三類訊息。
• 實現afterConnectionClosed方法，用來處理當使用者離開房間後的資訊銷燬工作。從Map中清除該使用者。
• 實現sendMessageToUser、sendMessageToRoomUsers兩個向客戶端傳送訊息的方法。直接通過Session即可傳送結構化資料到客戶端。sendMessageToUser實現了點對點訊息的傳送，sendMessageToRoomUsers實現了廣播訊息的傳送。

客戶端實現

客戶端我們就使用HTML5為我們提供的WebSocket JS介面即可。

<html>
    <script type="text/javascript">
        function ToggleConnectionClicked() {
            if (SocketCreated && (ws.readyState == 0 || ws.readyState == 1)) {
                lockOn("離開聊天室...");
                SocketCreated = false;
                isUserloggedout = true;
                var msg = JSON.stringify({'command':'leave', 'roomId':groom , 'name': gname,
                    'info':'離開房間'});
                ws.send(msg);
                ws.close();
            } else if(document.getElementById("roomId").value == "請輸入房間號!") {
                Log("請輸入房間號！");
            } else {
                lockOn("進入聊天室...");
                Log("準備連線到聊天伺服器 ...");
                groom = document.getElementById("roomId").value;
                gname = document.getElementById("txtName").value;
                try {
                    if ("WebSocket" in window) {
                        ws = new WebSocket(
                            'ws://localhost:8080/webSocket/INFO={"command":"enter","name":"'+ gname + '","roomId":"' + groom + '"}');
                    }
                    else if("MozWebSocket" in window) {
                        ws = new MozWebSocket(
                            'ws://localhost:8080/webSocket/INFO={"command":"enter","name":"'+ gname + '","roomId":"' + groom + '"}');
                    }
                    SocketCreated = true;
                    isUserloggedout = false;
                } catch (ex) {
                    Log(ex, "ERROR");
                    return;
                }
                document.getElementById("ToggleConnection").innerHTML = "斷開";
                ws.onopen = WSonOpen;
                ws.onmessage = WSonMessage;
                ws.onclose = WSonClose;
                ws.onerror = WSonError;
            }
        };


        function WSonOpen() {
            lockOff();
            Log("連線已經建立。", "OK");
            $("#SendDataContainer").show();
            var msg = JSON.stringify({'command':'enter', 'roomId':groom , 'name': "all",
                'info': gname + "加入聊天室"})
            ws.send(msg);
        };
</html>
複製程式碼

要點解析：

• 發起服務端連線時，注意地址資訊：'ws://localhost:8080/webSocket/INFO={"command":"enter","name":"'+ gname + '","roomId":"' + groom + '"}'，這裡我們在INFO後加入了使用者個人資訊，服務端收到後，即可根據這個資訊標記此會話。
• 連線建立後，傳送給房間內其他人一條加入資訊。通過ws.send()方法實現。

至此程式碼部分就介紹完了，過多的程式碼就不再堆疊了，更詳細的程式碼，請參見後面的Github地址。

本篇總結

通過一個相對完整的網頁聊天室例子，介紹了我們自己使用WebSocket時的幾個細節：

• 服務端想在建立連線，即握手階段搞事情，實現HandshakeInterceptor。
• 服務端想在建立連線之後和處理客戶端發來的訊息，實現WebSocketHandler。
• 服務端通過WebSocketSession即可向客戶端傳送訊息，通過使用者和Session的繫結，實現對應關係。

想加深理解的同學，還是要深入到程式碼中仔細體會。限於篇幅，而且在文章中加入大量程式碼本身也不容易讀下去。所以大家還是要實際對著程式碼理解比較好。

本篇涉及到的程式碼

完整程式碼實現-小小網頁聊天室

歡迎持續關注原創，喜歡的別忘了收藏關注，碼字實在太累，你們的鼓勵就是我堅持的動力！

小銘出品，必屬精品

歡迎關注xNPE技術論壇，更多原創乾貨每日推送。