diff --git a/server/lib/calculateUserClientsForOrgs.ts b/server/lib/calculateUserClientsForOrgs.ts index 39585500a..5770029de 100644 --- a/server/lib/calculateUserClientsForOrgs.ts +++ b/server/lib/calculateUserClientsForOrgs.ts @@ -339,19 +339,6 @@ export async function calculateUserClientsForOrgs( continue; } - // Get exit nodes for this org - const exitNodesList = await getExitNodes(orgId); - - if (exitNodesList.length === 0) { - logger.warn( - `Skipping org ${orgId} for OLM ${olm.olmId} (user ${userId}): no exit nodes found` - ); - continue; - } - - const randomExitNode = - exitNodesList[Math.floor(Math.random() * exitNodesList.length)]; - // Get next available subnet const { value: newSubnet, release: releaseSubnetLock } = await getNextAvailableClientSubnet(orgId, trx); @@ -370,7 +357,6 @@ export async function calculateUserClientsForOrgs( const newClientData: InferInsertModel = { userId, orgId: userOrg.orgId, - exitNodeId: randomExitNode.exitNodeId, name: olm.name || "User Client", subnet: updatedSubnet, olmId: olm.olmId, diff --git a/server/lib/exitNodes/calculateExitNodeWeight.ts b/server/lib/exitNodes/calculateExitNodeWeight.ts new file mode 100644 index 000000000..a740e7a89 --- /dev/null +++ b/server/lib/exitNodes/calculateExitNodeWeight.ts @@ -0,0 +1,40 @@ +import { db, sites, clients } from "@server/db"; +import { and, eq, count } from "drizzle-orm"; + +// (MAX_CONNECTIONS - current_connections) / MAX_CONNECTIONS) +// higher = more desirable +// like saying, this node has x% of its capacity left +export async function calculateExitNodeWeight( + exitNodeId: number, + maxConnections: number | null | undefined +): Promise { + if (maxConnections === null || maxConnections === undefined) { + return 1; + } + + const [[siteConnections], [clientConnections]] = await Promise.all([ + db + .select({ count: count() }) + .from(sites) + .where( + and(eq(sites.exitNodeId, exitNodeId), eq(sites.online, true)) + ), + db + .select({ count: count() }) + .from(clients) + .where( + and( + eq(clients.exitNodeId, exitNodeId), + eq(clients.online, true) + ) + ) + ]); + + const currentConnections = siteConnections.count + clientConnections.count; + + if (currentConnections >= maxConnections) { + return null; + } + + return (maxConnections - currentConnections) / maxConnections; +} diff --git a/server/lib/exitNodes/exitNodes.ts b/server/lib/exitNodes/exitNodes.ts index f405a1114..823fdde97 100644 --- a/server/lib/exitNodes/exitNodes.ts +++ b/server/lib/exitNodes/exitNodes.ts @@ -1,6 +1,5 @@ import { db, exitNodes, Transaction } from "@server/db"; import logger from "@server/logger"; -import { ExitNodePingResult } from "@server/routers/newt"; import { eq } from "drizzle-orm"; export async function verifyExitNodeOrgAccess( @@ -52,6 +51,16 @@ export async function listExitNodes( return allExitNodes; } +export type ExitNodePingResult = { + exitNodeId: number; + latencyMs: number; + weight: number; + error?: string; + exitNodeName: string; + endpoint: string; + wasPreviouslyConnected: boolean; +}; + export function selectBestExitNode( pingResults: ExitNodePingResult[] ): ExitNodePingResult | null { diff --git a/server/lib/exitNodes/index.ts b/server/lib/exitNodes/index.ts index d1477a683..459995074 100644 --- a/server/lib/exitNodes/index.ts +++ b/server/lib/exitNodes/index.ts @@ -2,3 +2,4 @@ export * from "./exitNodes"; export * from "./exitNodeComms"; export * from "./subnet"; export * from "./getCurrentExitNodeId"; +export * from "./calculateExitNodeWeight"; diff --git a/server/private/lib/exitNodes/exitNodes.ts b/server/private/lib/exitNodes/exitNodes.ts index 1f9517725..976f52cff 100644 --- a/server/private/lib/exitNodes/exitNodes.ts +++ b/server/private/lib/exitNodes/exitNodes.ts @@ -25,7 +25,6 @@ import { Transaction } from "@server/db"; import logger from "@server/logger"; -import { ExitNodePingResult } from "@server/routers/newt"; import { eq, and, or, ne, isNull, inArray } from "drizzle-orm"; import axios from "axios"; import config from "../config"; @@ -330,6 +329,16 @@ export async function listExitNodes( return exitNodesList; } +export type ExitNodePingResult = { + exitNodeId: number; + latencyMs: number; + weight: number; + error?: string; + exitNodeName: string; + endpoint: string; + wasPreviouslyConnected: boolean; +}; + /** * Selects the most suitable exit node from a list of ping results. * diff --git a/server/routers/client/createClient.ts b/server/routers/client/createClient.ts index 11c7c9ec2..d84a30301 100644 --- a/server/routers/client/createClient.ts +++ b/server/routers/client/createClient.ts @@ -255,20 +255,6 @@ export async function createClient( let newClient: Client | null = null; await db.transaction(async (trx) => { - // TODO: more intelligent way to pick the exit node - const exitNodesList = await listExitNodes(orgId); - const randomExitNode = - exitNodesList[Math.floor(Math.random() * exitNodesList.length)]; - - if (!randomExitNode) { - return next( - createHttpError( - HttpCode.NOT_FOUND, - `No exit nodes available. ${build == "saas" ? "Please contact support." : "You need to install gerbil to use the clients."}` - ) - ); - } - const [adminRole] = await trx .select() .from(roles) @@ -287,7 +273,6 @@ export async function createClient( .insert(clients) .values({ niceId, - exitNodeId: randomExitNode.exitNodeId, orgId, name, subnet: updatedSubnet, diff --git a/server/routers/client/createUserClient.ts b/server/routers/client/createUserClient.ts index 70027e4a5..bc1b25ebf 100644 --- a/server/routers/client/createUserClient.ts +++ b/server/routers/client/createUserClient.ts @@ -222,11 +222,6 @@ export async function createUserClient( let newClient: Client | null = null; await db.transaction(async (trx) => { - // TODO: more intelligent way to pick the exit node - const exitNodesList = await listExitNodes(orgId); - const randomExitNode = - exitNodesList[Math.floor(Math.random() * exitNodesList.length)]; - const [adminRole] = await trx .select() .from(roles) @@ -244,7 +239,6 @@ export async function createUserClient( [newClient] = await trx .insert(clients) .values({ - exitNodeId: randomExitNode.exitNodeId, orgId, niceId, name, diff --git a/server/routers/newt/handleNewtPingRequestMessage.ts b/server/routers/newt/handleNewtExitNodesRequestMessage.ts similarity index 64% rename from server/routers/newt/handleNewtPingRequestMessage.ts rename to server/routers/newt/handleNewtExitNodesRequestMessage.ts index f239dc4de..1adfd9f1c 100644 --- a/server/routers/newt/handleNewtPingRequestMessage.ts +++ b/server/routers/newt/handleNewtExitNodesRequestMessage.ts @@ -2,14 +2,17 @@ import { db, sites } from "@server/db"; import { MessageHandler } from "@server/routers/ws"; import { exitNodes, Newt } from "@server/db"; import logger from "@server/logger"; -import { ne, eq, or, and, count } from "drizzle-orm"; +import { eq } from "drizzle-orm"; import { listExitNodes } from "#dynamic/lib/exitNodes"; +import { calculateExitNodeWeight } from "@server/lib/exitNodes"; -export const handleNewtPingRequestMessage: MessageHandler = async (context) => { +export const handleNewtExitNodesRequestMessage: MessageHandler = async ( + context +) => { const { message, client, sendToClient } = context; const newt = client as Newt; - logger.info("Handling ping request newt message!"); + logger.info("Handling exit nodes request newt message!"); if (!newt) { logger.warn("Newt not found"); @@ -54,32 +57,13 @@ export const handleNewtPingRequestMessage: MessageHandler = async (context) => { const exitNodesPayload = await Promise.all( exitNodesList.map(async (node) => { - // (MAX_CONNECTIONS - current_connections) / MAX_CONNECTIONS) - // higher = more desirable - // like saying, this node has x% of its capacity left + const weight = await calculateExitNodeWeight( + node.exitNodeId, + node.maxConnections + ); - let weight = 1; - const maxConnections = node.maxConnections; - if (maxConnections !== null && maxConnections !== undefined) { - const [currentConnections] = await db - .select({ - count: count() - }) - .from(sites) - .where( - and( - eq(sites.exitNodeId, node.exitNodeId), - eq(sites.online, true) - ) - ); - - if (currentConnections.count >= maxConnections) { - return null; - } - - weight = - (maxConnections - currentConnections.count) / - maxConnections; + if (weight === null) { + return null; } return { diff --git a/server/routers/newt/handleNewtRegisterMessage.ts b/server/routers/newt/handleNewtRegisterMessage.ts index 0dc8380c8..27956e845 100644 --- a/server/routers/newt/handleNewtRegisterMessage.ts +++ b/server/routers/newt/handleNewtRegisterMessage.ts @@ -7,6 +7,7 @@ import logger from "@server/logger"; import config from "@server/lib/config"; import { findNextAvailableCidr } from "@server/lib/ip"; import { + ExitNodePingResult, selectBestExitNode, verifyExitNodeOrgAccess } from "#dynamic/lib/exitNodes"; @@ -15,16 +16,6 @@ import { lockManager } from "#dynamic/lib/lock"; import { buildTargetConfigurationForNewtClient } from "./buildConfiguration"; import { canCompress } from "@server/lib/clientVersionChecks"; -export type ExitNodePingResult = { - exitNodeId: number; - latencyMs: number; - weight: number; - error?: string; - exitNodeName: string; - endpoint: string; - wasPreviouslyConnected: boolean; -}; - export const handleNewtRegisterMessage: MessageHandler = async (context) => { const { message, client, sendToClient } = context; const newt = client as Newt; diff --git a/server/routers/newt/index.ts b/server/routers/newt/index.ts index 1f6090358..5665a5420 100644 --- a/server/routers/newt/index.ts +++ b/server/routers/newt/index.ts @@ -5,7 +5,7 @@ export * from "./handleNewtRegisterMessage"; export * from "./handleReceiveBandwidthMessage"; export * from "./handleNewtGetConfigMessage"; export * from "./handleSocketMessages"; -export * from "./handleNewtPingRequestMessage"; +export * from "./handleNewtExitNodesRequestMessage"; export * from "./handleApplyBlueprintMessage"; export * from "./handleNewtPingMessage"; export * from "./handleNewtDisconnectingMessage"; diff --git a/server/routers/olm/handleOlmExitNodesRequestMessage.ts b/server/routers/olm/handleOlmExitNodesRequestMessage.ts new file mode 100644 index 000000000..bb8c7a48f --- /dev/null +++ b/server/routers/olm/handleOlmExitNodesRequestMessage.ts @@ -0,0 +1,93 @@ +import { db, clients } from "@server/db"; +import { MessageHandler } from "@server/routers/ws"; +import { exitNodes, Olm } from "@server/db"; +import logger from "@server/logger"; +import { eq } from "drizzle-orm"; +import { listExitNodes } from "#dynamic/lib/exitNodes"; +import { calculateExitNodeWeight } from "@server/lib/exitNodes"; + +export const handleOlmExitNodesRequestMessage: MessageHandler = async ( + context +) => { + const { message, client: olmClient, sendToClient } = context; + const olm = olmClient as Olm; + + logger.info("Handling exit nodes request olm message!"); + + if (!olm) { + logger.warn("olm not found"); + return; + } + + // Get the olm's orgId through the client relationship + if (!olm.clientId) { + logger.warn("olm clientId not found"); + return; + } + + const [client] = await db + .select({ orgId: clients.orgId }) + .from(clients) + .where(eq(clients.clientId, olm.clientId)) + .limit(1); + + if (!client || !client.orgId) { + logger.warn("client not found"); + return; + } + + const { noCloud, chainId } = message.data; + + const exitNodesList = await listExitNodes( + client.orgId, + true, + noCloud || false, + olm.clientId + ); // filter for only the online ones + + let lastExitNodeId = null; + if (olm.clientId) { + const [lastExitNode] = await db + .select() + .from(clients) + .where(eq(clients.clientId, olm.clientId)) + .limit(1); + lastExitNodeId = lastExitNode?.exitNodeId || null; + } + + const exitNodesPayload = await Promise.all( + exitNodesList.map(async (node) => { + const weight = await calculateExitNodeWeight( + node.exitNodeId, + node.maxConnections + ); + + if (weight === null) { + return null; + } + + return { + exitNodeId: node.exitNodeId, + exitNodeName: node.name, + endpoint: node.endpoint, + weight, + wasPreviouslyConnected: node.exitNodeId === lastExitNodeId + }; + }) + ); + + // filter out null values + const filteredExitNodes = exitNodesPayload.filter((node) => node !== null); + + return { + message: { + type: "olm/ping/exitNodes", + data: { + exitNodes: filteredExitNodes, + chainId: chainId + } + }, + broadcast: false, // Send to all clients + excludeSender: false // Include sender in broadcast + }; +}; diff --git a/server/routers/olm/handleOlmRegisterMessage.ts b/server/routers/olm/handleOlmRegisterMessage.ts index d386fe74e..6b0002b83 100644 --- a/server/routers/olm/handleOlmRegisterMessage.ts +++ b/server/routers/olm/handleOlmRegisterMessage.ts @@ -22,6 +22,7 @@ import { canCompress } from "@server/lib/clientVersionChecks"; import config from "@server/lib/config"; import cache from "#dynamic/lib/cache"; // not using regional here because we need this in the register message handler before we know where the client is import { waitForClientRebuildIdle } from "@server/lib/rebuildClientAssociations"; +import { ExitNodePingResult, selectBestExitNode } from "#dynamic/lib/exitNodes"; const HOLEPUNCH_STALE_CHAIN_THRESHOLD = 18; const HOLEPUNCH_STALE_CHAIN_TTL_SECONDS = 1800; @@ -49,6 +50,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { olmAgent, orgId, userToken, + pingResults, fingerprint, postures, chainId @@ -284,7 +286,22 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { return; } - if (client.pubKey !== publicKey || client.archived) { + let exitNodeId: number | undefined; + if (pingResults) { + const bestPingResult = selectBestExitNode( + pingResults as ExitNodePingResult[] + ); + if (!bestPingResult) { + logger.warn("No suitable exit node found based on ping results"); + } + exitNodeId = bestPingResult?.exitNodeId; + } + + if ( + client.pubKey !== publicKey || + client.archived || + client.exitNodeId !== exitNodeId + ) { logger.info( "[handleOlmRegisterMessage] Public key mismatch. Updating public key and clearing session info...", { orgId: client.orgId, clientId: client.clientId } @@ -294,7 +311,8 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { .update(clients) .set({ pubKey: publicKey, - archived: false + archived: false, + exitNodeId: exitNodeId // this can be undefined if no exit node was selected, which is fine just means we cant talk to the node or connect to it }) .where(eq(clients.clientId, client.clientId)); diff --git a/server/routers/ws/messageHandlers.ts b/server/routers/ws/messageHandlers.ts index 496002142..6b8a14127 100644 --- a/server/routers/ws/messageHandlers.ts +++ b/server/routers/ws/messageHandlers.ts @@ -5,7 +5,7 @@ import { handleNewtGetConfigMessage, handleDockerStatusMessage, handleDockerContainersMessage, - handleNewtPingRequestMessage, + handleNewtExitNodesRequestMessage, handleApplyBlueprintMessage, handleNewtPingMessage, startNewtOfflineChecker, @@ -45,7 +45,7 @@ export const messageHandlers: Record = { "newt/receive-bandwidth": handleReceiveBandwidthMessage, "newt/socket/status": handleDockerStatusMessage, "newt/socket/containers": handleDockerContainersMessage, - "newt/ping/request": handleNewtPingRequestMessage, + "newt/ping/request": handleNewtExitNodesRequestMessage, "newt/blueprint/apply": handleApplyBlueprintMessage, "newt/healthcheck/status": handleHealthcheckStatusMessage, "ws/round-trip/complete": handleRoundTripMessage