From 9a2022a4fe2341385c58cfe7b481dc9fc2ac8f7a Mon Sep 17 00:00:00 2001
From: Owen <owen@txv.io>
Date: Wed, 11 Jun 2025 09:13:38 -0400
Subject: [PATCH] Working on new exit node picking

---
 server/routers/messageHandlers.ts             |  6 ++-
 .../newt/handleNewtPingRequestMessage.ts      | 38 +++++++++++++++++
 .../routers/newt/handleNewtRegisterMessage.ts | 41 +++++++++++--------
 3 files changed, 66 insertions(+), 19 deletions(-)
 create mode 100644 server/routers/newt/handleNewtPingRequestMessage.ts

diff --git a/server/routers/messageHandlers.ts b/server/routers/messageHandlers.ts
index 2f4cd660..92d2849f 100644
--- a/server/routers/messageHandlers.ts
+++ b/server/routers/messageHandlers.ts
@@ -3,7 +3,8 @@ import {
     handleReceiveBandwidthMessage,
     handleGetConfigMessage,
     handleDockerStatusMessage,
-    handleDockerContainersMessage
+    handleDockerContainersMessage,
+    handleNewtPingRequestMessage
 } from "./newt";
 import {
     handleOlmRegisterMessage,
@@ -21,7 +22,8 @@ export const messageHandlers: Record<string, MessageHandler> = {
     "olm/wg/relay": handleOlmRelayMessage,
     "olm/ping": handleOlmPingMessage,
     "newt/socket/status": handleDockerStatusMessage,
-    "newt/socket/containers": handleDockerContainersMessage
+    "newt/socket/containers": handleDockerContainersMessage,
+    "newt/ping/request": handleNewtPingRequestMessage,
 };
 
 startOfflineChecker(); // this is to handle the offline check for olms
diff --git a/server/routers/newt/handleNewtPingRequestMessage.ts b/server/routers/newt/handleNewtPingRequestMessage.ts
new file mode 100644
index 00000000..48d67e9d
--- /dev/null
+++ b/server/routers/newt/handleNewtPingRequestMessage.ts
@@ -0,0 +1,38 @@
+import { db } from "@server/db";
+import { MessageHandler } from "../ws";
+import { exitNodes, Newt } from "@server/db";
+import logger from "@server/logger";
+
+export const handleNewtPingRequestMessage: MessageHandler = async (context) => {
+    const { message, client, sendToClient } = context;
+    const newt = client as Newt;
+
+    logger.info("Handling ping request newt message!");
+
+    if (!newt) {
+        logger.warn("Newt not found");
+        return;
+    }
+
+    // TODO: pick which nodes to send and ping better than just all of them
+    const exitNodesList = await db
+        .select()
+        .from(exitNodes);
+
+    let exitNodesPayload = exitNodesList.map((node) => ({
+        exitNodeId: node.exitNodeId,
+        endpoint: node.endpoint,
+        weight: 0 // TODO: Implement weight calculation if needed depending on load
+    }));
+
+    return {
+        message: {
+            type: "newt/ping/exitNodes",
+            data: { 
+                exitNodes: exitNodesPayload
+            }
+        },
+        broadcast: false, // Send to all clients
+        excludeSender: false // Include sender in broadcast
+    };
+};
diff --git a/server/routers/newt/handleNewtRegisterMessage.ts b/server/routers/newt/handleNewtRegisterMessage.ts
index a8396554..491dab0a 100644
--- a/server/routers/newt/handleNewtRegisterMessage.ts
+++ b/server/routers/newt/handleNewtRegisterMessage.ts
@@ -1,16 +1,10 @@
 import { db } from "@server/db";
 import { MessageHandler } from "../ws";
-import {
-    exitNodes,
-    Newt,
-    resources,
-    sites,
-    Target,
-    targets
-} from "@server/db";
+import { exitNodes, Newt, resources, sites, Target, targets } from "@server/db";
 import { eq, and, sql, inArray } from "drizzle-orm";
 import { addPeer, deletePeer } from "../gerbil/peers";
 import logger from "@server/logger";
+import { exit } from "process";
 
 export const handleNewtRegisterMessage: MessageHandler = async (context) => {
     const { message, client, sendToClient } = context;
@@ -30,7 +24,7 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
 
     const siteId = newt.siteId;
 
-    const { publicKey } = message.data;
+    const { publicKey, exitNodeId } = message.data;
     if (!publicKey) {
         logger.warn("Public key not provided");
         return;
@@ -47,18 +41,31 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
         return;
     }
 
-    await db
-        .update(sites)
-        .set({
-            pubKey: publicKey
-        })
-        .where(eq(sites.siteId, siteId))
-        .returning();
+    let exitNodeIdToQuery = site.exitNodeId;
+    if (exitNodeId && site.exitNodeId !== exitNodeId) { // This effectively moves the exit node to the new one
+        exitNodeIdToQuery = exitNodeId; // Use the provided exitNodeId if it differs from the site's exitNodeId
+        await db
+            .update(sites)
+            .set({
+                pubKey: publicKey,
+                exitNodeId: exitNodeId
+            })
+            .where(eq(sites.siteId, siteId))
+            .returning();
+    } else {
+        await db
+            .update(sites)
+            .set({
+                pubKey: publicKey
+            })
+            .where(eq(sites.siteId, siteId))
+            .returning();
+    }
 
     const [exitNode] = await db
         .select()
         .from(exitNodes)
-        .where(eq(exitNodes.exitNodeId, site.exitNodeId))
+        .where(eq(exitNodes.exitNodeId, exitNodeIdToQuery))
         .limit(1);
 
     if (site.pubKey && site.pubKey !== publicKey) {