Handle unrelay and relaying better

server/lib/rebuildClientAssociations.ts

@@ -1011,76 +1011,18 @@ async function handleMessagesForClientSites(
                 continue;
             }
 
-            // Add peer to newt
+            await holepunchSiteAdd(
-            const isRelayed = true; // Default to relaying for new connections
+                // this will kick off the add peer process for the client
-            newtJobs.push(
+                client.clientId,
-                newtAddPeer(
+                {
-                    site.siteId,
+                    siteId: site.siteId,
-                    {
+                    exitNode: {
-                        publicKey: client.pubKey,
+                        publicKey: exitNode.publicKey,
-                        allowedIps: [`${client.subnet.split("/")[0]}/32`],
+                        endpoint: exitNode.endpoint
-                        endpoint: isRelayed ? "" : ""
+                    }
-                    },
+                },
-                    newt.newtId
+                olmId
-                )
             );
-
-            // Get all site resources for this site that the client has access to
-            const accessibleResources = await trx
-                .select()
-                .from(siteResources)
-                .innerJoin(
-                    clientSiteResourcesAssociationsCache,
-                    eq(
-                        siteResources.siteResourceId,
-                        clientSiteResourcesAssociationsCache.siteResourceId
-                    )
-                )
-                .where(
-                    and(
-                        eq(siteResources.siteId, site.siteId),
-                        eq(
-                            clientSiteResourcesAssociationsCache.clientId,
-                            client.clientId
-                        )
-                    )
-                );
-            try {
-                // Add peer to olm
-                olmJobs.push(
-                    olmAddPeer(
-                        client.clientId,
-                        {
-                            siteId: site.siteId,
-                            endpoint:
-                                isRelayed || !site.endpoint
-                                    ? `${exitNode.endpoint}:21820`
-                                    : site.endpoint,
-                            publicKey: site.publicKey,
-                            serverIP: site.address || "",
-                            serverPort: site.listenPort || 0,
-                            remoteSubnets: generateRemoteSubnets(
-                                accessibleResources.map(
-                                    ({ siteResources }) => siteResources
-                                )
-                            )
-                        },
-                        olmId
-                    )
-                );
-            } catch (error) {
-                // if the error includes not found then its just because the olm does not exist anymore or yet and its fine if we dont send
-                if (
-                    error instanceof Error &&
-                    error.message.includes("not found")
-                ) {
-                    logger.debug(
-                        `Olm data not found for client ${client.clientId}, skipping removal`
-                    );
-                } else {
-                    throw error;
-                }
-            }
         }
 
         // Update exit node destinations

server/routers/olm/handleOlmRegisterMessage.ts

@@ -237,7 +237,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
             );
         }
 
-        let endpoint = site.endpoint;
+        let relayEndpoint: string | undefined = undefined;
         if (relay) {
             const [exitNode] = await db
                 .select()
@@ -248,7 +248,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
                 logger.warn(`Exit node not found for site ${site.siteId}`);
                 continue;
             }
-            endpoint = `${exitNode.endpoint}:${config.getRawConfig().gerbil.clients_start_port}`;
+            relayEndpoint = `${exitNode.endpoint}:${config.getRawConfig().gerbil.clients_start_port}`;
         }
 
         const allSiteResources = await db // only get the site resources that this client has access to
