diff --git a/server/db/pg/schema/schema.ts b/server/db/pg/schema/schema.ts index 9ae0e76fa..145b35756 100644 --- a/server/db/pg/schema/schema.ts +++ b/server/db/pg/schema/schema.ts @@ -78,7 +78,8 @@ export const orgs = pgTable("orgs", { "settingsEnableGlobalNewtAutoUpdate" ) .notNull() - .default(false) + .default(false), + settingsJitModeLimit: integer("settingsJitModeLimit").notNull().default(250) }); export const orgDomains = pgTable("orgDomains", { diff --git a/server/db/sqlite/schema/schema.ts b/server/db/sqlite/schema/schema.ts index f14e9d8b1..a72e15556 100644 --- a/server/db/sqlite/schema/schema.ts +++ b/server/db/sqlite/schema/schema.ts @@ -78,7 +78,8 @@ export const orgs = sqliteTable("orgs", { { mode: "boolean" } ) .notNull() - .default(false) + .default(false), + settingsJitModeLimit: integer("settingsJitModeLimit").notNull().default(250) }); export const userDomains = sqliteTable("userDomains", { diff --git a/server/lib/rebuildClientAssociations.ts b/server/lib/rebuildClientAssociations.ts index 045ca3927..012c391bf 100644 --- a/server/lib/rebuildClientAssociations.ts +++ b/server/lib/rebuildClientAssociations.ts @@ -17,7 +17,8 @@ import { sites, Transaction, userOrgRoles, - userSiteResources + userSiteResources, + orgs } from "@server/db"; import { and, count, eq, inArray, isNotNull, ne } from "drizzle-orm"; @@ -234,7 +235,8 @@ export async function getClientSiteResourceAccess( .select({ clientId: clients.clientId, pubKey: clients.pubKey, - subnet: clients.subnet + subnet: clients.subnet, + orgId: clients.orgId }) .from(clients) .where( @@ -260,7 +262,8 @@ export async function getClientSiteResourceAccess( .select({ clientId: clients.clientId, pubKey: clients.pubKey, - subnet: clients.subnet + subnet: clients.subnet, + orgId: clients.orgId }) .from(clients) .where( @@ -463,7 +466,8 @@ async function rebuildClientAssociationsFromSiteResourceImpl( .select({ clientId: clients.clientId, pubKey: clients.pubKey, - subnet: clients.subnet + subnet: clients.subnet, + orgId: clients.orgId }) .from(clients) .where( @@ -568,7 +572,8 @@ async function rebuildClientAssociationsFromSiteResourceImpl( .select({ clientId: clients.clientId, pubKey: clients.pubKey, - subnet: clients.subnet + subnet: clients.subnet, + orgId: clients.orgId }) .from(clients) .where( @@ -720,11 +725,13 @@ async function handleMessagesForSiteClients( clientId: number; pubKey: string | null; subnet: string | null; + orgId: string; }[], existingClients: { clientId: number; pubKey: string | null; subnet: string | null; + orgId: string; }[], clientSitesToAdd: number[], clientSitesToRemove: number[], @@ -805,6 +812,7 @@ async function handleMessagesForSiteClients( clientId: number; pubKey: string | null; subnet: string | null; + orgId: string; } >(); @@ -860,6 +868,22 @@ async function handleMessagesForSiteClients( .map((r) => [r.clientId as number, r.olmId]) ); + // Batch-fetch the orgs for all clients we need to process so we don't + // issue a redundant query per client in the loop below + const orgIdsToProcess = Array.from( + new Set( + Array.from(clientsToProcess.values()).map((client) => client.orgId) + ) + ); + const orgRows = + orgIdsToProcess.length > 0 + ? await trx + .select() + .from(orgs) + .where(inArray(orgs.orgId, orgIdsToProcess)) + : []; + const orgByOrgId = new Map(orgRows.map((org) => [org.orgId, org])); + for (const client of clientsToProcess.values()) { // UPDATE THE NEWT if (!client.subnet || !client.pubKey) { @@ -899,7 +923,14 @@ async function handleMessagesForSiteClients( } if (isAdd) { - if (clientSiteCounts[client.clientId] > 250) { + const org = orgByOrgId.get(client.orgId); + + if (!org) { + logger.warn(`Client ${client.clientId} org not found`); + continue; + } + + if (clientSiteCounts[client.clientId] > org.settingsJitModeLimit) { // skip adding the peer if we have more than 250 sites because we are in jit mode anyway logger.info( `rebuildClientAssociations: Client ${client.clientId} has ${clientSiteCounts[client.clientId]} sites so skipping adding peer to newt and olm because it is likely in jit mode` @@ -1570,7 +1601,8 @@ export async function handleMessagingForUpdatedSiteResource( .select({ clientId: clientSiteResourcesAssociationsCache.clientId, pubKey: clients.pubKey, - subnet: clients.subnet + subnet: clients.subnet, + orgId: clients.orgId }) .from(clientSiteResourcesAssociationsCache) .innerJoin( @@ -2391,6 +2423,19 @@ async function handleMessagesForClientSites( .where(eq(clientSitesAssociationsCache.clientId, client.clientId)) .then((rows) => Number(rows[0].count)); + // client.orgId is constant for this call, so fetch the org once + // instead of re-querying it for every site in the loop below + const [org] = await trx + .select() + .from(orgs) + .where(eq(orgs.orgId, client.orgId)) + .limit(1); + + if (!org) { + logger.warn(`Client ${client.clientId} org not found`); + return; + } + for (const siteData of sitesData) { const site = siteData.sites; const exitNode = siteData.exitNodes; @@ -2451,7 +2496,7 @@ async function handleMessagesForClientSites( continue; } - if (totalSitesOnClient > 250) { + if (totalSitesOnClient > org.settingsJitModeLimit) { // skip adding the site if we have more than 250 because we are in jit mode anyway logger.info( `rebuildClientAssociations: Client ${client.clientId} has ${totalSitesOnClient} sites so skipping adding peer to newt and olm because it is likely in jit mode` @@ -3061,7 +3106,8 @@ export async function cleanupSiteAssociations( .select({ clientId: clients.clientId, pubKey: clients.pubKey, - subnet: clients.subnet + subnet: clients.subnet, + orgId: clients.orgId }) .from(clients) .where(inArray(clients.clientId, cachedClientIds)) diff --git a/server/routers/gerbil/updateHolePunch.ts b/server/routers/gerbil/updateHolePunch.ts index 54906e240..e6f6fbd95 100644 --- a/server/routers/gerbil/updateHolePunch.ts +++ b/server/routers/gerbil/updateHolePunch.ts @@ -8,7 +8,8 @@ import { sites, clientSitesAssociationsCache, exitNodes, - ExitNode + ExitNode, + orgs } from "@server/db"; import { db } from "@server/db"; import { eq, and, inArray } from "drizzle-orm"; @@ -112,7 +113,12 @@ export async function updateHolePunch( destinations: destinations }); } catch (error) { - if (!(error instanceof Error && error.message === "Exit node not allowed")) { + if ( + !( + error instanceof Error && + error.message === "Exit node not allowed" + ) + ) { logger.error(error); } return next( @@ -460,7 +466,18 @@ async function handleClientEndpointChange( return; } - if (sitesWithNewtsToUpdate.length > 250) { + const [org] = await db + .select() + .from(orgs) + .where(eq(orgs.orgId, client.orgId)) + .limit(1); + + if (!org) { + logger.warn(`Client ${clientId} org not found`); + return; + } + + if (sitesWithNewtsToUpdate.length > org.settingsJitModeLimit) { logger.warn( `Client ${clientId} has ${sitesWithNewtsToUpdate.length} connected sites so the client will be in jit mode anyway, skipping endpoint updates` ); diff --git a/server/routers/olm/getOlmToken.ts b/server/routers/olm/getOlmToken.ts index f7fdb81a8..c7c0c127b 100644 --- a/server/routers/olm/getOlmToken.ts +++ b/server/routers/olm/getOlmToken.ts @@ -8,7 +8,8 @@ import { ExitNode, exitNodes, sites, - clientSitesAssociationsCache + clientSitesAssociationsCache, + orgs } from "@server/db"; import { olms } from "@server/db"; import HttpCode from "@server/types/HttpCode"; @@ -225,7 +226,23 @@ export async function getOlmToken( ) .where(eq(clientSitesAssociationsCache.clientId, clientIdToUse!)); - if (clientSites.length > 250 && build == "saas") { + const [org] = await db + .select() + .from(orgs) + .where(eq(orgs.orgId, orgIdToUse)) + .limit(1); + + if (!org) { + logger.warn(`Client ${clientIdToUse} org not found`); + return next( + createHttpError( + HttpCode.INTERNAL_SERVER_ERROR, + "Client's org not found" + ) + ); + } + + if (clientSites.length > org.settingsJitModeLimit && build == "saas") { // set all of the cache rows isJitMode to true await db .update(clientSitesAssociationsCache) diff --git a/server/routers/olm/handleOlmRegisterMessage.ts b/server/routers/olm/handleOlmRegisterMessage.ts index 988e68afd..b65a78a70 100644 --- a/server/routers/olm/handleOlmRegisterMessage.ts +++ b/server/routers/olm/handleOlmRegisterMessage.ts @@ -277,7 +277,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { ); let jitMode = false; - if (sitesCount > 250 && build == "saas") { + if (sitesCount > org.settingsJitModeLimit && build == "saas") { // THIS IS THE MAX ON THE BUSINESS TIER // we have too many sites // If we have too many sites we need to drop into fully JIT mode by not sending any of the sites diff --git a/server/routers/olm/sync.ts b/server/routers/olm/sync.ts index 19d759769..5a681fef7 100644 --- a/server/routers/olm/sync.ts +++ b/server/routers/olm/sync.ts @@ -5,7 +5,8 @@ import { exitNodes, Olm, sites, - clientSitesAssociationsCache + clientSitesAssociationsCache, + orgs } from "@server/db"; import { buildSiteConfigurationForOlmClient } from "./buildConfiguration"; import { sendToClient } from "#dynamic/routers/ws"; @@ -36,8 +37,19 @@ export async function sendOlmSyncMessage(olm: Olm, client: Client) { { orgId: client.orgId } ); + const [org] = await db + .select() + .from(orgs) + .where(eq(orgs.orgId, client.orgId)) + .limit(1); + + if (!org) { + logger.warn(`Client ${client.clientId} org not found`); + return; + } + let jitMode = false; - if (sitesCount > 250 && build == "saas") { + if (sitesCount > org.settingsJitModeLimit && build == "saas") { // THIS IS THE MAX ON THE BUSINESS TIER // we have too many sites // If we have too many sites we need to drop into fully JIT mode by not sending any of the sites