Simple WebSocket on Node.js

WebSocket.server.js

const http = require('http')
const crypto = require('crypto')

// 当前请求数据状态
let STATE = {
    fin: null,
    opcode: null,
    masked: null,
    dataIndex: null,
    maskingKey: null,
    payloadData: null,
    payloadLength: null,
    remains: null
}
// 当前请求的数据分段
let DATA_LIST = []
// 所有连接 websocket 的用户
let SOCKS_LIST = []

/**
 * MAGIC_STRING
 */
const MAGIC_STRING = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

/**
 * HTTP 服务器部分
 */
const server = http
    .createServer(function (req, res) {
        res.end('\r\n');
    })
    .listen(3000, '0.0.0.0', _ => console.log('create server success on port 3000 successfully'))

/**
 * connection
 */
server.on('connection', _ => {
    console.log('http connection incoming');
});
/**
 * Upgrade 请求处理
 * http get 101 switch protocol
 */
server
    .on('upgrade', (req, socket) => {

        try {

            console.log(`upgrade webSocket on time ${new Date().toLocaleString()}`)

            // 计算返回的key
            const resKey = crypto.createHash('sha1')
                .update(req.headers['sec-websocket-key'] + MAGIC_STRING)
                .digest('base64');

            // 构造响应头
            const resHeaders = ([
                'HTTP/1.1 101 Switching Protocols',
                'Upgrade: websocket',
                'Connection: Upgrade',
                'Sec-WebSocket-Origin:' + req.headers.origin,
                'Sec-WebSocket-Accept: ' + resKey
            ]).concat('', '').join('\r\n');

            // upgrade 成功转为 websocket 协议后才可以接收数据
            socket
                .on('data', message => {

                    // 1、解码
                    const incoming = readMessage(message)
                    console.log(`socket get message ${JSON.stringify(incoming)}`)

                    const {opcode, payloadData, remains, payloadLength} = STATE

                    // 2 检查 opcode
                    if (opcode === 8) {
                        // 2-1 如果状态码为8说明要关闭连接
                        socket.end()
                        return
                    } else if (opcode === 10) {
                        // 2-2 如果是心跳pong,回一个ping
                        return
                    } else {

                        DATA_LIST.push(payloadData)

                        // 2-3 检查数据分段是否完成
                        if (remains === 0) {

                            const data = parseData(Buffer.concat(DATA_LIST, payloadLength))
                            console.log(`socket get message parsed ${data}`)

                            // 如果 socket 没有关闭，通知 所有的 client 已经接收完数据
                            // 可以给每一个 socket 添加 namespace 来区分开特定的 socket
                            if (opcode !== 8) {
                                SOCKS_LIST.forEach(socket => socket.write(writeMessage('200 with data ' + JSON.stringify(data))))
                            }

                            // reset local data list set
                            DATA_LIST = []
                        }
                    }
                })
                .on('close', () => {
                    console.log('socket status close');
                })
                .on('end', () => {
                    console.log('socket status end');
                });

            SOCKS_LIST.push(socket)
            socket.write(resHeaders);

        } catch (e) {
            console.error(e)
            socket.end()
        }
    })
    .on('close', _ => {
        SOCKS_LIST = []
        DATA_LIST = []
        STATE = {}
        try {
            global.gc();
        } catch (e) {
            console.log("You must run program with 'node --expose-gc yourfile.js'");
            process.exit();
        }
    })


function parseData(allData) {
    let len = allData.length,
        i = 0;
    for (; i < len; i++) {
        allData[i] = allData[i] ^ STATE.maskingKey[i % 4];// 异或运算，使用maskingKey四个字节轮流进行计算
    }
    // 判断数据类型，如果为文本类型
    if (STATE.opcode === 1) allData = allData.toString();

    return allData;
}

function readMessage(message) {

    let data01 = message[0].toString(2), data02 = message[1].toString(2);// 第一个字节和第二个字节的二进制字符串
    let fin = data01.slice(0, 1); // 得到fin字符串
    let opcode = parseInt(data01.slice(4), 2); // 将opcode的4位二进制字符串转化为十进制
    let dataIndex = 2; // 初始数据下标，因为第一、二个字节肯定不是payload data
    let masked = data02.slice(0, 1); // masked 的值为1或者0
    let payloadLength = parseInt(data02.slice(1), 2); // payloadLength的值按照0-125， 126, 127分别取值
    let payloadData;
    let maskingKey; // 如果masked = 1则，表示是客户端发送过来的，需要用使用masking key掩码解析payload data


    if (payloadLength === 126) {
        dataIndex += 2;// 数据长度为后面2个字节，3、4
        payloadLength = message.readUInt16BE(2);// 使用 Node.js Buffer 方法。从3个字节开始左往右读16位（也就是两个字节3、4）.也可以使用data.readUIntBE(2, 2)
    } else if (payloadLength === 127) {
        dataIndex += 8;// 数据长度为后面8个字节，3、4、6、7、8、9、10、11
        payloadLength = message.readUInt32BE(2) + message.readUInt32BE(6); //先读取的3、4、6、7，然后读取的8、9、10、11。也可以使用data.readUIntBE(2,6)+data.readUIntBE(8,2)
    }

    // 判断是否有masking key
    if (masked === "1") {
        maskingKey = message.slice(dataIndex, dataIndex + 4);
        dataIndex += 4; // 重新定位数据位置
        payloadData = message.slice(dataIndex);
    }
    let remains = STATE.remains || payloadLength; // 剩余数据长度
    remains = remains - payloadData.length; // 还剩余多少长度的payload data

    return Object.assign(STATE, {
        fin,
        opcode,
        masked,
        dataIndex,
        maskingKey,
        payloadData,
        payloadLength,
        remains
    })
}

function writeMessage(data) {

    let dataType = Buffer.isBuffer(data);// 数据类型
    let dataBuf, // 需要发送的二进制数据
        dataLength,// 数据真实长度
        dataIndex = 2; // 数据的起始长度
    let frame; // 数据帧

    if (dataType) dataBuf = data;
    else dataBuf = Buffer.from(data); // 也可以不做类型判断，直接Buffer.form(data)
    dataLength = dataBuf.byteLength;

    // 计算payload data在frame中的起始位置
    dataIndex = dataIndex + (dataLength > 65535 ? 8 : (dataLength > 125 ? 2 : 0));

    frame = new Buffer.alloc(dataIndex + dataLength);

    // 第一个字节,fin = 1,opcode = 1
    frame[0] = parseInt(10000001, 2);

    // 长度超过65535的则由8个字节表示，因为4个字节能表达的长度为4294967295，已经完全够用，因此直接将前面4个字节置0
    if (dataLength > 65535) {
        frame[1] = 127; // 第二个字节
        frame.writeUInt32BE(0, 2);
        frame.writeUInt32BE(dataLength, 6);
    } else if (dataLength > 125) {
        frame[1] = 126;
        frame.writeUInt16BE(dataLength, 2);
    } else {
        frame[1] = dataLength;
    }

    // 服务端发送到客户端的数据
    frame.write(dataBuf.toString(), dataIndex);

    return frame;
}