Merge branch 'main' into copilot/fix-1112

server/routers/olm/handleOlmRegisterMessage.ts

@@ -1,4 +1,4 @@
-import { db } from "@server/db";
+import { db, ExitNode } from "@server/db";
 import { MessageHandler } from "../ws";
 import {
     clients,
@@ -28,7 +28,10 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         return;
     }
     const clientId = olm.clientId;
-    const { publicKey } = message.data;
+    const { publicKey, relay } = message.data;
+
+    logger.debug(`Olm client ID: ${clientId}, Public Key: ${publicKey}, Relay: ${relay}`);
+
     if (!publicKey) {
         logger.warn("Public key not provided");
         return;
@@ -58,9 +61,11 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         sendToClient(olm.olmId, {
             type: "olm/wg/holepunch",
             data: {
-                serverPubKey: exitNode.publicKey
+                serverPubKey: exitNode.publicKey,
+                endpoint: exitNode.endpoint,
             }
         });
+    
     }
 
     if (now - (client.lastHolePunch || 0) > 6) {
@@ -84,7 +89,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         await db
             .update(clientSites)
             .set({
-                isRelayed: false
+                isRelayed: relay == true
             })
             .where(eq(clientSites.clientId, olm.clientId));
     }
@@ -97,7 +102,15 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         .where(eq(clientSites.clientId, client.clientId));
 
     // Prepare an array to store site configurations
-    const siteConfigurations = [];
+    let siteConfigurations = [];
+    logger.debug(`Found ${sitesData.length} sites for client ${client.clientId}`);
+
+    if (sitesData.length === 0) {
+        sendToClient(olm.olmId, {
+            type: "olm/register/no-sites",
+            data: {}
+        });
+    }
 
     // Process each site
     for (const { sites: site } of sitesData) {
@@ -114,12 +127,12 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
             continue;
         }
 
-        if (site.lastHolePunch && now - site.lastHolePunch > 6) {
-            logger.warn(
-                `Site ${site.siteId} last hole punch is too old, skipping`
-            );
-            continue;
-        }
+        // if (site.lastHolePunch && now - site.lastHolePunch > 6 && relay) {
+        //     logger.warn(
+        //         `Site ${site.siteId} last hole punch is too old, skipping`
+        //     );
+        //     continue;
+        // }
 
         // If public key changed, delete old peer from this site
         if (client.pubKey && client.pubKey != publicKey) {
@@ -142,7 +155,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
             await addPeer(site.siteId, {
                 publicKey: publicKey,
                 allowedIps: [`${client.subnet.split('/')[0]}/32`], // we want to only allow from that client
-                endpoint: client.endpoint
+                endpoint: relay ? "" : client.endpoint
             });
         } else {
             logger.warn(
@@ -150,21 +163,36 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
             );
         }
 
+        let endpoint = site.endpoint;
+        if (relay) {
+            const [exitNode] = await db
+                .select()
+                .from(exitNodes)
+                .where(eq(exitNodes.exitNodeId, site.exitNodeId))
+                .limit(1);
+            if (!exitNode) {
+                logger.warn(`Exit node not found for site ${site.siteId}`);
+                continue;
+            }
+            endpoint = `${exitNode.endpoint}:21820`;
+        }
+
         // Add site configuration to the array
         siteConfigurations.push({
             siteId: site.siteId,
-            endpoint: site.endpoint,
+            endpoint: endpoint,
             publicKey: site.publicKey,
             serverIP: site.address,
-            serverPort: site.listenPort
+            serverPort: site.listenPort,
+            remoteSubnets: site.remoteSubnets
         });
     }
 
-    // If we have no valid site configurations, don't send a connect message
-    if (siteConfigurations.length === 0) {
-        logger.warn("No valid site configurations found");
-        return;
-    }
+    // REMOVED THIS SO IT CREATES THE INTERFACE AND JUST WAITS FOR THE SITES 
+    // if (siteConfigurations.length === 0) {
+    //     logger.warn("No valid site configurations found");
+    //     return;
+    // }
 
     // Return connect message with all site configurations
     return {

server/routers/olm/handleOlmRelayMessage.ts

@@ -1,7 +1,7 @@
-import { db } from "@server/db";
+import { db, exitNodes, sites } from "@server/db";
 import { MessageHandler } from "../ws";
 import { clients, clientSites, Olm } from "@server/db";
-import { eq } from "drizzle-orm";
+import { and, eq } from "drizzle-orm";
 import { updatePeer } from "../newt/peers";
 import logger from "@server/logger";
 
@@ -30,29 +30,67 @@ export const handleOlmRelayMessage: MessageHandler = async (context) => {
         .limit(1);
 
     if (!client) {
-        logger.warn("Site not found or does not have exit node");
+        logger.warn("Client not found");
         return;
     }
 
     // make sure we hand endpoints for both the site and the client and the lastHolePunch is not too old
     if (!client.pubKey) {
-        logger.warn("Site or client has no endpoint or listen port");
+        logger.warn("Client has no endpoint or listen port");
         return;
     }
 
     const { siteId } = message.data;
 
+    // Get the site
+    const [site] = await db
+        .select()
+        .from(sites)
+        .where(eq(sites.siteId, siteId))
+        .limit(1);
+
+    if (!site || !site.exitNodeId) {
+        logger.warn("Site not found or has no exit node");
+        return;
+    }
+
+    // get the site's exit node
+    const [exitNode] = await db
+        .select()
+        .from(exitNodes)
+        .where(eq(exitNodes.exitNodeId, site.exitNodeId))
+        .limit(1);
+
+    if (!exitNode) {
+        logger.warn("Exit node not found for site");
+        return;
+    }
+
     await db
         .update(clientSites)
         .set({
             isRelayed: true
         })
-        .where(eq(clientSites.clientId, olm.clientId));
+        .where(
+            and(
+                eq(clientSites.clientId, olm.clientId),
+                eq(clientSites.siteId, siteId)
+            )
+        );
 
     // update the peer on the exit node
     await updatePeer(siteId, client.pubKey, {
         endpoint: "" // this removes the endpoint
     });
 
+    sendToClient(olm.olmId, {
+        type: "olm/wg/peer/relay",
+        data: {
+            siteId: siteId,
+            endpoint: exitNode.endpoint,
+            publicKey: exitNode.publicKey
+        }
+    });
+
     return;
 };

server/routers/olm/peers.ts

@@ -12,6 +12,7 @@ export async function addPeer(
         endpoint: string;
         serverIP: string | null;
         serverPort: number | null;
+        remoteSubnets: string | null; // optional, comma-separated list of subnets that this site can access
     }
 ) {
     const [olm] = await db
@@ -30,7 +31,8 @@ export async function addPeer(
             publicKey: peer.publicKey,
             endpoint: peer.endpoint,
             serverIP: peer.serverIP,
-            serverPort: peer.serverPort
+            serverPort: peer.serverPort,
+            remoteSubnets: peer.remoteSubnets // optional, comma-separated list of subnets that this site can access
         }
     });
 
@@ -66,6 +68,7 @@ export async function updatePeer(
         endpoint: string;
         serverIP: string | null;
         serverPort: number | null;
+        remoteSubnets?: string | null; // optional, comma-separated list of subnets that
     }
 ) {
     const [olm] = await db
@@ -84,7 +87,8 @@ export async function updatePeer(
             publicKey: peer.publicKey,
             endpoint: peer.endpoint,
             serverIP: peer.serverIP,
-            serverPort: peer.serverPort
+            serverPort: peer.serverPort,
+            remoteSubnets: peer.remoteSubnets
         }
     });