Start changes for multi site clients

- Org subnet and assign sites and clients out of the same subnet group
  on each org
- Add join table for client on multiple sites
- Start to handle websocket endpoints for these multiple connections

server/routers/olm/handleOlmRegisterMessage.ts

@@ -1,91 +1,53 @@
 import db from "@server/db";
 import { MessageHandler } from "../ws";
-import { clients, exitNodes, Olm, olms, sites } from "@server/db/schema";
-import { eq } from "drizzle-orm";
+import { clients, clientSites, exitNodes, Olm, olms, sites } from "@server/db/schema";
+import { eq, inArray } from "drizzle-orm";
 import { addPeer, deletePeer } from "../newt/peers";
 import logger from "@server/logger";
 
 export const handleOlmRegisterMessage: MessageHandler = async (context) => {
     const { message, client: c, sendToClient } = context;
     const olm = c as Olm;
-
     logger.info("Handling register olm message!");
-
     if (!olm) {
         logger.warn("Olm not found");
         return;
     }
-
     if (!olm.clientId) {
-        logger.warn("Olm has no site!"); // TODO: Maybe we create the site here?
+        logger.warn("Olm has no client ID!");
         return;
     }
-
     const clientId = olm.clientId;
-
     const { publicKey } = message.data;
     if (!publicKey) {
         logger.warn("Public key not provided");
         return;
     }
-
+    
+    // Get the client
     const [client] = await db
         .select()
         .from(clients)
         .where(eq(clients.clientId, clientId))
         .limit(1);
-
-    if (!client || !client.siteId) {
-        logger.warn("Site not found or does not have exit node");
+    
+    if (!client) {
+        logger.warn("Client not found");
         return;
     }
-
-    const [site] = await db
+    
+    // Get all site associations for this client
+    const clientSiteAssociations = await db
         .select()
-        .from(sites)
-        .where(eq(sites.siteId, client.siteId))
-        .limit(1);
-
-    if (!site) {
-        logger.warn("Site not found or does not have exit node");
+        .from(clientSites)
+        .where(eq(clientSites.clientId, clientId));
+    
+    if (clientSiteAssociations.length === 0) {
+        logger.warn("Client is not associated with any sites");
         return;
     }
-
-    if (!site.exitNodeId) {
-        logger.warn("Site does not have exit node");
-        return;
-    }
-
-    const [exitNode] = await db
-        .select()
-        .from(exitNodes)
-        .where(eq(exitNodes.exitNodeId, site.exitNodeId))
-        .limit(1);
-
-    sendToClient(olm.olmId, {
-        type: "olm/wg/holepunch",
-        data: {
-            serverPubKey: exitNode.publicKey,
-        }
-    });
-
-    // make sure we hand endpoints for both the site and the client and the lastHolePunch is not too old
-    if (!site.endpoint || !client.endpoint) {
-        logger.warn("Site or client has no endpoint or listen port");
-        return;
-    }
-
-    const now = new Date().getTime() / 1000;
-    if (site.lastHolePunch && now - site.lastHolePunch > 6) {
-        logger.warn("Site last hole punch is too old");
-        return;
-    }
-
-    if (client.lastHolePunch && now - client.lastHolePunch > 6) {
-        logger.warn("Client last hole punch is too old");
-        return;
-    }
-
+    
+    // Update the client's public key
     await db
         .update(clients)
         .set({
@@ -93,35 +55,103 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         })
         .where(eq(clients.clientId, olm.clientId))
         .returning();
-
-    if (client.pubKey && client.pubKey !== publicKey) {
-        logger.info("Public key mismatch. Deleting old peer...");
-        await deletePeer(site.siteId, client.pubKey);
+    
+    // Check if public key changed and handle old peer deletion later
+    const pubKeyChanged = client.pubKey && client.pubKey !== publicKey;
+    
+    // Get all sites data
+    const siteIds = clientSiteAssociations.map(cs => cs.siteId);
+    const sitesData = await db
+        .select()
+        .from(sites)
+        .where(inArray(sites.siteId, siteIds));
+    
+    // Prepare an array to store site configurations
+    const siteConfigurations = [];
+    const now = new Date().getTime() / 1000;
+    
+    // Process each site
+    for (const site of sitesData) {
+        if (!site.exitNodeId) {
+            logger.warn(`Site ${site.siteId} does not have exit node, skipping`);
+            continue;
+        }
+        
+        // Get the exit node for this site
+        const [exitNode] = await db
+            .select()
+            .from(exitNodes)
+            .where(eq(exitNodes.exitNodeId, site.exitNodeId))
+            .limit(1);
+        
+        // Validate endpoint and hole punch status
+        if (!site.endpoint) {
+            logger.warn(`Site ${site.siteId} has no endpoint, skipping`);
+            continue;
+        }
+        
+        if (site.lastHolePunch && now - site.lastHolePunch > 6) {
+            logger.warn(`Site ${site.siteId} last hole punch is too old, skipping`);
+            continue;
+        }
+        
+        if (client.lastHolePunch && now - client.lastHolePunch > 6) {
+            logger.warn("Client last hole punch is too old, skipping all sites");
+            break;
+        }
+        
+        // If public key changed, delete old peer from this site
+        if (pubKeyChanged) {
+            logger.info(`Public key mismatch. Deleting old peer from site ${site.siteId}...`);
+            await deletePeer(site.siteId, client.pubKey);
+        }
+        
+        if (!site.subnet) {
+            logger.warn(`Site ${site.siteId} has no subnet, skipping`);
+            continue;
+        }
+        
+        // Add the peer to the exit node for this site
+        await addPeer(site.siteId, {
+            publicKey: publicKey,
+            allowedIps: [client.subnet],
+            endpoint: client.endpoint
+        });
+        
+        // Add site configuration to the array
+        siteConfigurations.push({
+            siteId: site.siteId,
+            endpoint: site.endpoint,
+            publicKey: site.publicKey,
+            serverIP: site.address,
+        });
+        
+        // Send holepunch message for each site
+        sendToClient(olm.olmId, {
+            type: "olm/wg/holepunch",
+            data: {
+                serverPubKey: exitNode.publicKey,
+                siteId: site.siteId
+            }
+        });
     }
-
-    if (!site.subnet) {
-        logger.warn("Site has no subnet");
+    
+    // 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;
     }
-
-    // add the peer to the exit node
-    await addPeer(site.siteId, {
-        publicKey: publicKey,
-        allowedIps: [client.subnet],
-        endpoint: client.endpoint
-    });
-
+    
+    // Return connect message with all site configurations
     return {
         message: {
             type: "olm/wg/connect",
             data: {
-                endpoint: site.endpoint,
-                publicKey: site.publicKey,
-                serverIP: site.address!.split("/")[0],
-                tunnelIP: `${client.subnet.split("/")[0]}/${site.address!.split("/")[1]}` // put the client ip in the same subnet as the site. TODO: Is this right? Maybe we need th make .subnet work properly!
+                sites: siteConfigurations,
+                tunnelIP: client.subnet,
             }
         },
-        broadcast: false, // Send to all olms
-        excludeSender: false // Include sender in broadcast
+        broadcast: false,
+        excludeSender: false
     };
-};
+};