Prevent overlapping exit node address selections

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<string | null> {
+    exitNode: ExitNode
+): Promise<{ value: string | null; release: () => Promise<void> }> {
     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;
+    }
 }

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)

server/routers/olm/handleOlmRegisterMessage.ts

@@ -314,6 +314,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
     }
 
     let clientSubnet = client.exitNodeSubnet;
+    let releaseSubnetLock: (() => Promise<void>) | 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...");