Worker Server Update - Adding Websocket Support

Add WebSocket support for long-running connections

- Implement WebSocket upgrade handling to support persistent connections
- Establish unique tunnels for each WebSocket connection
- Clean up tunnels upon WebSocket disconnect
- Maintain existing HTTP request functionality
This commit is contained in:
hypermc 2024-11-08 02:06:12 -05:00
parent f31c52a4f2
commit ff27041513

View File

@ -95,6 +95,67 @@ if (cluster.isMaster) {
}
});
// Add WebSocket upgrade handling
server.on('upgrade', async function (req, socket, head) {
try {
const split = req.headers.host.split('.');
const publicKey = Buffer.from(b32.decode.asBytes(split[0].toUpperCase()));
if (publicKey.length < 32) {
console.log("Invalid Connection!");
socket.destroy();
return;
}
// Create a unique tunnel for each WebSocket connection
const port = await getAvailablePort();
const tunnelServer = net.createServer(function (servsock) {
const socketConn = dhtServer.connect(publicKey);
let open = { local: true, remote: true };
servsock.on('data', (d) => socketConn.write(d));
socketConn.on('data', (d) => servsock.write(d));
const remoteend = () => {
if (open.remote) socketConn.end();
open.remote = false;
};
const localend = () => {
if (open.local) servsock.end();
open.local = false;
};
servsock.on('error', remoteend);
servsock.on('finish', remoteend);
servsock.on('end', remoteend);
socketConn.on('finish', localend);
socketConn.on('error', localend);
socketConn.on('end', localend);
});
tunnelServer.listen(port, "127.0.0.1", () => {
if (DEBUG === 1 || CONINFO === 1) console.log(`WebSocket tunnel server listening on port ${port}`);
proxy.ws(req, socket, head, {
target: 'http://127.0.0.1:' + port
}, function (e) {
console.log("Proxy WS Error: ", e);
socket.destroy();
});
// Clean up tunnel after WebSocket disconnects
socket.on('close', () => {
tunnelServer.close(() => {
if (DEBUG === 1 || CONINFO === 1) console.log("WebSocket tunnel closed after disconnect.");
});
});
});
} catch (e) {
console.error("Error Occurred: ", e);
socket.destroy();
}
});
server.listen(8081, () => {
console.log(`Worker ${process.pid} listening on port 8081`);
});
@ -122,4 +183,4 @@ if (cluster.isMaster) {
tryPort();
});
}
}
}