From 1b8dcd0a49340b7b51f2d6a51e9fa6d990a29bb4 Mon Sep 17 00:00:00 2001
From: hippoz <paperplanelive@gmail.com>
Date: Fri, 29 Jan 2021 02:40:52 +0200
Subject: [PATCH] add connect handler that decides wether socket is accepted or
 not

---
 example/example.js | 12 ++++++++-
 index.js           | 67 ++++++++++++++++++++++++++++++++--------------
 lib/Socket.js      | 17 ++++++++++++
 lib/constants.js   |  8 +++++-
 4 files changed, 82 insertions(+), 22 deletions(-)

diff --git a/example/example.js b/example/example.js
index bf83724..8cc6810 100644
--- a/example/example.js
+++ b/example/example.js
@@ -1,3 +1,13 @@
+const http = require('http');
+
 const Wormhole = require('../index');
 
-const wormhole = new Wormhole({ urls: [ '/hello' ] });
\ No newline at end of file
+const httpServer = http.createServer();
+
+const wormhole = new Wormhole({ urls: [ '/hello' ], httpServer });
+
+wormhole.on('connect', ({ socket, accept, reject }) => {
+    accept();
+});
+
+httpServer.listen(8080);
\ No newline at end of file
diff --git a/index.js b/index.js
index c843d41..0ba7220 100644
--- a/index.js
+++ b/index.js
@@ -1,4 +1,4 @@
-const http = require('http');
+const EventEmitter = require('events');
 
 const { createLog } = require('./lib/logger');
 const handshake = require('./lib/handshake');
@@ -7,22 +7,33 @@ const Socket = require('./lib/Socket');
 
 const handshakeLog = createLog([ 'Wormhole', 'Handshake' ]);
 
-class Wormhole {
-    constructor({ urls=[ '/bruh' ], port=8080 }) {
-        this._urls = urls;
-        this._port = port;
+class Wormhole extends EventEmitter {
+    constructor({ urls=[ '/bruh' ], httpServer }) {
+        super();
 
-        this._httpServer = http.createServer((req, res) => {
+        this._urls = urls;
+        this._httpServer = httpServer;
+
+        this._sockets = [];
+
+        this._httpServer.on('request', ((req, res) => {
             if (req.method === 'GET' && req.url && this._urls.includes(req.url)) {
-                handshakeLog(`Got connection request to ${req.url} on port ${this._port}`);
+                handshakeLog(`Got connection request to ${req.url}`);
+
+                let socket = new Socket({ socket: res.socket, initalState: constants.states.CONNECTING });
 
                 const failConnection = (status=400) => {
+                    socket._setConnectionState(constants.states.CLOSING);
+                    
                     res.writeHead(status);
                     res.end();
 
+                    socket._setConnectionState(constants.states.CLOSED);
+
                     console.trace();
                 };
 
+                // TODO: check origin header
                 const websocketKey = req.headers['sec-websocket-key'];
                 const upgradeHeader = req.headers['upgrade'];
                 const websocketVersion = req.headers['sec-websocket-version'];
@@ -32,26 +43,42 @@ class Wormhole {
 
                 const websocketAccept = handshake.generateWebsocketAcceptValue(websocketKey);
 
-                handshakeLog(websocketKey, websocketAccept);
-
                 if (websocketAccept) {
-                    res.writeHead(101, {
-                        'Upgrade': 'websocket',
-                        'Connection': 'Upgrade',
-                        'Sec-WebSocket-Accept': websocketAccept
-                    });
-                    res.end();
+                    const accept = () => {
+                        try {
+                            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.');
+                        }
+                        
+                        res.writeHead(101, {
+                            'Upgrade': 'websocket',
+                            'Connection': 'Upgrade',
+                            'Sec-WebSocket-Accept': websocketAccept
+                        });
+                        res.end();
+                        
+                        socket._setConnectionState(constants.states.OPEN);
+    
+                        return true;
+                    };
 
-                    const socket = new Socket({ socket: res.socket, initalState: 'CONNECTED' });
+                    const reject = (status=403) => {
+                        try {
+                            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.');
+                        }
+                        
+                        failConnection(status);
+                    };
 
-                    return true;
+                    return this.emit('connect', { socket, accept, reject });
                 }
 
                 return failConnection();
             }
-        });
-
-        this._httpServer.listen(port);
+        }));
     }
 }
 
diff --git a/lib/Socket.js b/lib/Socket.js
index 4e248a3..8399329 100644
--- a/lib/Socket.js
+++ b/lib/Socket.js
@@ -1,12 +1,29 @@
+const constants = require("./constants");
+
 class Socket {
     constructor({ initialState='CONNECTING', socket }) {
         this._state = initialState;
         this._socket = socket;
+        this._accepted = false;
+        this._fateDecided = false; // Wether the decision to accept or reject the socket was made
 
         this._socket.on('data', (e) => {
+            if (this._state !== constants.states.OPEN) return;
+
             console.log(e.toString());
         });
     }
 }
 
+Socket.prototype._setConnectionState = function(state) {
+    this._state = state;
+};
+
+Socket.prototype._setAccepted = function(state) {
+    if (this._fateDecided) throw new Error('Tried to decide fate (wether socket is accepted or not) more than 1 time');
+
+    this._fateDecided = true;
+    this._accepted = state;
+};
+
 module.exports = Socket;
\ No newline at end of file
diff --git a/lib/constants.js b/lib/constants.js
index 1830419..c4d135b 100644
--- a/lib/constants.js
+++ b/lib/constants.js
@@ -1,5 +1,11 @@
 module.exports = {
     handshakeGUID: '258EAFA5-E914-47DA-95CA-C5AB0DC85B11',
     upgradeHeaderRequirement: 'websocket',
-    websocketVersionRequirement: '13'
+    websocketVersionRequirement: '13',
+    states: {
+        CONNECTING: 'CONNECTING',
+        OPEN: 'OPEN',
+        CLOSING: 'CLOSING',
+        CLOSED: 'CLOSED'
+    }
 }
\ No newline at end of file