diff --git a/worker-server.js b/worker-server.js index aebe41c..959cde1 100644 --- a/worker-server.js +++ b/worker-server.js @@ -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(); }); } -} +} \ No newline at end of file