Improve holepunching

server/routers/olm/getOlmToken.ts

@@ -1,9 +1,9 @@
 import { generateSessionToken } from "@server/auth/sessions/app";
-import { db } from "@server/db";
+import { clients, db, ExitNode, exitNodes, sites, clientSitesAssociationsCache } from "@server/db";
 import { olms } from "@server/db";
 import HttpCode from "@server/types/HttpCode";
 import response from "@server/lib/response";
-import { eq } from "drizzle-orm";
+import { eq, inArray } from "drizzle-orm";
 import { NextFunction, Request, Response } from "express";
 import createHttpError from "http-errors";
 import { z } from "zod";
@@ -15,11 +15,13 @@ import {
 import { verifyPassword } from "@server/auth/password";
 import logger from "@server/logger";
 import config from "@server/lib/config";
+import { listExitNodes } from "#dynamic/lib/exitNodes";
 
 export const olmGetTokenBodySchema = z.object({
     olmId: z.string(),
     secret: z.string(),
-    token: z.string().optional()
+    token: z.string().optional(),
+    orgId: z.string().optional()
 });
 
 export type OlmGetTokenBody = z.infer<typeof olmGetTokenBodySchema>;
@@ -40,7 +42,7 @@ export async function getOlmToken(
         );
     }
 
