Java NIO 实现 WebSocket 协议

文章参考：https://blog.youkuaiyun.com/Kurozaki_Kun/article/details/78843783

WebSocket协议

WebSocket是一种在单个TCP连接上进行全双工通信的协议。 WebSocket使得客户端和服务器之间的数据交换变得更加简单，允许服务端主动向客户端推送数据。在WebSocket API中，浏览器和服务器只需要完成一次握手，两者之间就直接可以创建持久性的连接，并进行双向数据传输。

WebSocket协议相比于Http协议来说，最大的特点就是可以实现服务端主动向客户端发送消息。在WebSocket出现之前，如果客户端想实时获取服务端的消息，就需要使用AJAX轮询，查询是否有消息，这样就很消耗服务器资源和带宽。但是用WebSocket就可以实现服务端主动向客户端发送数据，并且只需要占用一个TCP连接，节省了资源和带宽。

WebSocket连接建立过程

为了建立一个WebSocket连接，客户端浏览器首先要向服务器发起一个HTTP请求，这个请求和通常的HTTP请求不同，包含了一些附加的头信息，其中附加头信息“Upgrade: WebSocket” 表明这是一个申请协议升级的HTTP请求。服务器端解析这些附加的信息头，然后生成应答消息返回给客户端，客户端和服务端的WebSocket连接就建立了。之后就可以使用WebSocket协议的格式来双向发送消息。

建立连接时发送的HTTP请求头：

返回的HTTP响应头：

在响应头中的 Sec-WebSocket-Accept 时通过Sec-WebSocket-Key构造出来的。首先在Sec-WebSocket-Key后接上一个258EAFA5-E914-47DA-95CA-C5AB0DC85B11，然后再进行SHA1摘要得到160位数据在，在使用BASE64进行编码，最后得到的就是Sec-WebSocket-Accept。

WebSocket数据发送过程

WebSocket数据发送的帧格式如下所示：

FIN - 1bit

在数据发送的过程中，可能会分片发送，FIN表示是否为最后一个分片。如果发生了分片，则1表示时最后一个分片；不能再分片的情况下，这个标志总是为1。

RSV1 RSV2 RSV3 - 1bit each

用于扩展，不使用扩展时需要为全0；非零时通信双方必须协商好扩展。这里我们用不上。

OPCODE - 4bits

用于表示所传送数据的类型，也就是payload中的数据。

数值	含义
0x0	附加数据帧
0x1	文本数据帧
0x2	二进制数据帧
0x3-0x7	保留
0x8	关闭连接帧
0x9	ping帧
0xA	pong帧
0xB-0xF	保留

MASK - 1bit

用于表示payload是否被进行了掩码运算，1表示使用掩码，0表示不使用掩码。从客户端发送向服务端的数据帧必须使用掩码。

Payload length 7 bits，7+16 bits or 7+64 bits

用于表示payload的长度，有以下三种情况：

Payload length 表示的大小	payload的长度
0 - 125	Payload length 大小
126	之后的2个字节表示的无符号整数
127	之后的8个字节表示的无符号整数

Masking-key - 0 or 4 bytes

32 bit长的掩码，如果MASK为1，则帧中就存在这一个字段，在解析payload时，需要进行使用32长掩码进行异或操作，之后才能得到正确结果。

Java NIO 实现

利用Java NIO 来实现一个聊天室。部分代码如下。

NIO的常规代码：

selector.select(1000);
Set<SelectionKey> selectionKeys = selector.selectedKeys();
Iterator<SelectionKey> it = selectionKeys.iterator();
while (it.hasNext()) {
    SelectionKey key = it.next();
    it.remove();
    if (key.isAcceptable()) {
        handleAccept(key);
    }

    if (key.isReadable()) {
        handleRead(key);
    }
}

接受连接：

public void handleAccept(SelectionKey key) {
    ServerSocketChannel ssc = (ServerSocketChannel) key.channel();
    SocketChannel sc;
    try {
        sc = ssc.accept();
        sc.configureBlocking(false);
        sc.register(selector, SelectionKey.OP_READ);
        System.out.println(String.format("[server] -- client %s connected.", sc.getRemoteAddress().toString()));
    } catch (IOException e) {
        System.out.println(String.format("[server] -- error occur when accept: %s.", e.getMessage()));
        key.cancel();
    }
}

读取通道中的数据：

