Implement class Frame

metarhia · Nov 6, 2023 · 1f26327 · 1f26327
1 parent edf9759
commit 1f26327
Showing 1 changed file with 111 additions and 40 deletions.
diff --git a/lib/websocket.js b/lib/websocket.js
@@ -17,29 +17,117 @@ const OPCODE_SHORT = 0x81;
 const LEN_16_BIT = 126;
 const MAX_16_BIT = 65536;
 const LEN_64_BIT = 127;
+const MAX_64_BIT = 0x7fffffffffffffffn + 1n;
 
-const calcOffset = (frame, length) => {
-  if (length < LEN_16_BIT) return [2, 6];
-  if (length === LEN_16_BIT) return [4, 8];
-  return [10, 14];
+const DEFAULT_OPTIONS = {
+  server: true,
 };
 
-const parseFrame = (frame) => {
-  const length = frame[1] ^ 0x80;
-  const [maskOffset, dataOffset] = calcOffset(frame, length);
-  const mask = frame.subarray(maskOffset, maskOffset + MASK_LENGTH);
-  const data = frame.subarray(dataOffset);
-  return { mask, data };
-};
+class Frame {
+  #frame = null;
+  #mask = null;
+  #dataOffset = 6;
+  #length = 0n;
+  #server = false;
+
+  constructor(data, options) {
+    if (!Buffer.isBuffer(data) && typeof data !== 'string') {
+      throw new Error('Unsupported');
+    }
+    const opt = {
+      ...DEFAULT_OPTIONS,
+      ...options,
+    };
+    this.#server = opt.server;
 
-const unmask = (buffer, mask) => {
-  const data = Buffer.allocUnsafe(buffer.length);
-  buffer.copy(data);
-  for (let i = 0; i < data.length; i++) {
-    data[i] ^= mask[i & 3];
+    if (Buffer.isBuffer(data) && data.length !== 0) {
+      this.#frame = data;
+      this.#decode();
+      return;
+    }
+
+    if (typeof data === 'string' && data.length !== 0) {
+      this.#encode(data);
+      return;
+    }
   }
-  return data;
-};
+
+  #hasMask() {
+    if ((this.#frame[1] & 0x80) === 0x80) return true;
+    return false;
+  }
+
+  #decode() {
+    if (this.#server && !this.#hasMask())
+      throw new Error('1002 (protocol error)');
+    const maskLength = this.#hasMask() ? MASK_LENGTH : 0;
+    this.#length = BigInt(this.#frame[1] & 0x7f);
+    switch (this.#length) {
+      case 127n:
+        this.#dataOffset = 2 + 8 + maskLength;
+        this.#mask = this.#frame.subarray(2 + 8, 10 + MASK_LENGTH);
+        this.#length = this.#frame.readBigUInt64BE(2);
+        break;
+      case 126n:
+        this.#dataOffset = 2 + 2 + maskLength;
+        this.#mask = this.#frame.subarray(2 + 2, 4 + MASK_LENGTH);
+        this.#length = BigInt(this.#frame.readUInt16BE(2));
+        break;
+      default:
+        this.#dataOffset = 2 + maskLength;
+        this.#mask = this.#frame.subarray(2, 2 + MASK_LENGTH);
+    }
+
+    for (let i = 0n; i < this.#length; ++i) {
+      this.#frame[BigInt(this.#dataOffset) + i] ^=
+        this.#mask[i & 0x0000000000000003n];
+    }
+  }
+
+  #encode(text) {
+    const data = Buffer.from(text);
+    this.#frame = Buffer.alloc(2);
+    this.#frame[0] = 0x81; // FIN = 1, RSV = 0b000, opcode = 0b0001 (text frame)
+    const length = data.length;
+    if (length < LEN_16_BIT) {
+      this.#frame[1] = length;
+    } else if (length < MAX_16_BIT) {
+      const len = Buffer.alloc(2);
+      len.writeUint16BE(length, 0);
+      this.#frame[1] = LEN_16_BIT;
+      this.#frame = Buffer.concat([this.#frame, len]);
+    } else if (length < MAX_64_BIT) {
+      const len = Buffer.alloc(8);
+      len.writeBigUInt64BE(BigInt(length), 0);
+      this.#frame[1] = LEN_64_BIT;
+      this.#frame = Buffer.concat([this.#frame, len]);
+    } else {
+      throw new Error('text value is too long to encode in one frame!');
+    }
+    if (!this.#server) throw new Error('Unsupported');
+    this.#frame = Buffer.concat([this.#frame, data]);
+  }
+
+  toString() {
+    return this.#frame.toString(
+      'utf8',
+      this.#dataOffset,
+      Number(BigInt(this.#dataOffset) + this.#length),
+    );
+  }
+
+  get data() {
+    return this.#frame.subarray(this.#dataOffset);
+  }
+
+  get frame() {
+    return this.#frame;
+  }
+
+  get mask() {
+    return this.#mask;
+  }
+}
 
 class Connection {
   constructor(socket) {
@@ -56,34 +144,17 @@ class Connection {
   }
 
   send(text) {
-    const data = Buffer.from(text);
-    let meta = Buffer.alloc(2);
-    const length = data.length;
-    meta[0] = OPCODE_SHORT;
-    if (length < LEN_16_BIT) {
-      meta[1] = length;
-    } else if (length < MAX_16_BIT) {
-      const len = Buffer.from([(length & 0xff00) >> 8, length & 0x00ff]);
-      meta = Buffer.concat([meta, len]);
-      meta[1] = LEN_16_BIT;
-    } else {
-      const len = Buffer.alloc(8);
-      len.writeBigInt64BE(BigInt(length), 0);
-      meta = Buffer.concat([meta, len]);
-      meta[1] = LEN_64_BIT;
-    }
-    const frame = Buffer.concat([meta, data]);
-    this.socket.write(frame);
+    const frame = new Frame(text);
+    this.socket.write(frame.frame);
   }
 
   receive(data) {
     console.log('data: ', data[0], data.length);
     if (data[0] !== OPCODE_SHORT) return;
-    const frame = parseFrame(data);
-    const msg = unmask(frame.data, frame.mask);
-    const text = msg.toString();
-    this.send(`Echo "${text}"`);
+    const frame = new Frame(data);
+    const text = frame.toString();
     console.log('Message:', text);
+    this.send(`Echo "${text}"`);
   }
 
   accept(key) {