Fix relay ws message

server/routers/newt/peers.ts

@@ -50,3 +50,29 @@ export async function deletePeer(siteId: number, publicKey: string) {
 
     logger.info(`Deleted peer ${publicKey} from newt ${newt.newtId}`);
 }
+
+export async function updatePeer(siteId: number, publicKey: string, peer: {
+    allowedIps?: string[];
+    endpoint?: string;
+}) {
+    const [site] = await db.select().from(sites).where(eq(sites.siteId, siteId)).limit(1);
+    if (!site) {
+        throw new Error(`Exit node with ID ${siteId} not found`);
+    }
+
+    // get the newt on the site
+    const [newt] = await db.select().from(newts).where(eq(newts.siteId, siteId)).limit(1);
+    if (!newt) {
+        throw new Error(`Newt not found for site ${siteId}`);
+    }
+
+    sendToClient(newt.newtId, {
+        type: 'newt/wg/peer/update',
+        data: {
+            publicKey,
+            ...peer
+        }
+    });
+
+    logger.info(`Updated peer ${publicKey} on newt ${newt.newtId}`);
+}

server/routers/olm/handleOlmRelayMessage.ts

@@ -1,8 +1,8 @@
 import db from "@server/db";
 import { MessageHandler } from "../ws";
-import { clients, Olm, olms, sites } from "@server/db/schema";
+import { clients, clientSites, Olm, olms, sites } from "@server/db/schema";
 import { eq } from "drizzle-orm";
-import { addPeer, deletePeer } from "../newt/peers";
+import { updatePeer } from "../newt/peers";
 import logger from "@server/logger";
 
 export const handleOlmRelayMessage: MessageHandler = async (context) => {
@@ -29,17 +29,6 @@ export const handleOlmRelayMessage: MessageHandler = async (context) => {
         .where(eq(clients.clientId, clientId))
         .limit(1);
 
-    if (!client || !client.siteId) {
-        logger.warn("Site not found or does not have exit node");
-        return;
-    }
-
-    const [site] = await db
-        .select()
-        .from(sites)
-        .where(eq(sites.siteId, client.siteId))
-        .limit(1);
-
     if (!client) {
         logger.warn("Site not found or does not have exit node");
         return;
@@ -51,19 +40,22 @@ export const handleOlmRelayMessage: MessageHandler = async (context) => {
         return;
     }
 
-    if (!site.subnet) {
-        logger.warn("Site has no subnet");
-        return;
+    const sitesData = await db
+        .select()
+        .from(sites)
+        .innerJoin(clientSites, eq(sites.siteId, clientSites.siteId))
+        .where(eq(clientSites.clientId, client.clientId));
+
+    let jobs: Array<Promise<void>> = [];
+    for (const site of sitesData) {
+        // update the peer on the exit node
+        const job = updatePeer(site.sites.siteId, client.pubKey, {
+            endpoint: "" // this removes the endpoint
+        });
+        jobs.push(job);
     }
 
-    await deletePeer(site.siteId, client.pubKey);
+    await Promise.all(jobs);
 
-    // add the peer to the exit node
-    await addPeer(site.siteId, {
-        publicKey: client.pubKey,
-        allowedIps: [client.subnet],
-        endpoint: "" 
-    });
-
-    return
+    return;
 };