From ff27041513894f5e712aa1d57aca7a18423c88ca Mon Sep 17 00:00:00 2001
From: hypermc <hyper@ssh.surf>
Date: Fri, 8 Nov 2024 02:06:12 -0500
Subject: [PATCH] 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
---
 worker-server.js | 63 +++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 62 insertions(+), 1 deletion(-)

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