diff --git a/server/lib/exitNodes/getUniqueSubnetForExitNode.ts b/server/lib/exitNodes/getUniqueSubnetForExitNode.ts index 728049eb1..1c5f0f660 100644 --- a/server/lib/exitNodes/getUniqueSubnetForExitNode.ts +++ b/server/lib/exitNodes/getUniqueSubnetForExitNode.ts @@ -1,41 +1,62 @@ -import { db, ExitNode, Transaction, sites, clients } from "@server/db"; +import { ExitNode, primaryDb, sites, clients } from "@server/db"; import { eq } from "drizzle-orm"; import config from "@server/lib/config"; import { findNextAvailableCidr } from "@server/lib/ip"; import { lockManager } from "#dynamic/lib/lock"; +/** + * Picks the next free site/client subnet on the given exit node. This acquires + * a per-exit-node lock that the caller MUST release (via the returned + * `release`) only after the chosen subnet has been persisted to the database + * (or the operation has failed) — otherwise a concurrent caller on another + * node could pick the same subnet. If no subnet is available the lock is + * released before returning and `value` is null. + */ export async function getUniqueSubnetForExitNode( - exitNode: ExitNode, - trx: Transaction | typeof db = db -): Promise { + exitNode: ExitNode +): Promise<{ value: string | null; release: () => Promise }> { const lockKey = `subnet-allocation:${exitNode.exitNodeId}`; + const acquired = await lockManager.acquireLockWithRetry(lockKey, 6000); + if (!acquired) { + throw new Error(`Failed to acquire lock: ${lockKey}`); + } + const release = () => lockManager.releaseLock(lockKey, acquired); - return await lockManager.withLock( - lockKey, - async () => { - const [sitesQuery, clientsQuery] = await Promise.all([ - trx - .select({ subnet: sites.exitNodeSubnet }) - .from(sites) - .where(eq(sites.exitNodeId, exitNode.exitNodeId)), - trx - .select({ subnet: clients.exitNodeSubnet }) - .from(clients) - .where(eq(clients.exitNodeId, exitNode.exitNodeId)) - ]); + try { + // read from the primary so we see subnets committed by the previous lock holder + const [sitesQuery, clientsQuery] = await Promise.all([ + primaryDb + .select({ subnet: sites.exitNodeSubnet }) + .from(sites) + .where(eq(sites.exitNodeId, exitNode.exitNodeId)), + primaryDb + .select({ subnet: clients.exitNodeSubnet }) + .from(clients) + .where(eq(clients.exitNodeId, exitNode.exitNodeId)) + ]); - const blockSize = config.getRawConfig().gerbil.site_block_size; - const subnets = [...sitesQuery, ...clientsQuery] - .map((row) => row.subnet) - .filter( - (subnet): subnet is string => - !!subnet && - /^(\d{1,3}\.){3}\d{1,3}\/\d{1,2}$/.test(subnet) - ); - subnets.push(exitNode.address.replace(/\/\d+$/, `/${blockSize}`)); + const blockSize = config.getRawConfig().gerbil.site_block_size; + const subnets = [...sitesQuery, ...clientsQuery] + .map((row) => row.subnet) + .filter( + (subnet): subnet is string => + !!subnet && /^(\d{1,3}\.){3}\d{1,3}\/\d{1,2}$/.test(subnet) + ); + subnets.push(exitNode.address.replace(/\/\d+$/, `/${blockSize}`)); - return findNextAvailableCidr(subnets, blockSize, exitNode.address); - }, - 5000 // 5 second lock TTL - subnet allocation should be quick - ); + const subnet = findNextAvailableCidr( + subnets, + blockSize, + exitNode.address + ); + if (!subnet) { + await release(); + return { value: null, release: async () => {} }; + } + + return { value: subnet, release }; + } catch (e) { + await release(); + throw e; + } } diff --git a/server/routers/newt/handleNewtRegisterMessage.ts b/server/routers/newt/handleNewtRegisterMessage.ts index 4d90aa0aa..db21a6661 100644 --- a/server/routers/newt/handleNewtRegisterMessage.ts +++ b/server/routers/newt/handleNewtRegisterMessage.ts @@ -112,7 +112,8 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => { return; } - const newSubnet = await getUniqueSubnetForExitNode(exitNode); + const { value: newSubnet, release: releaseSubnetLock } = + await getUniqueSubnetForExitNode(exitNode); if (!newSubnet) { logger.error( @@ -124,15 +125,20 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => { siteSubnet = newSubnet; - await db - .update(sites) - .set({ - pubKey: publicKey, - exitNodeId: exitNodeId, - exitNodeSubnet: newSubnet - }) - .where(eq(sites.siteId, siteId)) - .returning(); + try { + await db + .update(sites) + .set({ + pubKey: publicKey, + exitNodeId: exitNodeId, + exitNodeSubnet: newSubnet + }) + .where(eq(sites.siteId, siteId)) + .returning(); + } finally { + // only release once the subnet is persisted so no other node can pick it + await releaseSubnetLock(); + } } else { await db .update(sites) diff --git a/server/routers/olm/handleOlmRegisterMessage.ts b/server/routers/olm/handleOlmRegisterMessage.ts index f583729ef..73078a4a6 100644 --- a/server/routers/olm/handleOlmRegisterMessage.ts +++ b/server/routers/olm/handleOlmRegisterMessage.ts @@ -314,6 +314,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { } let clientSubnet = client.exitNodeSubnet; + let releaseSubnetLock: (() => Promise) | null = null; if ( exitNodeId && (client.exitNodeId !== exitNodeId || !client.exitNodeSubnet) @@ -341,7 +342,8 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { // TODO: IF WE DO NOT HAVE AN INFERENCE RESOURCE DO WE NEED TO BE HOLDING A SUBNET ON THE CLIENT? - const newSubnet = await getUniqueSubnetForExitNode(exitNode); + const { value: newSubnet, release } = + await getUniqueSubnetForExitNode(exitNode); if (!newSubnet) { logger.error( @@ -353,6 +355,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { } clientSubnet = newSubnet; + releaseSubnetLock = release; } if ( @@ -365,16 +368,21 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { "[handleOlmRegisterMessage] Public key mismatch. Updating public key and clearing session info...", { orgId: client.orgId, clientId: client.clientId } ); - // Update the client's public key - await db - .update(clients) - .set({ - pubKey: publicKey, - 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 - exitNodeSubnet: clientSubnet - }) - .where(eq(clients.clientId, client.clientId)); + try { + // Update the client's public key + await db + .update(clients) + .set({ + pubKey: publicKey, + 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 + exitNodeSubnet: clientSubnet + }) + .where(eq(clients.clientId, client.clientId)); + } finally { + // only release once the subnet is persisted so no other node can pick it + await releaseSubnetLock?.(); + } // set isRelay to false for all of the client's sites to reset the connection metadata await db @@ -397,6 +405,9 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { ); } + // safety net in case a subnet was locked but the update above was skipped (release is idempotent) + await releaseSubnetLock?.(); + if (client.pubKey && client.pubKey !== publicKey && client.exitNodeId) { // test the old client to see if its different then remove logger.info("Public key mismatch. Deleting old peer...");