From 152fb47ca4cd4cdbf4dd9e7c0297e8804ea80663 Mon Sep 17 00:00:00 2001 From: Owen Date: Tue, 2 Dec 2025 11:17:08 -0500 Subject: [PATCH] Handle unrelay and relaying better --- server/lib/rebuildClientAssociations.ts | 80 +++------------- .../routers/olm/handleOlmRegisterMessage.ts | 7 +- server/routers/olm/handleOlmRelayMessage.ts | 11 +-- .../olm/handleOlmServerPeerAddMessage.ts | 2 + server/routers/olm/handleOlmUnRelayMessage.ts | 96 +++++++++++++++++++ server/routers/olm/index.ts | 3 +- server/routers/olm/peers.ts | 1 + server/routers/org/deleteOrg.ts | 2 +- server/routers/ws/messageHandlers.ts | 4 +- 9 files changed, 125 insertions(+), 81 deletions(-) create mode 100644 server/routers/olm/handleOlmUnRelayMessage.ts diff --git a/server/lib/rebuildClientAssociations.ts b/server/lib/rebuildClientAssociations.ts index 810acdef..00156c01 100644 --- a/server/lib/rebuildClientAssociations.ts +++ b/server/lib/rebuildClientAssociations.ts @@ -1011,76 +1011,18 @@ async function handleMessagesForClientSites( continue; } - // Add peer to newt - const isRelayed = true; // Default to relaying for new connections - newtJobs.push( - newtAddPeer( - site.siteId, - { - publicKey: client.pubKey, - allowedIps: [`${client.subnet.split("/")[0]}/32`], - endpoint: isRelayed ? "" : "" - }, - newt.newtId - ) + await holepunchSiteAdd( + // this will kick off the add peer process for the client + client.clientId, + { + siteId: site.siteId, + exitNode: { + publicKey: exitNode.publicKey, + endpoint: exitNode.endpoint + } + }, + 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 diff --git a/server/routers/olm/handleOlmRegisterMessage.ts b/server/routers/olm/handleOlmRegisterMessage.ts index 7cde2c76..696da748 100644 --- a/server/routers/olm/handleOlmRegisterMessage.ts +++ b/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, diff --git a/server/routers/olm/handleOlmRelayMessage.ts b/server/routers/olm/handleOlmRelayMessage.ts index 5479ccbb..595b35ba 100644 --- a/server/routers/olm/handleOlmRelayMessage.ts +++ b/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 diff --git a/server/routers/olm/handleOlmServerPeerAddMessage.ts b/server/routers/olm/handleOlmServerPeerAddMessage.ts index 3d0d61b2..2e5009eb 100644 --- a/server/routers/olm/handleOlmServerPeerAddMessage.ts +++ b/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 diff --git a/server/routers/olm/handleOlmUnRelayMessage.ts b/server/routers/olm/handleOlmUnRelayMessage.ts new file mode 100644 index 00000000..5f47a095 --- /dev/null +++ b/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 + }; +}; diff --git a/server/routers/olm/index.ts b/server/routers/olm/index.ts index 0fc65d92..e671dd42 100644 --- a/server/routers/olm/index.ts +++ b/server/routers/olm/index.ts @@ -7,4 +7,5 @@ export * from "./deleteUserOlm"; export * from "./listUserOlms"; export * from "./deleteUserOlm"; export * from "./getUserOlm"; -export * from "./handleOlmServerPeerAddMessage"; \ No newline at end of file +export * from "./handleOlmServerPeerAddMessage"; +export * from "./handleOlmUnRelayMessage"; \ No newline at end of file diff --git a/server/routers/olm/peers.ts b/server/routers/olm/peers.ts index 7651f0a9..87c634cc 100644 --- a/server/routers/olm/peers.ts +++ b/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 diff --git a/server/routers/org/deleteOrg.ts b/server/routers/org/deleteOrg.ts index 098c5c41..35dc7503 100644 --- a/server/routers/org/deleteOrg.ts +++ b/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 diff --git a/server/routers/ws/messageHandlers.ts b/server/routers/ws/messageHandlers.ts index b92e7530..acd1aef0 100644 --- a/server/routers/ws/messageHandlers.ts +++ b/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 = { "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,