clean up code and split more things into classes

index.js

@@ -1,85 +1,3 @@
-const EventEmitter = require('events');
+const Wormhole = require('./lib/Wormhole');
-
-const { createLog } = require('./lib/logger');
-const handshake = require('./lib/handshake');
-const constants = require('./lib/constants');
-const Socket = require('./lib/Socket');
-
-const handshakeLog = createLog([ 'Wormhole', 'Handshake' ]);
-
-class Wormhole extends EventEmitter {
-    constructor({ urls=[ '/bruh' ], httpServer }) {
-        super();
-
-        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}`);
-
-                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'];
-
-                if (upgradeHeader !== constants.upgradeHeaderRequirement) return failConnection();
-                if (websocketVersion !== constants.websocketVersionRequirement) return failConnection();
-
-                const websocketAccept = handshake.generateWebsocketAcceptValue(websocketKey);
-
-                if (websocketAccept) {
-                    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 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 this.emit('connect', { socket, accept, reject });
-                }
-
-                return failConnection();
-            }
-        }));
-    }
-}
 
 module.exports = Wormhole;

lib/ConnectingSocket.js

@@ -0,0 +1,78 @@
+const constants = require('./constants');
+const handshake = require('./handshake');
+
+class ConnectingSocket {
+    constructor({ socket, res, websocketKey, upgradeHeader, websocketVersion }) {
+        this.websocketKey = websocketKey;
+        this.upgradeHeader = upgradeHeader;
+        this.websocketVersion = websocketVersion;
+
+        this.socket = socket;
+        this.res = res;
+
+        this.connectionSuccess = undefined;
+    }
+}
+
+ConnectingSocket.prototype._handshakeEndWithStatus = function(status=400) {
+    this.res.writeHead(status);
+    this.res.end();
+};
+ConnectingSocket.prototype._handshakeAccept = function(websocketAccept) {
+    this.res.writeHead(101, {
+        'Upgrade': 'websocket',
+        'Connection': 'Upgrade',
+        'Sec-WebSocket-Accept': websocketAccept
+    });
+    this.res.end();
+};
+ConnectingSocket.prototype._failConnection = function(status=400) {
+    this.socket._setConnectionState(constants.states.CLOSING);
+    this._handshakeEndWithStatus(status);
+    this.socket._setConnectionState(constants.states.CLOSED);
+
+    this.connectionSuccess = false;
+    return true;
+};
+
+
+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.');
+    }
+    
+    const websocketAccept = handshake.generateWebsocketAcceptValue(this.websocketKey);
+    if (!websocketAccept) return this._failConnection(400);
+
+    this._handshakeAccept(websocketAccept);
+    this.socket._setConnectionState(constants.states.OPEN);
+
+    this.connectionSuccess = true;
+    return true;
+};
+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.');
+    }
+    
+    this._failConnection(status);
+
+    return true;
+};
+
+
+ConnectingSocket.prototype.connectionFunctions = function() {
+    const headersValid = handshake.validateHeaders({ upgradeHeader: this.upgradeHeader, websocketVersion: this.websocketVersion });
+    if (!headersValid) return;
+
+    return {
+        accept: this.acceptConnection.bind(this),
+        reject: this.rejectConnection.bind(this)
+    }
+};
+
+module.exports = ConnectingSocket;

lib/Wormhole.js

@@ -0,0 +1,41 @@
+const EventEmitter = require('events');
+
+const { createLog } = require('./logger');
+const constants = require('./constants');
+const Socket = require('./Socket');
+const ConnectingSocket = require('./ConnectingSocket');
+
+const handshakeLog = createLog([ 'Wormhole', 'Handshake' ]);
+
+class Wormhole extends EventEmitter {
+    constructor({ urls=[ '/bruh' ], httpServer }) {
+        super();
+
+        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}`);
+
+                const websocketKey = req.headers['sec-websocket-key'];
+                const upgradeHeader = req.headers['upgrade'];
+                const websocketVersion = req.headers['sec-websocket-version'];
+
+                let socket = new Socket({ socket: res.socket, initalState: constants.states.CONNECTING });
+                let connectingSocket = new ConnectingSocket({ res, socket, upgradeHeader, websocketKey, websocketVersion });
+                
+                const connectionFunctions = connectingSocket.connectionFunctions();
+                if (!connectionFunctions) return connectingSocket.rejectConnection(400);
+                
+                const { accept, reject } = connectionFunctions;
+
+                this.emit('connect', { socket, accept, reject });
+            }
+        });
+    }
+}
+
+module.exports = Wormhole;

lib/handshake.js

@@ -1,17 +1,23 @@
 const crypto = require('crypto');
 
-const { handshakeGUID } = require('./constants');
+const constants = require('./constants');
 
 const generateWebsocketAcceptValue = (websocketKey) => {
-    if (typeof websocketKey !== 'string' || typeof handshakeGUID !== 'string') {
+    if (typeof websocketKey !== 'string' || typeof constants.handshakeGUID !== 'string') {
         // TODO: maybe throw error?
         return;
     }
 
-    const concatenated = websocketKey + handshakeGUID;
+    const concatenated = websocketKey + constants.handshakeGUID;
     const sha1HashInBase64 = crypto.createHash('sha1').update(concatenated, 'binary').digest('base64');
     
     return sha1HashInBase64;
 };
 
-module.exports = { generateWebsocketAcceptValue };
+const validateHeaders = ({ upgradeHeader, websocketVersion }) => {
+    if (upgradeHeader !== constants.upgradeHeaderRequirement) return false;
+    if (websocketVersion !== constants.websocketVersionRequirement) return false;
+    return true;
+};
+
+module.exports = { generateWebsocketAcceptValue, validateHeaders };