Java NIO實現WebSocket服務器

簡介

在HTTP請求中，服務器往往處於被動的一方，通常都是客戶端向服務器發送請求時，服務器纔會做出響應，服務器並不會主動向客戶端推送消息。因此WebSocket API就爲此誕生。WebSocket API是HTML5中的一大特色，能夠使得建立連接的雙方在任意時刻相互推送消息，這意味着不同於HTTP，服務器服務器也可以主動向客戶端推送消息了。

關於WebSocket的介紹，可以參考下一篇博文http://blog.csdn.net/zwto1/article/details/52493119#websocket%E5%8E%9F%E7%90%86

WebSocket協議的格式

爲了實現一個能與H5的WebSocket API通信的服務器，我們需要先熟悉WebSocket數據包的格式。定的格式。

握手數據包

在一個連接建立以後，建立連接的雙方纔可以互相推送消息。雙方通過握手即可建立一個連接。握手數據包的格式如下：

客戶端向服務器發起請求

可以見到，客戶端請求連接建立的數據包是一個字符串，而且第一行表明這實際上是一個HTTP報文。其中Connection: Upgrade以及Upgrade: websocket兩字段就是用來告知服務器這是一個WebSocket握手請求。

服務器還要關心的一個字段是Sec-WebSocket-Key（倒數第二行），其值是一個隨機base64字符串，服務器怎麼處理該字符串請往下看。

服務器迴應請求

可以看到HTTP狀態碼爲101，同樣，服務端也帶有Connection和Upgrade字段來表明這是一個WebSocket數據包。

Sec-WebSocket-Accept字段是對請求報文中Sec-WebSocket-Key字段進行摘要運算的結果。其運算過程如下
1、將Sec-WebSocket-Key字段的值與字符串258EAFA5-E914-47DA-95CA-C5AB0DC85B11拼接。
2、對拼接後的字符串進行sha1運算，得到160位摘要（二進制）。
3、以base64的形式表示得到的摘要。
客戶端會進行同樣的運算，並且與服務器返回來的字段作對比，如果發現二者不相同，連接就無法建立了。

通信數據幀

通信數據幀的格式如下（參考官方文檔https://www.rfc-editor.org/rfc/rfc6455.txt）

其中各個字段的含義如下
FIN: 1bit，表示這是否爲分片的最後一個數據幀。這是考慮到發送的數據有可能被分片的情況，如果存在分片，將此字段置1就表明這是最後一個分片。如果不存在分片，此字段恆爲1。因爲只有一個分片就一定是最後一個分片。

RSV1, RSV2, RSV3: 各1bit，全0。現在暫時用不上，爲了將來可能用於功能拓展保留的字段。

Opcode: 4bits
指出數據的類型，值的解釋如下

值	含義
0x0	附加數據幀
0x1	文本數據幀
0x2	二進制數據幀
0x3-0x7	暫無定義
0x8	關閉連接
0x9	表示ping
0xA	表示pong
0aB-0xF	暫無定義

MASK: 1bit
表明是否對數據進行掩碼運算，置1表示使用掩碼。從客戶端向服務器發送的數據必須使用掩碼。

Payload length: 7 bits, 7+16 bits, or 7+64 bits
表明數據的長度。
如果長度在0-125內，這7bits就表示數據的長度；
如果值爲126，緊接着後面2字節（16bits）才表示數據的長度；
如果值爲127，後面8字節（64bits）表示數據的長度。

Masking-key: 無 或 4 字節
如果掩碼字段（MASK）置0，就不需要Masking-key。如果掩碼字段爲1，這4字節就是Masking-key，用它與數據部分進行異或運算。

Payload Data: 數據部分，長度可變。

關於其他詳細說明可以參考官方文檔，例如消息分片規則等。

實現一個WebSocket服務器（羣聊天室例子）

爲了更加深刻的理解這樣一個協議，這裏沒有使用Java已經封裝好操作的類庫。

