Seperate config gen into functions

server/routers/newt/handleGetConfigMessage.ts

@@ -7,7 +7,8 @@ import {
     ExitNode,
     exitNodes,
     siteResources,
-    clientSiteResourcesAssociationsCache
+    clientSiteResourcesAssociationsCache,
+    Site
 } from "@server/db";
 import { clients, clientSitesAssociationsCache, Newt, sites } from "@server/db";
 import { eq } from "drizzle-orm";
@@ -130,6 +131,38 @@ export const handleGetConfigMessage: MessageHandler = async (context) => {
         }
     }
 
+    const { peers, targets } = await buildClientConfigurationForNewtClient(
+        site,
+        exitNode
+    );
+
+    // Build the configuration response
+    const configResponse = {
+        ipAddress: site.address,
+        peers,
+        targets
+    };
+
+    logger.debug("Sending config: ", configResponse);
+
+    return {
+        message: {
+            type: "newt/wg/receive-config",
+            data: {
+                ...configResponse
+            }
+        },
+        broadcast: false,
+        excludeSender: false,
+    };
+};
+
+export async function buildClientConfigurationForNewtClient(
+    site: Site,
+    exitNode?: ExitNode
+) {
+    const siteId = site.siteId;
+
     // Get all clients connected to this site
     const clientsRes = await db
         .select()
@@ -278,22 +311,8 @@ export const handleGetConfigMessage: MessageHandler = async (context) => {
         targetsToSend.push(...resourceTargets);
     }
 
-    // Build the configuration response
-    const configResponse = {
-        ipAddress: site.address,
+    return {
         peers: validPeers,
         targets: targetsToSend
     };
-
-    logger.debug("Sending config: ", configResponse);
-    return {
-        message: {
-            type: "newt/wg/receive-config",
-            data: {
-                ...configResponse
-            }
-        },
-        broadcast: false,
-        excludeSender: false
-    };
-};
+}

server/routers/newt/handleNewtPingMessage.ts

@@ -0,0 +1,141 @@
+import { db } from "@server/db";
+import { disconnectClient } from "#dynamic/routers/ws";
+import { getClientConfigVersion, MessageHandler } from "@server/routers/ws";
+import { clients, Newt } from "@server/db";
+import { eq, lt, isNull, and, or } from "drizzle-orm";
+import logger from "@server/logger";
+import { validateSessionToken } from "@server/auth/sessions/app";
+import { checkOrgAccessPolicy } from "#dynamic/lib/checkOrgAccessPolicy";
+import { sendTerminateClient } from "../client/terminate";
+import { encodeHexLowerCase } from "@oslojs/encoding";
+import { sha256 } from "@oslojs/crypto/sha2";
+
+// Track if the offline checker interval is running
+// let offlineCheckerInterval: NodeJS.Timeout | null = null;
+// const OFFLINE_CHECK_INTERVAL = 30 * 1000; // Check every 30 seconds
+// const OFFLINE_THRESHOLD_MS = 2 * 60 * 1000; // 2 minutes
+
+/**
+ * Starts the background interval that checks for clients that haven't pinged recently
+ * and marks them as offline
+ */
+// export const startNewtOfflineChecker = (): void => {
+//     if (offlineCheckerInterval) {
+//         return; // Already running
+//     }
+
+//     offlineCheckerInterval = setInterval(async () => {
+//         try {
+//             const twoMinutesAgo = Math.floor(
+//                 (Date.now() - OFFLINE_THRESHOLD_MS) / 1000
+//             );
+
+//             // TODO: WE NEED TO MAKE SURE THIS WORKS WITH DISTRIBUTED NODES ALL DOING THE SAME THING
+
+//             // Find clients that haven't pinged in the last 2 minutes and mark them as offline
+//             const offlineClients = await db
+//                 .update(clients)
+//                 .set({ online: false })
+//                 .where(
+//                     and(
+//                         eq(clients.online, true),
+//                         or(
+//                             lt(clients.lastPing, twoMinutesAgo),
+//                             isNull(clients.lastPing)
+//                         )
+//                     )
+//                 )
+//                 .returning();
+
+//             for (const offlineClient of offlineClients) {
+//                 logger.info(
+//                     `Kicking offline newt client ${offlineClient.clientId} due to inactivity`
+//                 );
+
+//                 if (!offlineClient.newtId) {
+//                     logger.warn(
+//                         `Offline client ${offlineClient.clientId} has no newtId, cannot disconnect`
+//                     );
+//                     continue;
+//                 }
+
+//                 // Send a disconnect message to the client if connected
+//                 try {
+//                     await sendTerminateClient(
+//                         offlineClient.clientId,
+//                         offlineClient.newtId
+//                     ); // terminate first
+//                     // wait a moment to ensure the message is sent
+//                     await new Promise((resolve) => setTimeout(resolve, 1000));
+//                     await disconnectClient(offlineClient.newtId);
+//                 } catch (error) {
+//                     logger.error(
+//                         `Error sending disconnect to offline newt ${offlineClient.clientId}`,
+//                         { error }
+//                     );
+//                 }
+//             }
+//         } catch (error) {
+//             logger.error("Error in offline checker interval", { error });
+//         }
+//     }, OFFLINE_CHECK_INTERVAL);
+
+//     logger.debug("Started offline checker interval");
+// };
+
+/**
+ * Stops the background interval that checks for offline clients
+ */
+// export const stopNewtOfflineChecker = (): void => {
+//     if (offlineCheckerInterval) {
+//         clearInterval(offlineCheckerInterval);
+//         offlineCheckerInterval = null;
+//         logger.info("Stopped offline checker interval");
+//     }
+// };
+
+/**
+ * Handles ping messages from clients and responds with pong
+ */
+export const handleNewtPingMessage: MessageHandler = async (context) => {
+    const { message, client: c, sendToClient } = context;
+    const newt = c as Newt;
+
+    if (!newt) {
+        logger.warn("Newt not found");
+        return;
+    }
+
+    // get the version
+    const configVersion = await getClientConfigVersion(newt.newtId);
+
+    if (message.configVersion && configVersion != message.configVersion) {
+        logger.warn(`Newt ping with outdated config version: ${message.configVersion} (current: ${configVersion})`);
+
+        // TODO: sync the client
+    }
+
+    // try {
+    //     // Update the client's last ping timestamp
+    //     await db
+    //         .update(clients)
+    //         .set({
+    //             lastPing: Math.floor(Date.now() / 1000),
+    //             online: true
+    //         })
+    //         .where(eq(clients.clientId, newt.clientId));
+    // } catch (error) {
+    //     logger.error("Error handling ping message", { error });
+    // }
+
+    return {
+        message: {
+            type: "pong",
+            data: {
+                timestamp: new Date().toISOString()
+            }
+        },
+        broadcast: false,
+        excludeSender: false
+    };
+};

