Fix updating sites on exit nodes

server/routers/client/updateClient.ts

@@ -1,6 +1,6 @@
 import { Request, Response, NextFunction } from "express";
 import { z } from "zod";
-import { db, exitNodes, sites } from "@server/db";
+import { Client, db, exitNodes, sites } from "@server/db";
 import { clients, clientSites } from "@server/db";
 import response from "@server/lib/response";
 import HttpCode from "@server/types/HttpCode";
@@ -105,7 +105,6 @@ export async function updateClient(
             );
         }
 
-        if (siteIds) {
         let sitesAdded = [];
         let sitesRemoved = [];
 
@@ -117,50 +116,82 @@ export async function updateClient(
 
         const existingSiteIds = existingSites.map((site) => site.siteId);
 
+        const siteIdsToProcess = siteIds || [];
         // Determine which sites were added and removed
-            sitesAdded = siteIds.filter(
+        sitesAdded = siteIdsToProcess.filter(
             (siteId) => !existingSiteIds.includes(siteId)
         );
         sitesRemoved = existingSiteIds.filter(
-                (siteId) => !siteIds.includes(siteId)
+            (siteId) => !siteIdsToProcess.includes(siteId)
         );
 
+        let updatedClient: Client | undefined = undefined;
+        let sitesData: any; // TODO: define type somehow from the query below
+        await db.transaction(async (trx) => {
+            // Update client name if provided
+            if (name) {
+                await trx
+                    .update(clients)
+                    .set({ name })
+                    .where(eq(clients.clientId, clientId));
+            }
+
+            // Update site associations if provided
+            // Remove sites that are no longer associated
+            for (const siteId of sitesRemoved) {
+                await trx
+                    .delete(clientSites)
+                    .where(
+                        and(
+                            eq(clientSites.clientId, clientId),
+                            eq(clientSites.siteId, siteId)
+                        )
+                    );
+            }
+
+            // Add new site associations
+            for (const siteId of sitesAdded) {
+                await trx.insert(clientSites).values({
+                    clientId,
+                    siteId
+                });
+            }
+
+            // Fetch the updated client
+            [updatedClient] = await trx
+                .select()
+                .from(clients)
+                .where(eq(clients.clientId, clientId))
+                .limit(1);
+
+            // get all sites for this client and join with exit nodes with site.exitNodeId
+            sitesData = await trx
+                .select()
+                .from(sites)
+                .innerJoin(clientSites, eq(sites.siteId, clientSites.siteId))
+                .leftJoin(exitNodes, eq(sites.exitNodeId, exitNodes.exitNodeId))
+                .where(eq(clientSites.clientId, client.clientId));
+        });
+
         logger.info(
             `Adding ${sitesAdded.length} new sites to client ${client.clientId}`
         );
         for (const siteId of sitesAdded) {
             if (!client.subnet || !client.pubKey) {
-                    logger.debug(
+                logger.debug("Client subnet, pubKey or endpoint is not set");
-                        "Client subnet, pubKey or endpoint is not set"
-                    );
                 continue;
             }
 
-                // TODO: WE NEED TO HANDLE THIS BETTER. RIGHT NOW WE ARE JUST GUESSING BASED ON THE OTHER SITES
+            // TODO: WE NEED TO HANDLE THIS BETTER. WE ARE DEFAULTING TO RELAYING FOR NEW SITES
             // BUT REALLY WE NEED TO TRACK THE USERS PREFERENCE THAT THEY CHOSE IN THE CLIENTS
+            // AND TRIGGER A HOLEPUNCH OR SOMETHING TO GET THE ENDPOINT AND HP TO THE NEW SITES
             const isRelayed = true;
 
-                // get the clientsite
-                const [clientSite] = await db
-                    .select()
-                    .from(clientSites)
-                    .where(
-                        and(
-                            eq(clientSites.clientId, client.clientId),
-                            eq(clientSites.siteId, siteId)
-                        )
-                    )
-                    .limit(1);
-
-                if (!clientSite || !clientSite.endpoint) {
-                    logger.debug("Client site is missing or has no endpoint");
-                    continue;
-                }
-
             const site = await newtAddPeer(siteId, {
                 publicKey: client.pubKey,
                 allowedIps: [`${client.subnet.split("/")[0]}/32`], // we want to only allow from that client
-                    endpoint: isRelayed ? "" : clientSite.endpoint
+                // endpoint: isRelayed ? "" : clientSite.endpoint
+                endpoint: isRelayed ? "" : "" // we are not HPing yet so no endpoint
             });
 
             if (!site) {
@@ -191,15 +222,13 @@ export async function updateClient(
                     .limit(1);
 
                 if (!exitNode) {
-                        logger.warn(
+                    logger.warn(`Exit node not found for site ${site.siteId}`);
-                            `Exit node not found for site ${site.siteId}`
-                        );
                     return null;
                 }
 
                 endpoint = `${exitNode.endpoint}:21820`;
             } else {
-                    if (!endpoint) {
+                if (!site.endpoint) {
                     logger.warn(
                         `Site ${site.siteId} has no endpoint, skipping`
                     );
@@ -228,57 +257,24 @@ export async function updateClient(
             }
             const site = await newtDeletePeer(siteId, client.pubKey);
             if (!site) {
-                    logger.debug(
+                logger.debug("Failed to delete peer from newt - missing site");
-                        "Failed to delete peer from newt - missing site"
-                    );
                 continue;
             }
             if (!site.endpoint || !site.publicKey) {
                 logger.debug("Site endpoint or publicKey is not set");
                 continue;
             }
-                await olmDeletePeer(
+            await olmDeletePeer(client.clientId, site.siteId, site.publicKey);
-                    client.clientId,
+        }
-                    site.siteId,
+
-                    site.publicKey
+        if (!updatedClient || !sitesData) {
+            return next(
+                createHttpError(
+                    HttpCode.INTERNAL_SERVER_ERROR,
+                    `Failed to update client`
+                )
             );
         }
-        }
-
-        await db.transaction(async (trx) => {
-            // Update client name if provided
-            if (name) {
-                await trx
-                    .update(clients)
-                    .set({ name })
-                    .where(eq(clients.clientId, clientId));
-            }
-
-            // Update site associations if provided
-            if (siteIds) {
-                // Delete existing site associations
-                await trx
-                    .delete(clientSites)
-                    .where(eq(clientSites.clientId, clientId));
-
-                // Create new site associations
-                if (siteIds.length > 0) {
-                    await trx.insert(clientSites).values(
-                        siteIds.map((siteId) => ({
-                            clientId,
-                            siteId
-                        }))
-                    );
-                }
-            }
-
-            // get all sites for this client and join with exit nodes with site.exitNodeId
-            const sitesData = await db
-                .select()
-                .from(sites)
-                .innerJoin(clientSites, eq(sites.siteId, clientSites.siteId))
-                .leftJoin(exitNodes, eq(sites.exitNodeId, exitNodes.exitNodeId))
-                .where(eq(clientSites.clientId, client.clientId));
 
         let exitNodeDestinations: {
             reachableAt: string;
@@ -316,8 +312,7 @@ export async function updateClient(
                     type: site.exitNodes?.type || "",
                     sourceIp: site.clientSites.endpoint.split(":")[0] || "",
                     sourcePort:
-                            parseInt(site.clientSites.endpoint.split(":")[1]) ||
+                        parseInt(site.clientSites.endpoint.split(":")[1]) || 0,
-                            0,
                     destinations: [
                         {
                             destinationIP: site.sites.subnet.split("/")[0],
@@ -368,13 +363,6 @@ export async function updateClient(
             });
         }
 
-            // Fetch the updated client
-            const [updatedClient] = await trx
-                .select()
-                .from(clients)
-                .where(eq(clients.clientId, clientId))
-                .limit(1);
-
         return response(res, {
             data: updatedClient,
             success: true,
@@ -382,7 +370,6 @@ export async function updateClient(
             message: "Client updated successfully",
             status: HttpCode.OK
         });
-        });
     } catch (error) {
         logger.error(error);
         return next(