Fix TOCTOU Race Condition in registerDebouncedRun

Author: ericallam

internal-packages/run-engine/src/engine/systems/debounceSystem.ts

@@ -121,6 +121,22 @@ end
 return { 0, value }
       `,
     });
+
+    // Atomically sets runId only if current value equals expected pending claim.
+    // This prevents the TOCTOU race condition where between GET (check claim) and SET (register),
+    // another server could claim and register a different run, which would get overwritten.
+    // Returns 1 if set succeeded, 0 if claim mismatch (lost the claim).
+    this.redis.defineCommand("registerIfClaimOwned", {
+      numberOfKeys: 1,
+      lua: `
+local value = redis.call('GET', KEYS[1])
+if value == ARGV[1] then
+  redis.call('SET', KEYS[1], ARGV[2], 'PX', ARGV[3])
+  return 1
+end
+return 0
+      `,
+    });
   }
 
   /**
@@ -593,32 +609,40 @@ return { 0, value }
       async (span) => {
         const redisKey = this.getDebounceRedisKey(environmentId, taskIdentifier, debounceKey);
 
+        // Calculate TTL: delay until + buffer
+        const ttlMs = Math.max(
+          delayUntil.getTime() - Date.now() + 60_000, // Add 1 minute buffer
+          60_000
+        );
+
         if (claimId) {
-          // Verify we still own the pending claim before overwriting
-          const currentValue = await this.redis.get(redisKey);
-          if (currentValue !== `pending:${claimId}`) {
+          // Use atomic Lua script to verify claim and set runId in one operation.
+          // This prevents the TOCTOU race where another server could claim and register
+          // between our GET check and SET.
+          const result = await this.redis.registerIfClaimOwned(
+            redisKey,
+            `pending:${claimId}`,
+            runId,
+            ttlMs.toString()
+          );
+
+          if (result === 0) {
             // We lost the claim - another server took over or it expired
             this.$.logger.warn("registerDebouncedRun: lost claim, not registering", {
               runId,
               environmentId,
               taskIdentifier,
               debounceKey,
               claimId,
-              currentValue,
             });
             span.setAttribute("claimLost", true);
             return false;
           }
+        } else {
+          // No claim to verify, just set directly
+          await this.redis.set(redisKey, runId, "PX", ttlMs);
         }
 
-        // Calculate TTL: delay until + buffer
-        const ttlMs = Math.max(
-          delayUntil.getTime() - Date.now() + 60_000, // Add 1 minute buffer
-          60_000
-        );
-
-        await this.redis.set(redisKey, runId, "PX", ttlMs);
-
         this.$.logger.debug("registerDebouncedRun: stored debounce key mapping", {
           runId,
           environmentId,
@@ -751,5 +775,22 @@ declare module "@internal/redis" {
       key: string,
       callback?: Callback<[number, string | null]>
     ): Result<[number, string | null], Context>;
+
+    /**
+     * Atomically sets runId only if current value equals expected pending claim.
+     * Prevents TOCTOU race condition between claim verification and registration.
+     * @param key - The Redis key
+     * @param expectedClaim - Expected value "pending:{claimId}"
+     * @param runId - The new value (run ID) to set
+     * @param ttlMs - TTL in milliseconds
+     * @returns 1 if set succeeded, 0 if claim mismatch
+     */
+    registerIfClaimOwned(
+      key: string,
+      expectedClaim: string,
+      runId: string,
+      ttlMs: string,
+      callback?: Callback<number>
+    ): Result<number, Context>;
   }
 }

internal-packages/run-engine/src/engine/tests/debounce.test.ts

@@ -1939,5 +1939,107 @@ describe("RunEngine debounce", () => {
       }
     }
   );
+
+  containerTest(
+    "registerDebouncedRun: atomic claim prevents overwrite when claim is lost",
+    async ({ prisma, redisOptions }) => {
+      // This test verifies the fix for the TOCTOU race condition in registerDebouncedRun.
+      // The race occurs when:
+      // 1. Server A claims debounce key with claimId-A
+      // 2. Server B claims same key with claimId-B (after A's claim expires)
+      // 3. Server B registers runId-B successfully
+      // 4. Server A attempts to register runId-A with stale claimId-A
+      // Without the fix, step 4 would overwrite runId-B. With the fix, it fails atomically.
+
+      const { createRedisClient } = await import("@internal/redis");
+
+      const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
+
+      const engine = new RunEngine({
+        prisma,
+        worker: {
+          redis: redisOptions,
+          workers: 1,
+          tasksPerWorker: 10,
+          pollIntervalMs: 100,
+        },
+        queue: {
+          redis: redisOptions,
+        },
+        runLock: {
+          redis: redisOptions,
+        },
+        machines: {
+          defaultMachine: "small-1x",
+          machines: {
+            "small-1x": {
+              name: "small-1x" as const,
+              cpu: 0.5,
+              memory: 0.5,
+              centsPerMs: 0.0001,
+            },
+          },
+          baseCostInCents: 0.0001,
+        },
+        debounce: {
+          maxDebounceDurationMs: 60_000,
+        },
+        tracer: trace.getTracer("test", "0.0.0"),
+      });
+
+      // Create a separate Redis client to simulate "another server" modifying keys directly
+      const simulatedServerRedis = createRedisClient({
+        ...redisOptions,
+        keyPrefix: `${redisOptions.keyPrefix ?? ""}debounce:`,
+      });
+
+      try {
+        const taskIdentifier = "test-task";
+        const debounceKey = "race-test-key";
+        const environmentId = authenticatedEnvironment.id;
+        const delayUntil = new Date(Date.now() + 60_000);
+
+        await setupBackgroundWorker(engine, authenticatedEnvironment, taskIdentifier);
+
+        // Construct the Redis key (same format as DebounceSystem.getDebounceRedisKey)
+        const redisKey = `${environmentId}:${taskIdentifier}:${debounceKey}`;
+
+        // Step 1: Server A claims the key with claimId-A
+        const claimIdA = "claim-server-A";
+        await simulatedServerRedis.set(redisKey, `pending:${claimIdA}`, "PX", 60_000);
+
+        // Step 2 & 3: Simulate Server B claiming and registering (after A's claim "expires")
+        // In reality, this simulates the race where B's claim overwrites A's pending claim
+        const runIdB = "run_server_B";
+        await simulatedServerRedis.set(redisKey, runIdB, "PX", 60_000);
+
+        // Verify Server B's registration is in place
+        const valueAfterB = await simulatedServerRedis.get(redisKey);
+        expect(valueAfterB).toBe(runIdB);
+
+        // Step 4: Server A attempts to register with its stale claimId-A
+        // This should FAIL because the key no longer contains "pending:claim-server-A"
+        const runIdA = "run_server_A";
+        const registered = await engine.debounceSystem.registerDebouncedRun({
+          runId: runIdA,
+          environmentId,
+          taskIdentifier,
+          debounceKey,
+          delayUntil,
+          claimId: claimIdA, // Stale claim ID
+        });
+
+        // Step 5: Verify Server A's registration failed
+        expect(registered).toBe(false);
+
+        // Step 6: Verify Redis still contains runId-B (not overwritten by Server A)
+        const finalValue = await simulatedServerRedis.get(redisKey);
+        expect(finalValue).toBe(runIdB);
+      } finally {
+        await simulatedServerRedis.quit();
+        await engine.quit();
+      }
+    }
+  );
 });