public void handleRead(SelectionKey key) {
    SocketChannel sc = (SocketChannel) key.channel();
    Client client = (Client) key.attachment();
    ByteBuffer byteBuffer = ByteBuffer.allocate(1024);
    // 如果是第一次连接进来，就需要创建一个客户端对象，存储起来
    if (client == null) {
        client = new Client(sc);
        clients.add(client);
        key.attach(client);
        byteBuffer.clear();
        // 如果连接还没有建立，就是要HTTP建立连接
        try {
            sc.read(byteBuffer);
            byteBuffer.flip();
            String response = WebSocketHandler.getResponse(new String(byteBuffer.array()));
            byteBuffer.clear();
            byteBuffer.put(response.getBytes());
            byteBuffer.flip();
            while (byteBuffer.hasRemaining()) {
                sc.write(byteBuffer);
            }
        } catch (IOException e) {
            System.out.println(String.format("[server] -- error occur when read: %s.", e.getMessage()));
        }
        String message = "[系统消息] " + client.toString() + " 加入了群聊";
        broadcast(message.getBytes(), client);
    }
    byteBuffer.clear();
    int read = 0;
    try {
        read = sc.read(byteBuffer);
        if (read > 0) {
            byteBuffer.flip();
            int opcode = byteBuffer.get() & 0x0f;
            // 8表示客户端关闭了连接
            if (opcode == 8) {
                System.out.println(String.format("[server] -- client %s connection close.", sc.getRemoteAddress()));
                clients.remove(client);
                String message = "[系统消息] " + client.toString() + " 退出了群聊";
                broadcast(message.getBytes(), client);
                sc.close();
                key.cancel();
                return;
            }
			// 只考虑了最简单的payload长度情况。
            int len = byteBuffer.get();
            len &= 0x7f;
            byte[] mask = new byte[4];
            byteBuffer.get(mask);
            byte[] payload = new byte[len];
            byteBuffer.get(payload);

            for (int i = 0; i < payload.length; i++) {
                payload[i] ^= mask[i % 4];
            }

            System.out.println(String
                    .format("[server] -- client: [%s], send: [%s].", client.toString(), new String(payload)));
            String message = String.format("[%s]: %s", client.toString(), new String(payload));
            broadcast(message.getBytes(), client);

        } else if (read == -1) {
            System.out.println(String.format("[server] -- client %s connection close.", sc.getRemoteAddress()));
            clients.remove(client);
            String message = "[系统消息] " + client.toString() + " 退出了群聊";
            broadcast(message.getBytes(), client);
            sc.close();
            key.cancel();
        }
    } catch (IOException e) {
        System.out.println(String.format("[server] -- error occur when read: %s.", e.getMessage()));
    }

}

使用HTTP建立WebSocket连接。

public class WebSocketHandler {

    private static String APPEND_STRING = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

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

        public String get(String key) {
            return properties.get(key);
        }
    }

    private WebSocketHandler() {}

    private static Header phrase(String request) {
        Header header = new Header();
        String[] pros = request.split("\r\n");
        for (String pro : pros) {
            if (pro.contains(":")) {
                int index = pro.indexOf(":");
                String key = pro.substring(0, index).trim();
                String value = pro.substring(index + 1).trim();
                header.properties.put(key, value);
            }
        }
        return header;
    }

    public static String getResponse(String request) {
        Header header = phrase(request);
        String acceptKey = header.get("Sec-WebSocket-Key") + APPEND_STRING;
        MessageDigest sha1;
        try {
            sha1 = MessageDigest.getInstance("sha1");
            sha1.update(acceptKey.getBytes());
            acceptKey = new String(Base64.getEncoder().encode(sha1.digest()));
        } catch (NoSuchAlgorithmException e) {
            System.out.println("fail to encode " + e.getMessage());
            return null;
        }
        StringBuilder stringBuilder = new StringBuilder();
        stringBuilder.append("HTTP/1.1 101 Switching Protocols\r\n").append("Upgrade: websocket\r\n")
                     .append("Connection: Upgrade\r\n").append("Sec-WebSocket-Accept: " + acceptKey + "\r\n")
                     .append("\r\n");
        return stringBuilder.toString();
    }

}

客户端对象

/**
 * @author XinHui Chen
 * @date 2020/2/8 19:20
 */
public class Client {
    private SocketChannel socketChannel = null;

    private String id = null;

    public SocketChannel getSocketChannel() {
        return socketChannel;
    }

    public String getId() {
        return id;
    }

    Client(SocketChannel socketChannel) {
        this.socketChannel = socketChannel;
        this.id = UUID.randomUUID().toString();
    }

    @Override
    public String toString() {
        try {
            return id + " " + socketChannel.getRemoteAddress().toString();
        } catch (IOException e) {
            System.out.println(e.getMessage());
            return null;
        }
    }
}

结果

使用网页和控制台与服务端建立WebSocket连接，发送数据。两个都能成功显示。