基於NIO監聽端口

基於NIO中的ServerSocketChannel，實現一個接收並讀取Socket內容的服務端套路如下。

public class WebSocketServer {

    private Selector serverSelector;
    private WebSocketListener socketListener;
    private boolean isRunning = true;

    public WebSocketServer(int serverPort, WebSocketListener socketListener) throws IOException {
        //初始化ServerSocketChannel
        ServerSocketChannel serverSocketChannel =
                ServerSocketChannel.open();
        serverSocketChannel.bind(new InetSocketAddress(serverPort));
        serverSocketChannel.configureBlocking(false);

        //創建選擇器
        serverSelector = Selector.open();

        //註冊ServerSocketChannel的ACCEPT事件至選擇器
        serverSocketChannel.register(serverSelector, SelectionKey.OP_ACCEPT);
        this.socketListener = socketListener;
    }

    public void run() throws IOException {
        while (isRunning) {
            int selectCount = serverSelector.select();
            if (selectCount == 0)
                continue;

            Iterator<SelectionKey> iterator = serverSelector.selectedKeys().iterator();
            while (iterator.hasNext()) {
                SelectionKey selectKey = iterator.next();

                if (selectKey.isAcceptable()) {

                    //ACCEPT就緒，此時調用ServerSocketChannel的accept()方法可獲得連接的SocketChannel對象，將其READ事件註冊到選擇器，就可以讀取內容了。
                    ServerSocketChannel serverChannel = (ServerSocketChannel) selectKey.channel();
                    SocketChannel acceptSocketChannel = serverChannel.accept();
                  acceptSocketChannel.configureBlocking(false);  //記得設置爲非阻塞模式 
                    acceptSocketChannel.register(serverSelector, SelectionKey.OP_READ);

                } else if (selectKey.isReadable()) {
                    //TODO 讀取並處理數據
                }
                iterator.remove();
            }
        }
    }
}

會話管理

在本聊天室中，一個WebSocket連接就視爲一個會話，也就是一個用戶登錄。定義ClientSession來管理每一個連接的SocketChannel。

public class ClientSession {
    private SocketChannel socketChannel;
    private String sessionID;

    public ClientSession(SocketChannel channel) {
        this.socketChannel = channel;
        try {
            MessageDigest sha1 = MessageDigest.getInstance("sha1");
            sha1.update(Util.longToByteArray(System.currentTimeMillis()));
            BigInteger bi = new BigInteger(sha1.digest());
            sessionID = bi.toString(16);
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
    }

    public SocketChannel getSocketChannel() {
        return socketChannel;
    }

    public String getSessionID() {
        return sessionID;
    }
}

會話的幾種狀態

對於聊天室的一個用戶從建立連接到釋放連接，服務端最關心的無非是四個關鍵點：會話建立，收到來自客戶端的消息，會話關閉，拋出異常。可以將其抽象成接口。

interface WebSocketListener {
    void onOpen(ClientSession session) throws IOException;

    void onMessage(ClientSession session) throws IOException;

    void onException(ClientSession session, Exception ex);

    void onClose(ClientSession session) throws IOException;
}

處理SocketChannel

前面一開始的關於ServerSocketChannel代碼中，還沒對SocketChannel進行處理，現在來補上。

public class WebSocketServer {

    private Selector serverSelector;
    private WebSocketListener socketListener;
    private boolean isRunning = true;

    //...部分代碼省略