@@ -274,7 +274,8 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         // Add site configuration to the array
         siteConfigurations.push({
             siteId: site.siteId,
-            endpoint: endpoint,
+            relayEndpoint: relayEndpoint, // this can be undefined now if not relayed
+            endpoint: site.endpoint,
             publicKey: site.publicKey,
             serverIP: site.address,
             serverPort: site.listenPort,

server/routers/olm/handleOlmRelayMessage.ts

@@ -85,12 +85,11 @@ export const handleOlmRelayMessage: MessageHandler = async (context) => {
 
     return {
         message: {
-         type: "olm/wg/peer/relay",
+            type: "olm/wg/peer/relay",
-        data: {
+            data: {
-            siteId: siteId,
+                siteId: siteId,
-            endpoint: exitNode.endpoint,
+                relayEndpoint: exitNode.endpoint
-            publicKey: exitNode.publicKey
+            }
-        }
         },
         broadcast: false,
         excludeSender: false

server/routers/olm/handleOlmServerPeerAddMessage.ts

@@ -135,6 +135,8 @@ export const handleOlmServerPeerAddMessage: MessageHandler = async (
         return;
     }
 
+    // NOTE: here we are always starting direct to the peer and will relay later
+
     await newtAddPeer(siteId, {
         publicKey: client.pubKey,
         allowedIps: [`${client.subnet.split("/")[0]}/32`], // we want to only allow from that client

server/routers/olm/handleOlmUnRelayMessage.ts

@@ -0,0 +1,96 @@
+import { db, exitNodes, sites } from "@server/db";
+import { MessageHandler } from "@server/routers/ws";
+import { clients, clientSitesAssociationsCache, Olm } from "@server/db";
+import { and, eq } from "drizzle-orm";
+import { updatePeer as newtUpdatePeer } from "../newt/peers";
+import logger from "@server/logger";
+
+export const handleOlmUnRelayMessage: MessageHandler = async (context) => {
+    const { message, client: c, sendToClient } = context;
+    const olm = c as Olm;
+
+    logger.info("Handling unrelay 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?
+        return;
+    }
+
+    const clientId = olm.clientId;
+
+    const [client] = await db
+        .select()
+        .from(clients)
+        .where(eq(clients.clientId, clientId))
+        .limit(1);
+
+    if (!client) {
+        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("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) {
+        logger.warn("Site not found or has no exit node");
+        return;
+    }
+
+    const [clientSiteAssociation] = await db
+        .update(clientSitesAssociationsCache)
+        .set({
+            isRelayed: false
+        })
+        .where(
+            and(
+                eq(clientSitesAssociationsCache.clientId, olm.clientId),
+                eq(clientSitesAssociationsCache.siteId, siteId)
+            )
+        )
+        .returning();
+
+    if (!clientSiteAssociation) {
+        logger.warn("Client-Site association not found");
+        return;
+    }
+
+    if (!clientSiteAssociation.endpoint) {
+        logger.warn("Client-Site association has no endpoint, cannot unrelay");
+        return;
+    }
+
+    // update the peer on the exit node
+    await newtUpdatePeer(siteId, client.pubKey, {
+        endpoint: clientSiteAssociation.endpoint // this is the endpoint of the client to connect directly to the exit node
+    });
+
+    return {
+        message: {
+            type: "olm/wg/peer/unrelay",
+            data: {
+                siteId: siteId,
+                endpoint: site.endpoint
+            }
+        },
+        broadcast: false,
+        excludeSender: false
+    };
+};

server/routers/olm/index.ts

@@ -8,3 +8,4 @@ export * from "./listUserOlms";
 export * from "./deleteUserOlm";
 export * from "./getUserOlm";
 export * from "./handleOlmServerPeerAddMessage";
+export * from "./handleOlmUnRelayMessage";

server/routers/olm/peers.ts

@@ -103,6 +103,7 @@ export async function updatePeer(
             siteId: peer.siteId,
             publicKey: peer.publicKey,
             endpoint: peer.endpoint,
+            relayEndpoint: peer.serverIP,
             serverIP: peer.serverIP,
             serverPort: peer.serverPort,
             remoteSubnets: peer.remoteSubnets

server/routers/org/deleteOrg.ts

@@ -193,7 +193,7 @@ export async function deleteOrg(
         // Send termination messages outside of transaction to prevent blocking
         for (const newtId of deletedNewtIds) {
             const payload = {
-                type: `newt/terminate`,
+                type: `newt/wg/terminate`,
                 data: {}
             };
             // Don't await this to prevent blocking the response

server/routers/ws/messageHandlers.ts

@@ -12,7 +12,8 @@ import {
     handleOlmRelayMessage,
     handleOlmPingMessage,
     startOlmOfflineChecker,
-    handleOlmServerPeerAddMessage
+    handleOlmServerPeerAddMessage,
+    handleOlmUnRelayMessage
 } from "../olm";
 import { handleHealthcheckStatusMessage } from "../target";
 import { MessageHandler } from "./types";
@@ -21,6 +22,7 @@ export const messageHandlers: Record<string, MessageHandler> = {
     "olm/wg/server/peer/add": handleOlmServerPeerAddMessage,
     "olm/wg/register": handleOlmRegisterMessage,
     "olm/wg/relay": handleOlmRelayMessage,
+    "olm/wg/unrelay": handleOlmUnRelayMessage,
     "olm/ping": handleOlmPingMessage,
     "newt/wg/register": handleNewtRegisterMessage,
     "newt/wg/get-config": handleGetConfigMessage,