Implement Connection for websocket

PR-URL: #472
metarhia · Oct 30, 2023 · edf9759 · edf9759
1 parent fa2f2aa
commit edf9759
Showing 1 changed file with 54 additions and 46 deletions.
diff --git a/lib/websocket.js b/lib/websocket.js
@@ -7,6 +7,7 @@ const UPGRADE = [
   'HTTP/1.1 101 Switching Protocols',
   'Upgrade: websocket',
   'Connection: Upgrade',
+  'Sec-WebSocket-Accept: ',
 ].join(EOL);
 const MAGIC = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11';
 const MASK_LENGTH = 4;
@@ -17,12 +18,6 @@ const LEN_16_BIT = 126;
 const MAX_16_BIT = 65536;
 const LEN_64_BIT = 127;
 
-const acceptKey = (key) => {
-  const hash = crypto.createHash('sha1');
-  hash.update(key + MAGIC);
-  return hash.digest('base64');
-};
-
 const calcOffset = (frame, length) => {
   if (length < LEN_16_BIT) return [2, 6];
   if (length === LEN_16_BIT) return [4, 8];
@@ -37,27 +32,6 @@ const parseFrame = (frame) => {
   return { mask, data };
 };
 
-const sendMessage = (socket, 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]);
-  socket.write(frame);
-};
-
 const unmask = (buffer, mask) => {
   const data = Buffer.allocUnsafe(buffer.length);
   buffer.copy(data);
@@ -67,31 +41,65 @@ const unmask = (buffer, mask) => {
   return data;
 };
 
-const init = (server) => {
-  server.on('upgrade', (req, socket, head) => {
-    const receive = (data) => {
-      if (data[0] !== OPCODE_SHORT) return;
-      const frame = parseFrame(data);
-      const msg = unmask(frame.data, frame.mask);
-      const text = msg.toString();
-      sendMessage(socket, `Echo "${text}"`);
-    };
-
-    const key = req.headers['sec-websocket-key'];
-    const accept = acceptKey(key);
-    const packet = UPGRADE + EOL + `Sec-WebSocket-Accept: ${accept}`;
-    socket.write(packet + EOL + EOL);
-    receive(head);
-
-    socket.on('data', receive);
-
+class Connection {
+  constructor(socket) {
+    this.socket = socket;
+    socket.on('data', (data) => {
+      this.receive(data);
+    });
     socket.on('error', (error) => {
       console.log(error.code);
     });
-
     setInterval(() => {
       socket.write(PING);
     }, PING_TIMEOUT);
+  }
+
+  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);
+  }
+
+  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}"`);
+    console.log('Message:', text);
+  }
+
+  accept(key) {
+    const hash = crypto.createHash('sha1');
+    hash.update(key + MAGIC);
+    const packet = UPGRADE + hash.digest('base64');
+    this.socket.write(packet + EOL + EOL);
+  }
+}
+
+const init = (server) => {
+  server.on('upgrade', (req, socket, head) => {
+    const ws = new Connection(socket);
+    const key = req.headers['sec-websocket-key'];
+    ws.accept(key);
+    ws.receive(head);
   });
 };