    public void run() throws IOException {
        while (isRunning) {
            int selectCount = serverSelector.select();
            if (selectCount == 0)
                continue;

            Iterator<SelectionKey> iterator = serverSelector.selectedKeys().iterator();
            while (iterator.hasNext()) {
                SelectionKey selectKey = iterator.next();

                if (selectKey.isAcceptable()) {
                    //重複代碼省略
                } else if (selectKey.isReadable()) {
                    try {
                        SocketChannel socketChannel = (SocketChannel) selectKey.channel();
                        ClientSession session = (ClientSession) selectKey.attachment();　　//用前面定義的ClientSession來作爲SocketChannel的attach object，方便存儲關於SocketChannel的其他信息，容易管理。

                        if (session == null) {
                            //如果SocketChannel還沒有被ClientSession綁定，認爲這是一個新連接，需要完成握手
                            byte[] byteArray = Util.readByteArray(socketChannel);
                            System.out.println(new String(byteArray));
                            WSProtocol.Header header = WSProtocol.Header.decodeFromString(new String(byteArray));
                            String receiveKey = header.getHeader("Sec-WebSocket-Key");
                            String response = WSProtocol.getHandShakeResponse(receiveKey);
                            socketChannel.write(ByteBuffer.wrap(response.getBytes()));

                            ClientSession newSession = new ClientSession(socketChannel);
                            selectKey.attach(newSession);
                            socketListener.onOpen(newSession);  //會話打開
                        } else {
                            //收到數據，交給上面定義的接口處理
                            socketListener.onMessage(session);
                        }
                    } catch (IOException e) {
                        e.printStackTrace();
                        //出現異常，進行一系列處理
                        selectKey.channel().close();
                        selectKey.cancel();

                        ClientSession attSession = (ClientSession) selectKey.attachment();
                        socketListener.onException(attSession, e);  //拋出異常
                        socketListener.onClose(attSession);  //強制關閉拋出異常的連接
                    }
                }
                iterator.remove();
            }
        }
    }
}
//...省略下部分代碼

數據包的處理

數據包的處理放在單獨一個類裏面

class WSProtocol {

    static class Header {
        private Map<String, String> headers = new HashMap<>();

        String getHeader(String key) {
            return headers.get(key);
        }

        static Header decodeFromString(String headers) {
            Header header = new Header();

            Map<String, String> headerMap = new HashMap<>();
            String[] headerArray = headers.split("\r\n");
            for (String headerLine : headerArray) {
                if (headerLine.contains(":")) {
                    int splitPos = headerLine.indexOf(":");
                    String key = headerLine.substring(0, splitPos);
                    String value = headerLine.substring(splitPos + 1).trim();
                    headerMap.put(key, value);
                }
            }
            header.headers = headerMap;
            return header;
        }
    }

    static String getHandShakeResponse(String receiveKey) {
        String keyOrigin = receiveKey + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
        MessageDigest sha1;
        String accept = null;
        try {
            sha1 = MessageDigest.getInstance("sha1");
            sha1.update(keyOrigin.getBytes());
            accept = new String(Base64.getEncoder().encode(sha1.digest()));
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        String echoHeader = "";
        echoHeader += "HTTP/1.1 101 Switching Protocols\r\n";
        echoHeader += "Upgrade: websocket\r\n";
        echoHeader += "Connection: Upgrade\r\n";
        echoHeader += "Sec-WebSocket-Accept: " + accept + "\r\n";
        echoHeader += "\r\n";

        return echoHeader;
    }
}

另外，還需要實現接口來處理會話的幾種狀態。

class WebSocketListenerImpl implements WebSocketListener {

    private Map<String, ClientSession> connSessionMap = new HashMap<>();

    @Override
    public void onOpen(ClientSession session) throws IOException {
        connSessionMap.put(session.getSessionID(), session);
        sendBoardCast(session.getSocketChannel().socket().getInetAddress().getHostName() + ":" +
                session.getSocketChannel().socket().getPort() + " Join", session);
        Log.info("session open: " + session.getSessionID());
    }

