Handle unrelay and relaying better

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",
-        data: {
-            siteId: siteId,
-            endpoint: exitNode.endpoint,
-            publicKey: exitNode.publicKey
-        }
+            type: "olm/wg/peer/relay",
+            data: {
+                siteId: siteId,
+                relayEndpoint: exitNode.endpoint
+            }
         },
         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

@@ -7,4 +7,5 @@ export * from "./deleteUserOlm";
 export * from "./listUserOlms";
 export * from "./deleteUserOlm";
 export * from "./getUserOlm";
-export * from "./handleOlmServerPeerAddMessage";
+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,