Working on getting targets connected to newt

server/routers/newt/handleRegisterMessage.ts

@@ -1,42 +1,116 @@
 import db from "@server/db";
 import { MessageHandler } from "../ws";
-import { sites } from "@server/db/schema";
-import { eq } from "drizzle-orm";
+import { exitNodes, resources, sites, targets } from "@server/db/schema";
+import { eq, inArray } from "drizzle-orm";
+import { addPeer, deletePeer } from "../gerbil/peers";
+import logger from "@server/logger";
+import { findNextAvailableCidr } from "@server/utils/ip";
+import { exit } from "process";
 
 export const handleRegisterMessage: MessageHandler = async (context) => {
     const { message, newt, sendToClient } = context;
-    
+
     if (!newt) {
-        console.log("Newt not found");
+        logger.warn("Newt not found");
         return;
     }
 
     if (!newt.siteId) {
-        console.log("Newt has no site!"); // TODO: Maybe we create the site here?
+        logger.warn("Newt has no site!"); // TODO: Maybe we create the site here?
         return;
     }
-    
+
     const siteId = newt.siteId;
-        
-    // get the site
-    const site = await db
+
+    const { publicKey } = message.data;
+    if (!publicKey) {
+        logger.warn("Public key not provided");
+        return;
+    }
+
+    // const [site] = await db
+    //     .select()
+    //     .from(sites)
+    //     .where(eq(sites.siteId, siteId))
+    //     .limit(1);
+
+    const [site] = await db
+    .update(sites)
+        .set({
+            pubKey: publicKey
+        })
+        .where(eq(sites.siteId, siteId))
+        .returning();
+
+
+    if (!site || !site.exitNodeId) {
+        logger.warn("Site not found or does not have exit node");
+        return;
+    }
+
+    const [exitNode] = await db
+        .select()
+        .from(exitNodes)
+        .where(eq(exitNodes.exitNodeId, site.exitNodeId))
+        .limit(1);
+
+    if (site.pubKey && site.pubKey !== publicKey) {
+        logger.info("Public key mismatch. Deleting old peer...");
+        await deletePeer(site.exitNodeId, site.pubKey);
+    }
+
+    if (!site.subnet) {
+        logger.warn("Site has no subnet");
+        return;
+    }
+
+    // add the peer to the exit node
+    await addPeer(site.exitNodeId, {
+        publicKey: publicKey,
+        allowedIps: [site.subnet],
+    });
+
+    const siteResources = await db.select().from(resources).where(eq(resources.siteId, siteId));
+
+    // get the targets from the resourceIds
+    const siteTargets = await db
     .select()
-    .from(sites)
-    .where(eq(sites.siteId, siteId))
-    .limit(1);
+    .from(targets)
+    .where(
+        inArray(
+            targets.resourceId,
+            siteResources.map(resource => resource.resourceId)
+        )
+    );
 
+    const udpTargets = siteTargets
+        .filter((target) => target.protocol === "udp")
+        .map((target) => {
+            return `${target.internalPort ? target.internalPort + ":" : ""}${target.ip}:${target.port}`;
+        });
+
+    const tcpTargets = siteTargets
+        .filter((target) => target.protocol === "tcp")
+        .map((target) => {
+            return `${target.internalPort ? target.internalPort + ":" : ""}${target.ip}:${target.port}`;
+        });
 
-        const { publicKey } = message.data;
 
     return {
         message: {
-            type: 'newt/wg/connect',
+            type: "newt/wg/connect",
             data: {
-                publicKey: 'publicKey',
-
-            }
+                endpoint: exitNode.endpoint,
+                publicKey: exitNode.publicKey,
+                serverIP: exitNode.address,
+                tunnelIP: site.subnet,
+                targets: {
+                    udp: udpTargets,
+                    tcp: tcpTargets,
+                }
+            },
         },
-        broadcast: false,  // Send to all clients
-        excludeSender: false  // Include sender in broadcast
+        broadcast: false, // Send to all clients
+        excludeSender: false, // Include sender in broadcast
     };
-};
+};