    @Override
    public void onMessage(ClientSession session) throws IOException {

        SocketChannel socketChannel = session.getSocketChannel();

        byte[] bytesData = Util.readByteArray(socketChannel);

        //opcode爲8，對方主動斷開連接
        if ((bytesData[0] & 0xf) == 8) {
            throw new IOException("session disconnect.");
        }

        byte payloadLength = (byte) (bytesData[1] & 0x7f);
        byte[] mask = Arrays.copyOfRange(bytesData, 2, 6);
        byte[] payloadData = Arrays.copyOfRange(bytesData, 6, bytesData.length);
        for (int i = 0; i < payloadData.length; i++) {
            payloadData[i] = (byte) (payloadData[i] ^ mask[i % 4]);
        }

        String echoData =
                "[" + session.getSocketChannel().socket().getInetAddress().getHostAddress() + ":" +
                        session.getSocketChannel().socket().getPort() + "]" +
                        (new String(payloadData));

        sendBoardCast(echoData, session);
    }

    @Override
    public void onException(ClientSession session, Exception ex) {
        Log.info("exception catch: " + ex.getMessage());
    }

    @Override
    public void onClose(ClientSession session) throws IOException {
        connSessionMap.remove(session.getSessionID());

        sendBoardCast(session.getSocketChannel().socket().getInetAddress().getHostName() + ":" +
                session.getSocketChannel().socket().getPort() + " Leave", session);

        Log.info("closed sessionId = " + session.getSessionID());
    }

    private void sendBoardCast(String message, ClientSession ownSession) throws IOException {
        Iterator<ClientSession> iterator = connSessionMap.values().iterator();
        while (iterator.hasNext()) {
            ClientSession nextSession = iterator.next();
            if (nextSession == ownSession) {
                continue;
            }
            byte[] boardCastData = new byte[2 + message.getBytes().length];
            boardCastData[0] = (byte) 0x81;
            boardCastData[1] = (byte) message.getBytes().length;
            System.arraycopy(message.getBytes(), 0, boardCastData, 2, message.getBytes().length);

            nextSession.getSocketChannel().write(ByteBuffer.wrap(boardCastData));
        }
    }
}

這裏考慮的極其簡單，都是數據長度不超過126，且都是文本不分片的情況，有興趣的可以按照WebSocket的文檔將數據操作的過程寫完整。

最後寫一張簡單的測試頁面

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>Chat</title>
</head>

<script>
    var socket = null;
    var isConn = false;

    socket = new WebSocket('ws://127.0.0.1:8888');
    socket.onerror = function (err) {
        console.log(err);
        addError('連接錯誤')
    };
    socket.onopen = function () {
        isConn = true;
        addMessage('連接成功');
        console.log('open');
    };
    socket.onmessage = function (event) {
        console.log(event.data);
        addMessage(event.data);
    };
    socket.onclose = function () {
        console.log('close')
    };


    function sendMessage() {
        var sendText = document.getElementById('input_text').value;
        if (!isConn) {
            addError('發送失敗');
        } else {
            if (!sendText) {
                addError('不要發送空消息');
            } else {
                addMessage('<label style="font-style: oblique">' + '[我]' + sendText + '</label>');
                socket.send(sendText);
            }
        }
    }

    function addMessage(message) {
        var textShow = document.getElementById('show-message');
        textShow.innerHTML += message + '<br>'
    }

    function addError(error) {
        var textShow = document.getElementById('show-message');
        textShow.innerHTML += '<label style="color: red">[Error] ' + error + '</label><br>'
    }
</script>
<style>
    button {
        border-radius: 5px;
        padding: 8px;
        color: white;
        border: none;
    }

    input {
        border-radius: 4px;
        width: 300px;
        padding: 6px;
        border: 1px solid dodgerblue;
    }

    #show-message {
        border-radius: 4px;
        height: 320px;
        width: 480px;
        border: 1px double darkgray;
        padding: 8px;
    }
</style>
<body>
<div id="show-message">
</div>
<div style="margin-top: 20px">
    <input id="input_text" type="text" placeholder="message">
    <button onclick="sendMessage()" style="background-color: dodgerblue">Send</button>
</div>
</body>
</html>

運行效果