-    const { olmId, secret, token } = parsedBody.data;
+    const { olmId, secret, token, orgId } = parsedBody.data;
 
     try {
         if (token) {
@@ -61,11 +63,12 @@ export async function getOlmToken(
             }
         }
 
-        const existingOlmRes = await db
+        const [existingOlm] = await db
             .select()
             .from(olms)
             .where(eq(olms.olmId, olmId));
-        if (!existingOlmRes || !existingOlmRes.length) {
+
+        if (!existingOlm) {
             return next(
                 createHttpError(
                     HttpCode.BAD_REQUEST,
@@ -74,12 +77,11 @@ export async function getOlmToken(
             );
         }
 
-        const existingOlm = existingOlmRes[0];
-
         const validSecret = await verifyPassword(
             secret,
             existingOlm.secretHash
         );
+
         if (!validSecret) {
             if (config.getRawConfig().app.log_failed_attempts) {
                 logger.info(
@@ -96,11 +98,78 @@ export async function getOlmToken(
         const resToken = generateSessionToken();
         await createOlmSession(resToken, existingOlm.olmId);
 
+        let orgIdToUse = orgId;
+        if (!orgIdToUse) {
+            if (!existingOlm.clientId) {
+                return next(
+                    createHttpError(
+                        HttpCode.BAD_REQUEST,
+                        "Olm is not associated with a client, orgId is required"
+                    )
+                );
+            }
+
+            const [client] = await db
+                .select()
+                .from(clients)
+                .where(eq(clients.clientId, existingOlm.clientId))
+                .limit(1);
+            
+            if (!client) {
+                return next(
+                    createHttpError(
+                        HttpCode.BAD_REQUEST,
+                        "Olm's associated client not found, orgId is required"
+                    )
+                );
+            }
+
+            orgIdToUse = client.orgId;
+        }
+
+        // Get all exit nodes from sites where the client has peers
+        const clientSites = await db
+            .select()
+            .from(clientSitesAssociationsCache)
+            .innerJoin(
+                sites,
+                eq(sites.siteId, clientSitesAssociationsCache.siteId)
+            )
+            .where(eq(clientSitesAssociationsCache.clientId, existingOlm.clientId!));
+
+        // Extract unique exit node IDs
+        const exitNodeIds = Array.from(
+            new Set(
+                clientSites
+                    .map(({ sites: site }) => site.exitNodeId)
+                    .filter((id): id is number => id !== null)
+            )
+        );
+
+        let allExitNodes: ExitNode[] = [];
+        if (exitNodeIds.length > 0) {
+            allExitNodes = await db
+                .select()
+                .from(exitNodes)
+                .where(inArray(exitNodes.exitNodeId, exitNodeIds));
+        }
+
+        const exitNodesHpData = allExitNodes.map((exitNode: ExitNode) => {
+            return {
+                publicKey: exitNode.publicKey,
+                endpoint: exitNode.endpoint
+            };
+        });
+
         logger.debug("Token created successfully");
 
-        return response<{ token: string }>(res, {
+        return response<{
+            token: string;
+            exitNodes: { publicKey: string; endpoint: string }[];
+        }>(res, {
             data: {
-                token: resToken
+                token: resToken,
+                exitNodes: exitNodesHpData
             },
             success: true,
             error: false,

server/routers/olm/handleOlmPingMessage.ts

@@ -59,7 +59,7 @@ export const startOlmOfflineChecker = (): void => {
 
                 // Send a disconnect message to the client if connected
                 try {
-                    await sendTerminateClient(offlineClient.clientId); // terminate first
+                    await sendTerminateClient(offlineClient.clientId, offlineClient.olmId); // terminate first
                     // wait a moment to ensure the message is sent
                     await new Promise(resolve => setTimeout(resolve, 1000));
                     await disconnectClient(offlineClient.olmId);

server/routers/olm/handleOlmRegisterMessage.ts

@@ -34,21 +34,28 @@ import { generateRemoteSubnets } from "@server/lib/ip";
 import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations";
 import { checkOrgAccessPolicy } from "@server/lib/checkOrgAccessPolicy";
 import { validateSessionToken } from "@server/auth/sessions/app";
+import config from "@server/lib/config";
 
 export const handleOlmRegisterMessage: MessageHandler = async (context) => {
     logger.info("Handling register olm message!");
     const { message, client: c, sendToClient } = context;
     const olm = c as Olm;
 
-    const now = new Date().getTime() / 1000;
+    const now = Math.floor(Date.now() / 1000);
 
     if (!olm) {
         logger.warn("Olm not found");
         return;
     }
 
-    const { publicKey, relay, olmVersion, orgId, doNotCreateNewClient, token: userToken } =
-        message.data;
+    const {
+        publicKey,
+        relay,
+        olmVersion,
+        orgId,
+        doNotCreateNewClient,
+        token: userToken
+    } = message.data;
 
     let client: Client | undefined;
     let org: Org | undefined;
@@ -63,7 +70,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
                     olm.name || "User Device",
                     // doNotCreateNewClient ? true : false
                     true // for now never create a new client automatically because we create the users clients when they are added to the org
-                    // this means that the rebuildClientAssociationsFromClient call below issue is not a problem 
+                    // this means that the rebuildClientAssociationsFromClient call below issue is not a problem
                 );
 
             client = clientRes;
@@ -113,12 +120,14 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
             `Switching olm client ${olm.olmId} to org ${orgId} for user ${olm.userId}`
         );
 
-        await db
-            .update(olms)
-            .set({
-                clientId: client.clientId
-            })
-            .where(eq(olms.olmId, olm.olmId));
+        if (olm.clientId !== client.clientId) { // we only need to do this if the client is changing
+            await db
+                .update(olms)
+                .set({
+                    clientId: client.clientId
+                })
+                .where(eq(olms.olmId, olm.olmId));
+        }
     } else {
         if (!olm.clientId) {
             logger.warn("Olm has no client ID!");
@@ -159,41 +168,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         return;
     }
 
-    if (client.exitNodeId) {
-        // TODO: FOR NOW WE ARE JUST HOLEPUNCHING ALL EXIT NODES BUT IN THE FUTURE WE SHOULD HANDLE THIS BETTER
-
-        // Get the exit node
-        const allExitNodes = await listExitNodes(client.orgId, true); // FILTER THE ONLINE ONES
-
-        const exitNodesHpData = allExitNodes.map((exitNode: ExitNode) => {
-            return {
-                publicKey: exitNode.publicKey,
-                endpoint: exitNode.endpoint
-            };
-        });
-
-        // Send holepunch message
-        await sendToClient(olm.olmId, {
-            type: "olm/wg/holepunch/all",
-            data: {
-                exitNodes: exitNodesHpData
-            }
-        });
-
-        if (!olmVersion) {
-            // THIS IS FOR BACKWARDS COMPATIBILITY
-            // THE OLDER CLIENTS DID NOT SEND THE VERSION
-            await sendToClient(olm.olmId, {
-                type: "olm/wg/holepunch",
-                data: {
-                    serverPubKey: allExitNodes[0].publicKey,
-                    endpoint: allExitNodes[0].endpoint
-                }
-            });
-        }
-    }
-
-    if (olmVersion) {
+    if (olmVersion && olm.version !== olmVersion) {
         await db
             .update(olms)
             .set({
@@ -202,10 +177,14 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
             .where(eq(olms.olmId, olm.olmId));
     }
 
-    // if (now - (client.lastHolePunch || 0) > 6) {
-    //     logger.warn("Client last hole punch is too old, skipping all sites");
-    //     return;
-    // }
+    // this prevents us from accepting a register from an olm that has not hole punched yet.
+    // the olm will pump the register so we can keep checking
+    if (now - (client.lastHolePunch || 0) > 5) {
+        logger.warn(
+            "Client last hole punch is too old; skipping this register"
+        );
+        return;
+    }
 
     if (client.pubKey !== publicKey) {
         logger.info(
@@ -319,7 +298,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
                 logger.warn(`Exit node not found for site ${site.siteId}`);
                 continue;
             }
-            endpoint = `${exitNode.endpoint}:21820`;
+            endpoint = `${exitNode.endpoint}:${config.getRawConfig().gerbil.clients_start_port}`;
         }
 
         const allSiteResources = await db // only get the site resources that this client has access to

server/routers/olm/handleOlmRelayMessage.ts

@@ -2,7 +2,7 @@ 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 } from "../newt/peers";
+import { updatePeer as newtUpdatePeer } from "../newt/peers";
 import logger from "@server/logger";
 
 export const handleOlmRelayMessage: MessageHandler = async (context) => {
@@ -79,18 +79,20 @@ export const handleOlmRelayMessage: MessageHandler = async (context) => {
         );
 
     // update the peer on the exit node
-    await updatePeer(siteId, client.pubKey, {
-        endpoint: "" // this removes the endpoint
+    await newtUpdatePeer(siteId, client.pubKey, {
+        endpoint: "" // this removes the endpoint so the exit node knows to relay
     });
 
-    sendToClient(olm.olmId, {
-        type: "olm/wg/peer/relay",
+    return {
+        message: {
+         type: "olm/wg/peer/relay",
         data: {
             siteId: siteId,
             endpoint: exitNode.endpoint,
             publicKey: exitNode.publicKey
         }
-    });
-
-    return;
+        },
+        broadcast: false,
+        excludeSender: false
+    };
 };

server/routers/olm/handleOlmServerPeerAddMessage.ts

@@ -0,0 +1,185 @@
+import {
+    Client,
+    clientSiteResourcesAssociationsCache,
+    db,
+    ExitNode,
+    Org,
+    orgs,
+    roleClients,
+    roles,
+    siteResources,
+    Transaction,
+    userClients,
+    userOrgs,
+    users
+} from "@server/db";
+import { MessageHandler } from "@server/routers/ws";
+import {
+    clients,
+    clientSitesAssociationsCache,
+    exitNodes,
+    Olm,
+    olms,
+    sites
+} from "@server/db";
+import { and, eq, inArray, isNotNull, isNull } from "drizzle-orm";
+import { addPeer, deletePeer } from "../newt/peers";
+import logger from "@server/logger";
+import { listExitNodes } from "#dynamic/lib/exitNodes";
+import {
+    generateAliasConfig,
+    getNextAvailableClientSubnet
+} from "@server/lib/ip";
+import { generateRemoteSubnets } from "@server/lib/ip";
+import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations";
+import { checkOrgAccessPolicy } from "@server/lib/checkOrgAccessPolicy";
+import { validateSessionToken } from "@server/auth/sessions/app";
+import config from "@server/lib/config";
+import {
+    addPeer as newtAddPeer,
+    deletePeer as newtDeletePeer
+} from "@server/routers/newt/peers";
+
+export const handleOlmServerPeerAddMessage: MessageHandler = async (
+    context
+) => {
+    logger.info("Handling register olm message!");
+    const { message, client: c, sendToClient } = context;
+    const olm = c as Olm;
+
+    const now = Math.floor(Date.now() / 1000);
+
+    if (!olm) {
+        logger.warn("Olm not found");
+        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.error(
+            `handleOlmServerPeerAddMessage: Site with ID ${siteId} not found`
+        );
+        return;
+    }
+
+    if (!site.endpoint) {
+        logger.error(
+            `handleOlmServerPeerAddMessage: Site with ID ${siteId} has no endpoint`
+        );
+        return;
+    }
+
+    // get the client
+
+    if (!olm.clientId) {
+        logger.error(
+            `handleOlmServerPeerAddMessage: Olm with ID ${olm.olmId} has no clientId`
+        );
+        return;
+    }
+
+    const [client] = await db
+        .select()
+        .from(clients)
+        .where(and(eq(clients.clientId, olm.clientId)))
+        .limit(1);
+
+    if (!client) {
+        logger.error(
+            `handleOlmServerPeerAddMessage: Client with ID ${olm.clientId} not found`
+        );
+        return;
+    }
+
+    if (!client.pubKey) {
+        logger.error(
+            `handleOlmServerPeerAddMessage: Client with ID ${client.clientId} has no public key`
+        );
+        return;
+    }
+
+    let endpoint: string | null = null;
+
+
+    const currentSessionSiteAssociationCaches = await db
+        .select()
+        .from(clientSitesAssociationsCache)
+        .where(
+            and(
+                eq(clientSitesAssociationsCache.clientId, client.clientId),
+                isNotNull(clientSitesAssociationsCache.endpoint),
+                eq(clientSitesAssociationsCache.publicKey, client.pubKey) // limit it to the current session its connected with otherwise the endpoint could be stale
+            )
+        );
+
+    // pick an endpoint
+    for (const assoc of currentSessionSiteAssociationCaches) {
+        if (assoc.endpoint) {
+            endpoint = assoc.endpoint;
+            break;
+        }
+    }
+
+    if (!endpoint) {
+        logger.error(
+            `handleOlmServerPeerAddMessage: No endpoint found for client ${client.clientId}`
+        );
+        return;
+    }
+
+    await newtAddPeer(siteId, {
+        publicKey: client.pubKey,
+        allowedIps: [`${client.subnet.split("/")[0]}/32`], // we want to only allow from that client
+        endpoint: endpoint // this is the client's endpoint with reference to the site's exit node
+    });
+
+    const allSiteResources = await db // only get the site resources that this client has access to
+        .select()
+        .from(siteResources)
+        .innerJoin(
+            clientSiteResourcesAssociationsCache,
+            eq(
+                siteResources.siteResourceId,
+                clientSiteResourcesAssociationsCache.siteResourceId
+            )
+        )
+        .where(
+            and(
+                eq(siteResources.siteId, site.siteId),
+                eq(
+                    clientSiteResourcesAssociationsCache.clientId,
+                    client.clientId
+                )
+            )
+        );
+
+    // Return connect message with all site configurations
+    return {
+        message: {
+            type: "olm/wg/peer/add",
+            data: {
+                siteId: site.siteId,
+                endpoint: site.endpoint,
+                publicKey: site.publicKey,
+                serverIP: site.address,
+                serverPort: site.listenPort,
+                remoteSubnets: generateRemoteSubnets(
+                    allSiteResources.map(({ siteResources }) => siteResources)
+                ),
+                aliases: generateAliasConfig(
+                    allSiteResources.map(({ siteResources }) => siteResources)
+                )
+            }
+        },
+        broadcast: false,
+        excludeSender: false
+    };
+};

server/routers/olm/index.ts

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

server/routers/olm/peers.ts

@@ -3,6 +3,7 @@ import { clients, olms, newts, sites } from "@server/db";
 import { eq } from "drizzle-orm";
 import { sendToClient } from "#dynamic/routers/ws";
 import logger from "@server/logger";
+import { exit } from "process";
 
 export async function addPeer(
     clientId: number,
@@ -110,3 +111,40 @@ export async function updatePeer(
 
     logger.info(`Added peer ${peer.publicKey} to olm ${olmId}`);
 }
+
+export async function initPeerAddHandshake(
+    clientId: number,
+    peer: {
+        siteId: number;
+        exitNode: {
+            publicKey: string;
+            endpoint: string;
+        };
+    },
+    olmId?: string
+) {
+    if (!olmId) {
+        const [olm] = await db
+            .select()
+            .from(olms)
+            .where(eq(olms.clientId, clientId))
+            .limit(1);
+        if (!olm) {
+            throw new Error(`Olm with ID ${clientId} not found`);
+        }
+        olmId = olm.olmId;
+    }
+
+    await sendToClient(olmId, {
+        type: "olm/wg/peer/holepunch/site/add",
+        data: {
+            siteId: peer.siteId,
+            exitNode: {
+                publicKey: peer.exitNode.publicKey,
+                endpoint: peer.exitNode.endpoint
+            }
+        }
+    });
+
+    logger.info(`Initiated peer add handshake for site ${peer.siteId} to olm ${olmId}`);
+}