From c065746cc85b73cddf7c1d03bb7e2d692dfcf36a Mon Sep 17 00:00:00 2001
From: hippoz <paperplanelive@gmail.com>
Date: Thu, 25 Mar 2021 20:24:43 +0200
Subject: [PATCH] amazing

---
 example/example.html    |  3 ++-
 example/example.js      | 13 ++++++-------
 lib/ConnectingSocket.js | 10 ++++++----
 lib/Parser.js           | 18 +++++++++--------
 lib/Socket.js           | 43 ++++++++++++++---------------------------
 lib/Wormhole.js         |  3 ++-
 6 files changed, 41 insertions(+), 49 deletions(-)

diff --git a/example/example.html b/example/example.html
index ee98a84..d58d730 100644
--- a/example/example.html
+++ b/example/example.html
@@ -7,7 +7,7 @@
 
     <script>
         let theFunniestStringEver = '';
-        for (let i = 0; i < 100000; i++) {
+        for (let i = 0; i < 20000; i++) {
             theFunniestStringEver += 'h';
         }
 
@@ -17,6 +17,7 @@
         // Connection opened
         socket.addEventListener('open', function (event) {
             socket.send(theFunniestStringEver);
+            socket.send(new Uint8Array([ 41, 41, 41, 41, 41 ]));
         });
 
         socket.addEventListener('error', console.error)
diff --git a/example/example.js b/example/example.js
index aa37c37..5fe2873 100644
--- a/example/example.js
+++ b/example/example.js
@@ -6,15 +6,14 @@ const httpServer = http.createServer();
 
 const wormhole = new Wormhole({ urls: [ '/hello' ], httpServer });
 
-wormhole.on('connect', ({ socket, accept, reject }) => {
-    accept();
-    const buf = Buffer.from([0x41, 0x41, 0x41, 0x41]);
-    socket.send(buf);
+wormhole.on('connect', ({ connectingSocket, accept, reject }) => {
+    const socket = accept();
+    socket.send(Buffer.from(new Array(1000000)));
     socket.on('text', (e) => {
-        //console.log('Got text frame', e);
+        console.log('Got text frame', e);
     });
-    socket.on('binary', () => {
-        //console.log('Got binary frame', e);
+    socket.on('binary', (e) => {
+        console.log('Got binary frame', e);
     });
 });
 
diff --git a/lib/ConnectingSocket.js b/lib/ConnectingSocket.js
index be81b93..213a0d7 100644
--- a/lib/ConnectingSocket.js
+++ b/lib/ConnectingSocket.js
@@ -2,13 +2,14 @@ const constants = require('./constants');
 const handshake = require('./handshake');
 
 class ConnectingSocket {
-    constructor({ socket, res, websocketKey, upgradeHeader, websocketVersion }) {
+    constructor({ socket, res, req, websocketKey, upgradeHeader, websocketVersion }) {
         this.websocketKey = websocketKey;
         this.upgradeHeader = upgradeHeader;
         this.websocketVersion = websocketVersion;
 
         this.socket = socket;
         this.res = res;
+        this.req = req;
 
         this.connectionSuccess = undefined;
     }
@@ -40,7 +41,7 @@ ConnectingSocket.prototype.acceptConnection = function() {
     try {
         this.socket._setAccepted(true);
     } catch(e) {
-        throw new Error('Tried to set socket fate (wether it is accept or not) more than once. Check if there are multiple listeners for the "connect" event or if you are somehow calling accept or reject multiple times.');
+        throw new Error('Tried to set socket fate (wether it is accepted or not) more than once. Check if there are multiple listeners for the "connect" event or if you are somehow calling accept or reject multiple times.');
     }
     
     const websocketAccept = handshake.generateWebsocketAcceptValue(this.websocketKey);
@@ -50,13 +51,14 @@ ConnectingSocket.prototype.acceptConnection = function() {
     this.socket._setConnectionState(constants.states.OPEN);
 
     this.connectionSuccess = true;
-    return true;
+
+    return this.socket;
 };
 ConnectingSocket.prototype.rejectConnection = function(status=403) {
     try {
         this.socket._setAccepted(false);
     } catch(e) {
-        throw new Error('Tried to set socket fate (wether it is accept or not) more than once. Check if there are multiple listeners for the "connect" event or if you are somehow calling accept or reject multiple times.');
+        throw new Error('Tried to set socket fate (wether it is accepted or not) more than once. Check if there are multiple listeners for the "connect" event or if you are somehow calling accept or reject multiple times.');
     }
     
     this._failConnection(status);
diff --git a/lib/Parser.js b/lib/Parser.js
index 05afa69..2931212 100644
--- a/lib/Parser.js
+++ b/lib/Parser.js
@@ -1,3 +1,8 @@
+const constants = require("./constants");
+const { createLog } = require('./logger');
+
+const parserLog = createLog([ 'Wormhole', 'Parser' ]);
+
 class WebsocketFrame {
     constructor() {
         // Just defining the structure of the WebsocketFrame,
@@ -36,10 +41,10 @@ WebsocketFrame.prototype.toBuffer = function() {
     let lenBuffer;
     if (this.PayloadLen === 126) {
         lenBuffer = Buffer.alloc(2);
-        buf.writeUint16BE(this.PayloadLenEx);
+        lenBuffer.writeUint16BE(this.PayloadLenEx);
     } else if (this.PayloadLen === 127) {
         lenBuffer = Buffer.alloc(8);
-        buf.writeBigUint64BE(this.PayloadLenEx);
+        lenBuffer.writeBigUint64BE(BigInt(this.PayloadLenEx));
     }
     lenBuffer && packet.push(lenBuffer);
 
@@ -55,8 +60,6 @@ const parseWebsocketFrame = (data) => {
 
     const frame = new WebsocketFrame();
 
-    console.log(firstByte.toString(2));
-
     // 1 byte - first byte
     frame.FIN = (firstByte >> 7); // >> 7 - Get most significant bit (FIN)
     frame.RSVx = (firstByte & 0x70) >> 4; // 0x70[0b01110000] - Get the 3 bits after the most significant bit (RSVx)
@@ -88,7 +91,7 @@ const parseWebsocketFrame = (data) => {
     // TODO: implement unmasked string decoding
     if (frame.MASK) {
         switch (frame.Opcode) {
-            case 1: { // Denotes a text frame   
+            case constants.opcodes.TEXT_FRAME: { // Denotes a text frame   
                 const payloadDataView = new DataView(frame.PayloadData);
                 const maskingKeyView = new DataView(frame.MaskingKey);
     
@@ -103,8 +106,7 @@ const parseWebsocketFrame = (data) => {
                 break;
             }
     
-            case 2: { // Denotes a binary frame
-                // TODO: untested
+            case constants.opcodes.BINARY_FRAME: { // Denotes a binary frame
                 const payloadDataView = new DataView(frame.PayloadData);
                 const maskingKeyView = new DataView(frame.MaskingKey);
     
@@ -114,7 +116,7 @@ const parseWebsocketFrame = (data) => {
                     decoded.push(payloadDataView.getUint8(i) ^ maskingKeyView.getUint8(i % 4));
                 }
                 
-                frame.DecodedPayloadData = decoded;
+                frame.DecodedPayloadData = Uint8Array.from(decoded);
 
                 break;
             }
diff --git a/lib/Socket.js b/lib/Socket.js
index 56ca63c..afc9197 100644
--- a/lib/Socket.js
+++ b/lib/Socket.js
@@ -2,6 +2,9 @@ const EventEmitter = require('events');
 
 const constants = require("./constants");
 const parser = require('./Parser');
+const { createLog } = require('./logger');
+
+const socketLog = createLog([ 'Wormhole', 'Socket' ]);
 
 class Socket extends EventEmitter {
     constructor({ initialState='CONNECTING', socket }) {
@@ -12,39 +15,24 @@ class Socket extends EventEmitter {
         this._accepted = false;
         this._fateDecided = false; // Wether the decision to accept or reject the socket was made
 
-        this._buffering = undefined;
-
         this._socket.on('data', (e) => {
-            if (!this._isConnected()) return;
+            if (!this._isConnected()) return socketLog('ERROR: Received data, but connection is not fully completed');
 
-            const packet = this._decodePayload(e.buffer);
-
-            console.log(packet);
-
-            if (this._buffering && packet.FIN === 1) {
-                try {
-                    this.emit('binary', Buffer.concat([this._buffering, packet.PayloadData]));
-                    this._buffering = undefined;
-                } catch (e) {
-                    console.error(e);
-                }
+            let packet;
+            try {
+                packet = this._decodePayload(e.buffer);
+            } catch(e) {
+                return socketLog('ERROR: Received data, but encountered error while parsing: ', e);
             }
 
-            if (packet.FIN === 0) {
-                if (!this._buffering) {
-                    this._buffering = Buffer.from(packet.PayloadData);
-                } else {
-                    Buffer.concat([this._buffering, Buffer.from(packet.PayloadData)]);
-                }
-                
-                console.log('Buffering', packet);
-                return;
-            }
+            if (!packet.FIN) return socketLog('WARN: Message fragmentation is not supported, but got fragmented message');
 
             if (packet.Opcode === constants.opcodes.TEXT_FRAME) {
                 this.emit('text', packet);
             } else if (packet.Opcode === constants.opcodes.BINARY_FRAME) {
                 this.emit('binary', packet);
+            } else {
+                return socketLog(`WARN: Received unsupported opcode 0x${packet.Opcode.toString(16)}`);
             }
         });
     }
@@ -56,7 +44,7 @@ Socket.prototype.send = function(payload) {
     } else if (payload instanceof Buffer) {
         this._sendBuffer(payload);
     } else {
-        throw new Error(`Unsupported type ${typeof payload}, supported types are: string, Buffer`);
+        throw new Error(`Unsupported payload type ${typeof payload}, supported types are: string, Buffer`);
     }
 };
 
@@ -70,8 +58,7 @@ Socket.prototype._sendBuffer = function(buffer) {
     const length = buffer.byteLength;
     if (length > 125) {
         frame.PayloadLen = 126;
-    } else if (length >= 65536) {
-        frame.PayloadLen = 127;
+        if (length >= 65535) frame.PayloadLen = 127;
     } else if (length <= 125) {
         frame.PayloadLen = length;
     }
@@ -97,7 +84,7 @@ Socket.prototype._setAccepted = function(state) {
 };
 
 Socket.prototype._isConnected = function() {
-    return (this._fateDecided && this._isConnected && this._accepted === true);
+    return (this._fateDecided && this._isConnected && this._accepted);
 }
 
 module.exports = Socket;
\ No newline at end of file
diff --git a/lib/Wormhole.js b/lib/Wormhole.js
index fe03658..9110911 100644
--- a/lib/Wormhole.js
+++ b/lib/Wormhole.js
@@ -18,6 +18,7 @@ class Wormhole extends EventEmitter {
 
         this._httpServer.on('request', (req, res) => {
             if (req.method === 'GET' && req.url && this._urls.includes(req.url)) {
+                // TODO: check for origin header
                 handshakeLog(`Got connection request to ${req.url}`);
 
                 const websocketKey = req.headers['sec-websocket-key'];
@@ -32,7 +33,7 @@ class Wormhole extends EventEmitter {
                 
                 const { accept, reject } = connectionFunctions;
 
-                this.emit('connect', { socket, accept, reject });
+                this.emit('connect', { connectingSocket, accept, reject });
             }
         });
     }