From d5d99a480465650e07bfdde64a5d65ca61567004 Mon Sep 17 00:00:00 2001 From: Owen Date: Mon, 29 Jun 2026 14:59:05 -0400 Subject: [PATCH] Add org rebuild rate limit --- server/lib/orgRebuildCounter.ts | 24 ++++ server/lib/rebuildClientAssociations.ts | 20 +++- server/private/lib/orgRebuildCounter.ts | 105 ++++++++++++++++++ server/private/routers/user/addUserRole.ts | 14 ++- .../private/routers/user/setUserOrgRoles.ts | 14 ++- server/routers/client/createClient.ts | 14 ++- server/routers/client/createUserClient.ts | 14 ++- server/routers/client/deleteClient.ts | 14 ++- .../rebuildClientAssociationsCacheRoute.ts | 11 +- server/routers/olm/deleteUserOlm.ts | 29 ++++- .../siteResource/addClientToSiteResource.ts | 14 ++- .../siteResource/addRoleToSiteResource.ts | 14 ++- .../siteResource/addUserToSiteResource.ts | 14 ++- .../batchAddClientToSiteResources.ts | 14 ++- .../siteResource/createSiteResource.ts | 14 ++- .../removeClientFromSiteResource.ts | 13 ++- .../removeRoleFromSiteResource.ts | 14 ++- .../removeUserFromSiteResource.ts | 14 ++- .../siteResource/setSiteResourceClients.ts | 14 ++- .../siteResource/setSiteResourceRoles.ts | 11 +- .../siteResource/setSiteResourceUsers.ts | 14 ++- .../siteResource/updateSiteResource.ts | 10 ++ server/routers/user/addUserRoleLegacy.ts | 14 ++- 23 files changed, 413 insertions(+), 20 deletions(-) create mode 100644 server/lib/orgRebuildCounter.ts create mode 100644 server/private/lib/orgRebuildCounter.ts diff --git a/server/lib/orgRebuildCounter.ts b/server/lib/orgRebuildCounter.ts new file mode 100644 index 000000000..40edd0aa0 --- /dev/null +++ b/server/lib/orgRebuildCounter.ts @@ -0,0 +1,24 @@ +export const ORG_REBUILD_CONCURRENCY_LIMIT = 10; + +const orgActiveRebuilds = new Map(); + +export async function incrementOrgRebuildCount(orgId: string): Promise { + orgActiveRebuilds.set(orgId, (orgActiveRebuilds.get(orgId) ?? 0) + 1); +} + +export async function decrementOrgRebuildCount(orgId: string): Promise { + const current = orgActiveRebuilds.get(orgId) ?? 0; + if (current <= 1) { + orgActiveRebuilds.delete(orgId); + } else { + orgActiveRebuilds.set(orgId, current - 1); + } +} + +export async function getOrgActiveRebuildCount(orgId: string): Promise { + return orgActiveRebuilds.get(orgId) ?? 0; +} + +export async function checkOrgRebuildRateLimit(orgId: string): Promise { + return (orgActiveRebuilds.get(orgId) ?? 0) >= ORG_REBUILD_CONCURRENCY_LIMIT; +} diff --git a/server/lib/rebuildClientAssociations.ts b/server/lib/rebuildClientAssociations.ts index 1f675e81e..150001662 100644 --- a/server/lib/rebuildClientAssociations.ts +++ b/server/lib/rebuildClientAssociations.ts @@ -45,11 +45,23 @@ import { } from "@server/routers/client/targets"; import { lockManager } from "#dynamic/lib/lock"; import { rebuildQueue } from "#dynamic/lib/rebuildQueue"; +import { + checkOrgRebuildRateLimit, + decrementOrgRebuildCount, + incrementOrgRebuildCount, + ORG_REBUILD_CONCURRENCY_LIMIT +} from "#dynamic/lib/orgRebuildCounter"; + +export { ORG_REBUILD_CONCURRENCY_LIMIT }; // TTL for rebuild-association locks. These functions can fan out into many // peer/proxy updates, so give them a generous window. const REBUILD_ASSOCIATIONS_LOCK_TTL_MS = 120000; +export async function isOrgRebuildRateLimited(orgId: string): Promise { + return checkOrgRebuildRateLimit(orgId); +} + const REBUILD_IDLE_POLL_INTERVAL_MS = 300; const REBUILD_IDLE_DEFAULT_TIMEOUT_MS = 130_000; // slightly longer than lock TTL const REBUILD_IDLE_HANDLER_TIMEOUT_MS = 5_000; @@ -271,6 +283,7 @@ export async function getClientSiteResourceAccess( export async function rebuildClientAssociationsFromSiteResource( siteResource: SiteResource ) { + await incrementOrgRebuildCount(siteResource.orgId); try { return await lockManager.withLock( `rebuild-client-associations:site-resource:${siteResource.siteResourceId}`, @@ -292,6 +305,8 @@ export async function rebuildClientAssociationsFromSiteResource( return { mergedAllClients: [] }; } throw err; + } finally { + await decrementOrgRebuildCount(siteResource.orgId); } } @@ -1638,8 +1653,9 @@ export async function handleMessagingForUpdatedSiteResource( export async function rebuildClientAssociationsFromClient( client: Client ): Promise { - const trx = primaryDb; + await incrementOrgRebuildCount(client.orgId); try { + const trx = primaryDb; return await lockManager.withLock( `rebuild-client-associations:client:${client.clientId}`, () => rebuildClientAssociationsFromClientImpl(client, trx), @@ -1660,6 +1676,8 @@ export async function rebuildClientAssociationsFromClient( return; } throw err; + } finally { + await decrementOrgRebuildCount(client.orgId); } } diff --git a/server/private/lib/orgRebuildCounter.ts b/server/private/lib/orgRebuildCounter.ts new file mode 100644 index 000000000..58d954f66 --- /dev/null +++ b/server/private/lib/orgRebuildCounter.ts @@ -0,0 +1,105 @@ +/* + * This file is part of a proprietary work. + * + * Copyright (c) 2025-2026 Fossorial, Inc. + * All rights reserved. + * + * This file is licensed under the Fossorial Commercial License. + * You may not use this file except in compliance with the License. + * Unauthorized use, copying, modification, or distribution is strictly prohibited. + * + * This file is not licensed under the AGPLv3. + */ + +import { redis } from "#private/lib/redis"; +import logger from "@server/logger"; + +export const ORG_REBUILD_CONCURRENCY_LIMIT = 5; + +// Safety-net TTL: slightly longer than the rebuild lock TTL (120 s). If a +// server process dies while holding a rebuild, this ensures the counter key +// eventually expires rather than staying inflated forever. +const ORG_REBUILD_COUNT_TTL_MS = 180000; +const KEY_PREFIX = "rebuild-org-count:"; + +// In-memory fallback used when Redis is unavailable. +const localFallback = new Map(); + +function isRedisReady(): boolean { + return !!(redis && redis.status === "ready"); +} + +export async function incrementOrgRebuildCount(orgId: string): Promise { + if (!isRedisReady()) { + localFallback.set(orgId, (localFallback.get(orgId) ?? 0) + 1); + return; + } + try { + const key = `${KEY_PREFIX}${orgId}`; + await redis!.incr(key); + // Always refresh the TTL so the key doesn't expire while rebuilds are + // still in progress. The TTL is purely a crash-recovery safety net. + await redis!.pexpire(key, ORG_REBUILD_COUNT_TTL_MS); + } catch (err) { + logger.warn( + `orgRebuildCounter: Redis increment failed for org ${orgId}, falling back to local:`, + err + ); + localFallback.set(orgId, (localFallback.get(orgId) ?? 0) + 1); + } +} + +export async function decrementOrgRebuildCount(orgId: string): Promise { + if (!isRedisReady()) { + const current = localFallback.get(orgId) ?? 0; + if (current <= 1) { + localFallback.delete(orgId); + } else { + localFallback.set(orgId, current - 1); + } + return; + } + try { + const key = `${KEY_PREFIX}${orgId}`; + const count = await redis!.decr(key); + if (count <= 0) { + await redis!.del(key); + } + } catch (err) { + logger.warn( + `orgRebuildCounter: Redis decrement failed for org ${orgId}, falling back to local:`, + err + ); + const current = localFallback.get(orgId) ?? 0; + if (current <= 1) { + localFallback.delete(orgId); + } else { + localFallback.set(orgId, current - 1); + } + } +} + +export async function getOrgActiveRebuildCount(orgId: string): Promise { + if (!isRedisReady()) { + return localFallback.get(orgId) ?? 0; + } + try { + const key = `${KEY_PREFIX}${orgId}`; + const val = await redis!.get(key); + return val ? parseInt(val, 10) : 0; + } catch (err) { + logger.warn( + `orgRebuildCounter: Redis get failed for org ${orgId}, falling back to local:`, + err + ); + return localFallback.get(orgId) ?? 0; + } +} + +export async function checkOrgRebuildRateLimit( + orgId: string +): Promise { + return ( + (await getOrgActiveRebuildCount(orgId)) >= ORG_REBUILD_CONCURRENCY_LIMIT + ); +} diff --git a/server/private/routers/user/addUserRole.ts b/server/private/routers/user/addUserRole.ts index ce5a6dd50..2065984a3 100644 --- a/server/private/routers/user/addUserRole.ts +++ b/server/private/routers/user/addUserRole.ts @@ -23,7 +23,10 @@ import createHttpError from "http-errors"; import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromClient, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; const addUserRoleParamsSchema = z.strictObject({ userId: z.string(), @@ -128,6 +131,15 @@ export async function addUserRole( ); } + if (await isOrgRebuildRateLimited(role.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + let newUserRole: { userId: string; orgId: string; diff --git a/server/private/routers/user/setUserOrgRoles.ts b/server/private/routers/user/setUserOrgRoles.ts index ef6bc1b4f..b583233d1 100644 --- a/server/private/routers/user/setUserOrgRoles.ts +++ b/server/private/routers/user/setUserOrgRoles.ts @@ -21,7 +21,10 @@ import HttpCode from "@server/types/HttpCode"; import createHttpError from "http-errors"; import logger from "@server/logger"; import { fromError } from "zod-validation-error"; -import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromClient, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; const setUserOrgRolesParamsSchema = z.strictObject({ orgId: z.string(), @@ -87,6 +90,15 @@ export async function setUserOrgRoles( ); } + if (await isOrgRebuildRateLimited(orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + const orgRoles = await db .select({ roleId: roles.roleId, isAdmin: roles.isAdmin }) .from(roles) diff --git a/server/routers/client/createClient.ts b/server/routers/client/createClient.ts index bef47245d..57604f0f0 100644 --- a/server/routers/client/createClient.ts +++ b/server/routers/client/createClient.ts @@ -24,7 +24,10 @@ import { isIpInCidr } from "@server/lib/ip"; import { listExitNodes } from "#dynamic/lib/exitNodes"; import { generateId } from "@server/auth/sessions/app"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromClient, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; import { getUniqueClientName } from "@server/db/names"; import { build } from "@server/build"; @@ -154,6 +157,15 @@ export async function createClient( ); } + if (await isOrgRebuildRateLimited(orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + const updatedSubnet = `${subnet}/${org.subnet.split("/")[1]}`; // we want the block size of the whole org // make sure the subnet is unique diff --git a/server/routers/client/createUserClient.ts b/server/routers/client/createUserClient.ts index 3c7d85018..70027e4a5 100644 --- a/server/routers/client/createUserClient.ts +++ b/server/routers/client/createUserClient.ts @@ -21,7 +21,10 @@ import { isValidIP } from "@server/lib/validators"; import { isIpInCidr } from "@server/lib/ip"; import { listExitNodes } from "#dynamic/lib/exitNodes"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromClient, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; import { getUniqueClientName } from "@server/db/names"; const paramsSchema = z @@ -146,6 +149,15 @@ export async function createUserClient( ); } + if (await isOrgRebuildRateLimited(orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + const updatedSubnet = `${subnet}/${org.subnet.split("/")[1]}`; // we want the block size of the whole org // make sure the subnet is unique diff --git a/server/routers/client/deleteClient.ts b/server/routers/client/deleteClient.ts index 62765c2c1..e6acead2e 100644 --- a/server/routers/client/deleteClient.ts +++ b/server/routers/client/deleteClient.ts @@ -9,7 +9,10 @@ import createHttpError from "http-errors"; import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromClient, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; import { sendTerminateClient } from "./terminate"; import { OlmErrorCodes } from "../olm/error"; @@ -76,6 +79,15 @@ export async function deleteClient( ); } + if (await isOrgRebuildRateLimited(client.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + // Only allow deletion of machine clients (clients without userId) if (client.userId) { return next( diff --git a/server/routers/client/rebuildClientAssociationsCacheRoute.ts b/server/routers/client/rebuildClientAssociationsCacheRoute.ts index 86cb5c485..ca149fa97 100644 --- a/server/routers/client/rebuildClientAssociationsCacheRoute.ts +++ b/server/routers/client/rebuildClientAssociationsCacheRoute.ts @@ -9,7 +9,7 @@ import createHttpError from "http-errors"; import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations"; +import { rebuildClientAssociationsFromClient, isOrgRebuildRateLimited } from "@server/lib/rebuildClientAssociations"; const paramsSchema = z.strictObject({ clientId: z.string().transform(Number).pipe(z.int().positive()) @@ -60,6 +60,15 @@ export async function rebuildClientAssociationsCacheRoute( ); } + if (await isOrgRebuildRateLimited(client.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + rebuildClientAssociationsFromClient(client).catch((e) => { logger.error( `Failed to rebuild client associations for client ${clientId}: ${e}` diff --git a/server/routers/olm/deleteUserOlm.ts b/server/routers/olm/deleteUserOlm.ts index addef32d9..d2d0821e8 100644 --- a/server/routers/olm/deleteUserOlm.ts +++ b/server/routers/olm/deleteUserOlm.ts @@ -9,7 +9,10 @@ import { z } from "zod"; import { fromError } from "zod-validation-error"; import logger from "@server/logger"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromClient, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; import { sendTerminateClient } from "../client/terminate"; import { OlmErrorCodes } from "./error"; @@ -64,6 +67,30 @@ export async function deleteUserOlm( const { olmId } = parsedParams.data; + // get the client first + const [client] = await db + .select() + .from(clients) + .where(eq(clients.olmId, olmId)); + + if (!client) { + return next( + createHttpError( + HttpCode.NOT_FOUND, + `No client found for olmId ${olmId}` + ) + ); + } + + if (await isOrgRebuildRateLimited(client.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + let deletedClient: Client | undefined; // Delete associated clients and the OLM in a transaction await db.transaction(async (trx) => { diff --git a/server/routers/siteResource/addClientToSiteResource.ts b/server/routers/siteResource/addClientToSiteResource.ts index 4a6dd141e..5e81cf8f7 100644 --- a/server/routers/siteResource/addClientToSiteResource.ts +++ b/server/routers/siteResource/addClientToSiteResource.ts @@ -8,7 +8,10 @@ import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { eq, and } from "drizzle-orm"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromSiteResource, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; const addClientToSiteResourceBodySchema = z .object({ @@ -128,6 +131,15 @@ export async function addClientToSiteResource( ); } + if (await isOrgRebuildRateLimited(siteResource.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + // Check if client already exists in site resource const existingEntry = await db .select() diff --git a/server/routers/siteResource/addRoleToSiteResource.ts b/server/routers/siteResource/addRoleToSiteResource.ts index 05186c351..d3dd164cc 100644 --- a/server/routers/siteResource/addRoleToSiteResource.ts +++ b/server/routers/siteResource/addRoleToSiteResource.ts @@ -9,7 +9,10 @@ import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { eq, and } from "drizzle-orm"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromSiteResource, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; const addRoleToSiteResourceBodySchema = z .object({ @@ -104,6 +107,15 @@ export async function addRoleToSiteResource( ); } + if (await isOrgRebuildRateLimited(siteResource.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + // verify the role exists and belongs to the same org const [role] = await db .select() diff --git a/server/routers/siteResource/addUserToSiteResource.ts b/server/routers/siteResource/addUserToSiteResource.ts index c35357993..0a93868ae 100644 --- a/server/routers/siteResource/addUserToSiteResource.ts +++ b/server/routers/siteResource/addUserToSiteResource.ts @@ -9,7 +9,10 @@ import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { eq, and } from "drizzle-orm"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromSiteResource, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; const addUserToSiteResourceBodySchema = z .object({ @@ -104,6 +107,15 @@ export async function addUserToSiteResource( ); } + if (await isOrgRebuildRateLimited(siteResource.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + // Check if user already exists in site resource const existingEntry = await db .select() diff --git a/server/routers/siteResource/batchAddClientToSiteResources.ts b/server/routers/siteResource/batchAddClientToSiteResources.ts index 1ebb3359d..3eb282310 100644 --- a/server/routers/siteResource/batchAddClientToSiteResources.ts +++ b/server/routers/siteResource/batchAddClientToSiteResources.ts @@ -15,7 +15,10 @@ import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { eq, and, inArray } from "drizzle-orm"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromClient, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; const batchAddClientToSiteResourcesParamsSchema = z .object({ @@ -186,6 +189,15 @@ export async function batchAddClientToSiteResources( ); } + if (await isOrgRebuildRateLimited(client.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + if (client.userId !== null) { return next( createHttpError( diff --git a/server/routers/siteResource/createSiteResource.ts b/server/routers/siteResource/createSiteResource.ts index d0d018f84..4b55921a2 100644 --- a/server/routers/siteResource/createSiteResource.ts +++ b/server/routers/siteResource/createSiteResource.ts @@ -21,7 +21,10 @@ import { } from "@server/lib/ip"; import { isLicensedOrSubscribed } from "#dynamic/lib/isLicencedOrSubscribed"; import { TierFeature, tierMatrix } from "@server/lib/billing/tierMatrix"; -import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromSiteResource, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; import response from "@server/lib/response"; import logger from "@server/logger"; import { OpenAPITags, registry } from "@server/openApi"; @@ -339,6 +342,15 @@ export async function createSiteResource( ); } + if (await isOrgRebuildRateLimited(org.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + // Only check if destination is an IP address const isIp = z .union([z.ipv4(), z.ipv6()]) diff --git a/server/routers/siteResource/removeClientFromSiteResource.ts b/server/routers/siteResource/removeClientFromSiteResource.ts index c53e214cd..e955dd5c6 100644 --- a/server/routers/siteResource/removeClientFromSiteResource.ts +++ b/server/routers/siteResource/removeClientFromSiteResource.ts @@ -8,7 +8,10 @@ import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { eq, and } from "drizzle-orm"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromSiteResource, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; const removeClientFromSiteResourceBodySchema = z .object({ @@ -106,6 +109,14 @@ export async function removeClientFromSiteResource( ); } + if (await isOrgRebuildRateLimited(siteResource.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } // Check if client exists and has a userId const [client] = await db .select() diff --git a/server/routers/siteResource/removeRoleFromSiteResource.ts b/server/routers/siteResource/removeRoleFromSiteResource.ts index 6dd978e24..509088597 100644 --- a/server/routers/siteResource/removeRoleFromSiteResource.ts +++ b/server/routers/siteResource/removeRoleFromSiteResource.ts @@ -9,7 +9,10 @@ import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { eq, and } from "drizzle-orm"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromSiteResource, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; const removeRoleFromSiteResourceBodySchema = z .object({ @@ -106,6 +109,15 @@ export async function removeRoleFromSiteResource( ); } + if (await isOrgRebuildRateLimited(siteResource.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + // Check if the role is an admin role const [roleToCheck] = await db .select() diff --git a/server/routers/siteResource/removeUserFromSiteResource.ts b/server/routers/siteResource/removeUserFromSiteResource.ts index 67e6ac960..3d8e1fd97 100644 --- a/server/routers/siteResource/removeUserFromSiteResource.ts +++ b/server/routers/siteResource/removeUserFromSiteResource.ts @@ -9,7 +9,10 @@ import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { eq, and } from "drizzle-orm"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromSiteResource, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; const removeUserFromSiteResourceBodySchema = z .object({ @@ -106,6 +109,15 @@ export async function removeUserFromSiteResource( ); } + if (await isOrgRebuildRateLimited(siteResource.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + // Check if user exists in site resource const existingEntry = await db .select() diff --git a/server/routers/siteResource/setSiteResourceClients.ts b/server/routers/siteResource/setSiteResourceClients.ts index a4bc5b69e..59dff6ef9 100644 --- a/server/routers/siteResource/setSiteResourceClients.ts +++ b/server/routers/siteResource/setSiteResourceClients.ts @@ -8,7 +8,10 @@ import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { eq, inArray } from "drizzle-orm"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromSiteResource, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; const setSiteResourceClientsBodySchema = z .object({ @@ -107,6 +110,15 @@ export async function setSiteResourceClients( ); } + if (await isOrgRebuildRateLimited(siteResource.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + // Check if any clients have a userId (associated with a user) if (clientIds.length > 0) { const clientsWithUsers = await db diff --git a/server/routers/siteResource/setSiteResourceRoles.ts b/server/routers/siteResource/setSiteResourceRoles.ts index cad6da53b..613f0aa63 100644 --- a/server/routers/siteResource/setSiteResourceRoles.ts +++ b/server/routers/siteResource/setSiteResourceRoles.ts @@ -9,7 +9,7 @@ import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { eq, and, ne, inArray } from "drizzle-orm"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations"; +import { rebuildClientAssociationsFromSiteResource, isOrgRebuildRateLimited } from "@server/lib/rebuildClientAssociations"; const setSiteResourceRolesBodySchema = z .object({ @@ -167,6 +167,15 @@ export async function setSiteResourceRoles( } }); + if (await isOrgRebuildRateLimited(siteResource.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + rebuildClientAssociationsFromSiteResource(siteResource).catch((e) => { logger.error( `Failed to rebuild client associations for site resource ${siteResourceId}. Error: ${e}` diff --git a/server/routers/siteResource/setSiteResourceUsers.ts b/server/routers/siteResource/setSiteResourceUsers.ts index cde5b4e66..4a423bf06 100644 --- a/server/routers/siteResource/setSiteResourceUsers.ts +++ b/server/routers/siteResource/setSiteResourceUsers.ts @@ -9,7 +9,10 @@ import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { eq } from "drizzle-orm"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromSiteResource, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; import { error } from "node:console"; const setSiteResourceUsersBodySchema = z @@ -109,6 +112,15 @@ export async function setSiteResourceUsers( ); } + if (await isOrgRebuildRateLimited(siteResource.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + await db.transaction(async (trx) => { await trx .delete(userSiteResources) diff --git a/server/routers/siteResource/updateSiteResource.ts b/server/routers/siteResource/updateSiteResource.ts index 434163f6f..22159f1c6 100644 --- a/server/routers/siteResource/updateSiteResource.ts +++ b/server/routers/siteResource/updateSiteResource.ts @@ -19,6 +19,7 @@ import { OpenAPITags, registry } from "@server/openApi"; import { isIpInCidr, portRangeStringSchema } from "@server/lib/ip"; import { handleMessagingForUpdatedSiteResource, + isOrgRebuildRateLimited, rebuildClientAssociationsFromSiteResource, waitForSiteResourceRebuildIdle } from "@server/lib/rebuildClientAssociations"; @@ -345,6 +346,15 @@ export async function updateSiteResource( ); } + if (await isOrgRebuildRateLimited(org.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + // Verify the site exists and belongs to the org const sitesToAssign = await db .select() diff --git a/server/routers/user/addUserRoleLegacy.ts b/server/routers/user/addUserRoleLegacy.ts index b3ff55a06..b7df41bf3 100644 --- a/server/routers/user/addUserRoleLegacy.ts +++ b/server/routers/user/addUserRoleLegacy.ts @@ -10,7 +10,10 @@ import createHttpError from "http-errors"; import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { OpenAPITags, registry } from "@server/openApi"; -import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations"; +import { + rebuildClientAssociationsFromClient, + isOrgRebuildRateLimited +} from "@server/lib/rebuildClientAssociations"; /** Legacy path param order: /role/:roleId/add/:userId */ const addUserRoleLegacyParamsSchema = z.strictObject({ @@ -127,6 +130,15 @@ export async function addUserRoleLegacy( ); } + if (await isOrgRebuildRateLimited(role.orgId)) { + return next( + createHttpError( + HttpCode.TOO_MANY_REQUESTS, + "Too many concurrent rebuild operations for this organization. Please retry after a moment." + ) + ); + } + let orgClientsToRebuild: Client[] = []; await db.transaction(async (trx) => {