Nodejs原生实现WebSocket服务

import { createHash } from 'node:crypto';

/* 对Sec-Websocket-Key进行加密，作为Sec-Websocket-Accept的值 */
export function encryptSocketKey(key) {
  const WEBSOCKET_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11';
  const sha1Hash = createHash('sha1');

  sha1Hash.update(`${ key }${ WEBSOCKET_GUID }`);

  return sha1Hash.digest('base64');
}

/* 处理接收到的数据 */
export function decodeWsFrame(data) {
  let start = 0;
  const frame = {
    isFinal: (data[start] & 0x80) === 0x80,
    opcode: data[start++] & 0xF,
    masked: (data[start] & 0x80) === 0x80,
    payloadLen: data[start++] & 0x7F,
    maskingKey: '',
    payloadData: null
  };

  if (frame.payloadLen === 126) {
    frame.payloadLen = (data[start++] << 8) + data[start++];
  } else if (frame.payloadLen === 127) {
    frame.payloadLen = 0;

    for (let i = 7; i >= 0; i--) {
      frame.payloadLen += (data[start++] << (i * 8));
    }
  }

  if (frame.payloadLen) {
    if (frame.masked) {
      const maskingKey = [data[start++], data[start++], data[start++], data[start++]];

      frame.maskingKey = maskingKey;
      frame.payloadData = data
        .subarray(start, start + frame.payloadLen)
        .map((byte, index) => byte ^ maskingKey[index % 4]);
    } else {
      frame.payloadData = data.slice(start, start + frame.payloadLen);
    }
  }

  return frame;
}

/* 处理发送的数据 */
export function encodeWsFrame(data) {
  const isFinal = data.isFinal !== undefined ? data.isFinal : true,
    opcode = data.opcode !== undefined ? data.opcode : 1,
    payloadData = data.payloadData ? Buffer.from(data.payloadData) : null,
    payloadLen = payloadData ? payloadData.length : 0;
  const frame = [];

  if (isFinal) {
    frame.push((1 << 7) + opcode);
  } else {
    frame.push(opcode);
  }

  if (payloadLen < 126) {
    frame.push(payloadLen);
  } else if (payloadLen < 65536) {
    frame.push(126, payloadLen >> 8, payloadLen & 0xFF);
  } else {
    frame.push(127);

    for (let i = 7; i >= 0; i--) {
      frame.push((payloadLen & (0xFF << (i * 8))) >> (i * 8));
    }
  }

  return payloadData
    ? Buffer.concat([Buffer.from(frame), payloadData])
    : Buffer.from(frame);
}

export function rStr(len: number): string {
  const str: string = 'QWERTYUIOPASDFGHJKLZXCVBNM1234567890';
  let result: string = '';

  for (let i: number = 0; i < len; i++) {
    const rIndex: number = Math.floor(Math.random() * str.length);

    result += str[rIndex];
  }

  return result;
}

import net from 'node:net';
import { encodeWsFrame, decodeWsFrame, encryptSocketKey } from './utils.js';

/* 建立请求 */
function createConnect(socket) {
  return new Promise((resolve, reject) => {
    socket.once('data', function(buffer) {
      const data = buffer.toString();
      const headers = data.split('\r\n')
        .splice(1)
        .filter((o) => o !== '')
        .reduce((result, header, index) => {
          const [key, value] = header.split(': ');

          result[key.toLowerCase()] = value;

          return result;
        }, {});

      // 判断是否为websocket并检查版本
      if (!(
        headers.upgrade === 'websocket'
        && headers['sec-websocket-version'] === '13'
      )) {
        socket.end();
        resolve(false);

        return;
      }

      const responseHeader = `HTTP/1.1 101 Switching Protocols\r
Upgrade: websocket\r
Connection: Upgrade\r
Sec-Websocket-Accept: ${ encryptSocketKey(headers['sec-websocket-key']) }\r\n\r\n`

      socket.write(responseHeader);
      resolve(true);
    });
  });
}

const server = net.createServer(async function(socket) {
  const createResult = await createConnect(socket);

  if (!createResult) return;

  socket.on('data', function(buffer) {
    const data = decodeWsFrame(buffer);
    const { opcode, payloadData } = data;

    if (opcode === 8) {
      socket.end();
    } else {
      socket.write(encodeWsFrame({
        opcode: 1,
        payloadData: `接收：${ payloadData.toString() }`
      }));
    }
  });
});

server.listen(5059);

import net from 'node:net';
import { setInterval } from 'node:timers';
import { encryptSocketKey, encodeWsFrame, decodeWsFrame, rStr } from './utils.js';

const client = new net.Socket();

/* 建立请求 */
function createConnect(client, uuid) {
  return new Promise((resolve, reject) => {
    client.once('data', function(buffer) {
      const data = buffer.toString();
      const headers = data.split('\r\n')
        .splice(1)
        .filter((o) => o !== '')
        .reduce((result, header, index) => {
          const [key, value] = header.split(': ');

          result[key.toLowerCase()] = value;

          return result;
        }, {});

      // 判断是否为websocket并检查版本
      if (!(
        headers.upgrade === 'websocket'
        && headers['sec-websocket-accept'] === encryptSocketKey(uuid)
      )) {
        client.end();
        resolve(false);

        return;
      }

      resolve(true);
    });
  });
}

client.connect(5059, '127.0.0.1', async function() {
  const uuid = btoa(rStr(16));

  client.write(`GET /index.html HTTP/1.1\r
Connection: Upgrade\r
Upgrade: websocket\r
Sec-WebSocket-Key: ${ uuid }\r
Sec-WebSocket-Version: 13\r\n\r\n`);

  const createResult = await createConnect(client, uuid);

  if (!createResult) return;

  let i = 0;

  client.on('data', function(buffer) {
    const data = decodeWsFrame(buffer);

    console.log(data.payloadData.toString());
  });

  setInterval(() => {
    console.log(`发送：${ i }`);
    client.write(encodeWsFrame({
      opcode: 1,
      payloadData: `${ i++ }`
    }));
  }, 3_000);
});