server/routers/newt/handleNewtRegisterMessage.ts

@@ -233,6 +233,35 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
             .where(eq(newts.newtId, newt.newtId));
     }
 
+    const { tcpTargets, udpTargets, validHealthCheckTargets } =
+        await buildTargetConfigurationForNewtClient(siteId);
+
+    logger.debug(
+        `Sending health check targets to newt ${newt.newtId}: ${JSON.stringify(validHealthCheckTargets)}`
+    );
+
+    return {
+        message: {
+            type: "newt/wg/connect",
+            data: {
+                endpoint: `${exitNode.endpoint}:${exitNode.listenPort}`,
+                relayPort: config.getRawConfig().gerbil.clients_start_port,
+                publicKey: exitNode.publicKey,
+                serverIP: exitNode.address.split("/")[0],
+                tunnelIP: siteSubnet.split("/")[0],
+                targets: {
+                    udp: udpTargets,
+                    tcp: tcpTargets
+                },
+                healthCheckTargets: validHealthCheckTargets
+            }
+        },
+        broadcast: false, // Send to all clients
+        excludeSender: false // Include sender in broadcast
+    };
+};
+
+export async function buildTargetConfigurationForNewtClient(siteId: number) {
     // Get all enabled targets with their resource protocol information
     const allTargets = await db
         .select({
@@ -337,30 +366,12 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
         (target) => target !== null
     );
 
-    logger.debug(
-        `Sending health check targets to newt ${newt.newtId}: ${JSON.stringify(validHealthCheckTargets)}`
-    );
-
     return {
-        message: {
-            type: "newt/wg/connect",
-            data: {
-                endpoint: `${exitNode.endpoint}:${exitNode.listenPort}`,
-                relayPort: config.getRawConfig().gerbil.clients_start_port,
-                publicKey: exitNode.publicKey,
-                serverIP: exitNode.address.split("/")[0],
-                tunnelIP: siteSubnet.split("/")[0],
-                targets: {
-                    udp: udpTargets,
-                    tcp: tcpTargets
-                },
-                healthCheckTargets: validHealthCheckTargets
-            }
-        },
-        broadcast: false, // Send to all clients
-        excludeSender: false // Include sender in broadcast
+        validHealthCheckTargets,
+        tcpTargets,
+        udpTargets
     };
-};
+}
 
 async function getUniqueSubnetForSite(
     exitNode: ExitNode,

server/routers/newt/index.ts

@@ -6,3 +6,4 @@ export * from "./handleGetConfigMessage";
 export * from "./handleSocketMessages";
 export * from "./handleNewtPingRequestMessage";
 export * from "./handleApplyBlueprintMessage";
+export * from "./handleNewtPingMessage";