node原生实现websocket通信

code

const net = require('net');
const {
  dealwithResData,
  dealwithReqData,
  createPongFrame,
  dealwithSocketHeader,
} = require('./utils');
// 创建TCP服务器
const server = net.createServer((socket) => {

  const onClose = () => {
    console.log('close');
  };
  
  const onMsg = (frame) => {
    const req = dealwithReqData(frame);
    console.log(req, '??????');
    const resText = dealwithResData('收到');
    socket.write(resText);
  };
  
  let skip = false;
  const hartbeat = () => {
    setInterval(() => {
      if (skip) {
        skip = false;
      }
      const pongFrame = createPongFrame();
      console.log('send');
      socket.write(pongFrame);
      socket.write(dealwithResData('auto send'));
    }, 3000);
  };
  socket.once('data', (buffer) => {
    // 接收到数据后进行握手处理
    const headers = buffer
      .toString()
      .split('\r\n')
      .reduce((acc, current) => {
        const [key, value] = current.split(': ');
        acc[key] = value;
        return acc;
      }, {});

    // 检查WebSocket的key是否存在
    if (headers['Sec-WebSocket-Key']) {
      // 发送握手的响应头
      socket.write(dealwithSocketHeader(headers));
      
      // hartbeat();
      // 握手完成后，可以继续监听数据事件来接收数据帧
      socket.on('data', (frame) => {
        // 检查FIN位是否为1，表示这是消息的最后一个帧
        const isFinalFrame = frame[0] & 0x80;
        // 获取操作码，判断是否是文本帧
        const opCode = frame[0] & 0x0f;
        console.log(frame[0], frame[1]);
        switch (opCode) {
          case 0x0:
            // 继续帧
            break;
          case 0x1:
            // 文本帧
            onMsg(frame);
            break;
          case 0x2:
            // 二进制帧
            break;
          case 0x8:
            // 连接关闭
            onClose();
            break;
          case 0x9:
            // ping
            // skip = true;
            // const pongFrame = createPongFrame();
            // socket.write(pongFrame);
            break;
          case 0xa:
            // pong
            break;
        }
      });
    }
  });

  // 监听socket的关闭事件
  socket.on('end', () => {
    console.log('Client disconnected');
  });
});

// 监听端口
server.listen(3000, () => {
  console.log('Server listening on port 3000');
});

utils

// utils.js
const crypto = require('crypto');
function dealwithResData(text) {
  // 文本数据的第一个字节：10000001，表示这是最后一个帧，并且是文本帧
  const firstByte = 0x81;
  const buffer = Buffer.from(text, 'utf8');
  let payloadLength = buffer.length;
  let secondByte = payloadLength;

  // 根据payload长度决定帧的大小
  let payloadLengthBytes = [];
  if (payloadLength > 65535) {
    secondByte = 127;
    payloadLengthBytes.push((payloadLength >> 56) & 255);
    payloadLengthBytes.push((payloadLength >> 48) & 255);
    payloadLengthBytes.push((payloadLength >> 40) & 255);
    payloadLengthBytes.push((payloadLength >> 32) & 255);
    payloadLengthBytes.push((payloadLength >> 24) & 255);
    payloadLengthBytes.push((payloadLength >> 16) & 255);
    payloadLengthBytes.push((payloadLength >> 8) & 255);
    payloadLengthBytes.push(payloadLength & 255);
  } else if (payloadLength > 125) {
    secondByte = 126;
    payloadLengthBytes.push((payloadLength >> 8) & 255);
    payloadLengthBytes.push(payloadLength & 255);
  }

  // 创建帧
  const frame = Buffer.concat([
    Buffer.from([firstByte, secondByte]),
    Buffer.from(payloadLengthBytes),
    buffer,
  ]);
  return frame;
}
const dealwithReqData = (frame) => {
  // 获取掩码位
  const isMasked = frame[1] & 0x80;

  let currentOffset = 2;
  let mask;
  if (isMasked) {
    mask = frame.slice(currentOffset, currentOffset + 4);
    currentOffset += 4;
  }

  // 提取数据
  const data = frame.slice(currentOffset);

  // 如果数据被掩码了，需要解码
  if (isMasked) {
    for (let i = 0; i < data.length; i++) {
      data[i] ^= mask[i % 4];
    }
  }

  // 将数据转换为UTF-8字符串
  const text = data.toString('utf8');
  return text;
};

function createPongFrame() {
  // 创建一个简单的Pong帧
  // Pong帧的opCode是0xA
  const pongFrame = Buffer.alloc(2);
  pongFrame[0] = 0b10001010; // FIN位设置为1，opCode设置为0xA
  pongFrame[1] = 0; // 假设没有负载数据，所以长度是0
  return pongFrame;
}

const dealwithSocketHeader = (headers) => {
  // 计算WebSocket的响应key
  const SEC_WEBSOCKET_ACCEPT = crypto
    .createHash('sha1')
    .update(headers['Sec-WebSocket-Key'] + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11')
    .digest('base64');

  // 发送握手的响应头
  return (
    'HTTP/1.1 101 Web Socket Protocol Handshake\r\n' +
    'Upgrade: websocket\r\n' +
    'Connection: Upgrade\r\n' +
    'Sec-WebSocket-Accept: ' +
    SEC_WEBSOCKET_ACCEPT +
    '\r\n' +
    '\r\n'
  );
};
module.exports = {
  dealwithResData,
  dealwithReqData,
  createPongFrame,
  dealwithSocketHeader,
};