From 3c8434fb8ae17139ac94e65636708d15136391cf Mon Sep 17 00:00:00 2001 From: hippoz Date: Wed, 3 Feb 2021 18:27:08 +0200 Subject: [PATCH] clean up code and split more things into classes --- index.js | 84 +---------------------------------------- lib/ConnectingSocket.js | 78 ++++++++++++++++++++++++++++++++++++++ lib/Wormhole.js | 41 ++++++++++++++++++++ lib/handshake.js | 14 +++++-- 4 files changed, 130 insertions(+), 87 deletions(-) create mode 100644 lib/ConnectingSocket.js create mode 100644 lib/Wormhole.js diff --git a/index.js b/index.js index 0ba7220..7b9f491 100644 --- a/index.js +++ b/index.js @@ -1,85 +1,3 @@ -const EventEmitter = require('events'); - -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(); - } - })); - } -} +const Wormhole = require('./lib/Wormhole'); module.exports = Wormhole; \ No newline at end of file diff --git a/lib/ConnectingSocket.js b/lib/ConnectingSocket.js new file mode 100644 index 0000000..be81b93 --- /dev/null +++ b/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; \ No newline at end of file diff --git a/lib/Wormhole.js b/lib/Wormhole.js new file mode 100644 index 0000000..fe03658 --- /dev/null +++ b/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; \ No newline at end of file diff --git a/lib/handshake.js b/lib/handshake.js index c6630a6..d4cc09d 100644 --- a/lib/handshake.js +++ b/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 }; \ No newline at end of file +const validateHeaders = ({ upgradeHeader, websocketVersion }) => { + if (upgradeHeader !== constants.upgradeHeaderRequirement) return false; + if (websocketVersion !== constants.websocketVersionRequirement) return false; + return true; +}; + +module.exports = { generateWebsocketAcceptValue, validateHeaders }; \ No newline at end of file