clean up code and split more things into classes

This commit is contained in:
hippoz 2021-02-03 18:27:08 +02:00
parent 6f6a89acbe
commit 3c8434fb8a
Signed by: hippoz
GPG key ID: 7C52899193467641
4 changed files with 130 additions and 87 deletions

View file

@ -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;

78
lib/ConnectingSocket.js Normal file
View file

@ -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;

41
lib/Wormhole.js Normal file
View file

@ -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;

View file

@ -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 };