Working on new exit node picking

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

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

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) {