fix: answer 503 when an impersonation link's use can't be recorded and keep used links as expiring keys #1066

The marker-write test from d28ee361c now expects the 503 path (a rejected promise) instead of false: refusing with "Token already used" sent admins looking for a replay that never happened, and the jti is now remembered in memory only after it was recorded, so a retry works once the store or state dir is back.

app/api/auth/impersonate/route.ts

@@ -8,7 +8,7 @@ import { normalizeJmapServerUrl } from '@/lib/auth/verify-jmap-auth';
 import { setStalwartAuthContextInStore } from '@/lib/stalwart/auth-context';
 import { recordLogin } from '@/lib/telemetry/login-tracker';
 import { ImpersonationJwtError, verifyImpersonationJwt } from '@/lib/impersonation/jwt';
-import { consumeImpersonationJti } from '@/lib/impersonation/replay-store';
+import { consumeImpersonationJti, isImpersonationReplayUnavailable } from '@/lib/impersonation/replay-store';
 import {
   readImpersonationConfig,
   resolveImpersonationServerUrl,
@@ -79,7 +79,15 @@ export async function GET(request: NextRequest) {
     return NextResponse.json({ error: 'Internal server error' }, { status: 500 });
   }
 
-  if (!(await consumeImpersonationJti(claims.jti, claims.exp))) {
+  let firstUse: boolean;
+  try {
+    firstUse = await consumeImpersonationJti(claims.jti, claims.exp);
+  } catch (err) {
+    if (!isImpersonationReplayUnavailable(err)) throw err;
+    // Not "already used": the use couldn't be recorded, and the link may work in a moment.
+    return NextResponse.json({ error: 'Impersonation is temporarily unavailable' }, { status: 503, headers: { 'Retry-After': '10' } });
+  }
+  if (!firstUse) {
     logger.warn('Impersonation JWT replay rejected', { jti: claims.jti });
     return NextResponse.json({ error: 'Token already used' }, { status: 401 });
   }

lib/__tests__/impersonation-replay-store.test.ts

@@ -49,7 +49,10 @@ describe('consumeImpersonationJti', () => {
     const blocked = path.join(stateDir, 'not-a-dir');
     writeFileSync(blocked, '');
     process.env.ADMIN_STATE_DIR = blocked;
-    expect(await consumeImpersonationJti('jti-unrecorded', exp, now)).toBe(false);
+    await expect(consumeImpersonationJti('jti-unrecorded', exp, now)).rejects.toMatchObject({ name: 'ImpersonationReplayUnavailableError' });
+    // Not remembered as used: once the marker can be written, the token works.
+    process.env.ADMIN_STATE_DIR = stateDir;
+    expect(await consumeImpersonationJti('jti-unrecorded', exp, now)).toBe(true);
   });
 
   it('prunes markers once their token can no longer verify', async () => {

lib/__tests__/impersonation-session.test.ts

@@ -5,6 +5,9 @@ import { tmpdir } from 'node:os';
 import path from 'node:path';
 import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
 import { NextRequest } from 'next/server';
+import { setStoreForTests } from '@/lib/shared-store';
+import { resetLocalStoreForTests } from '@/lib/shared-store/local';
+import { switchableSharedStore } from './helpers/shared-store-down';
 
 vi.mock('@/lib/auth/session-secret', () => ({
   getSessionSecret: () => 's'.repeat(64),
@@ -234,3 +237,28 @@ describe('impersonation handoff', () => {
     expect(jar.has('jmap_session')).toBe(false);
   });
 });
+
+describe('impersonation handoff while the shared store is down (#1066)', () => {
+  beforeAll(() => {
+    process.env.BULWARK_JWT_AUTH_SECRET = SECRET;
+    process.env.BULWARK_STALWART_MASTER_USER = 'master@example.org';
+    process.env.BULWARK_STALWART_MASTER_PASSWORD = MASTER_PASSWORD;
+  });
+  afterEach(() => {
+    setStoreForTests(null);
+    resetLocalStoreForTests();
+    vi.unstubAllGlobals();
+  });
+
+  it('answers 503, not "already used", and opens no session', async () => {
+    const store = switchableSharedStore();
+    store.down = true;
+    setStoreForTests(store);
+    jar.clear();
+    calls = [];
+    vi.stubGlobal('fetch', vi.fn(stalwartFetch));
+    const res = await impersonate();
+    expect(res.status).toBe(503);
+    expect(calls).toHaveLength(0);
+  });
+});

lib/__tests__/shared-impersonation-replay.test.ts

@@ -1,5 +1,5 @@
 // @vitest-environment node
-import { mkdtempSync, rmSync } from 'node:fs';
+import { mkdtempSync, readdirSync, rmSync } from 'node:fs';
 import { tmpdir } from 'node:os';
 import path from 'node:path';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
@@ -72,8 +72,18 @@ describe('impersonation replay with a shared store', () => {
     expect(results.filter(Boolean)).toHaveLength(1);
   });
 
-  it('refuses the token while the store is down', async () => {
+  it('refuses the token while the store is down, without remembering it as used', async () => {
     store.down = true;
+    await expect(consumeImpersonationJti('jti-down', exp, now)).rejects.toMatchObject({ name: 'ImpersonationReplayUnavailableError' });
+    store.down = false;
+    expect(await consumeImpersonationJti('jti-down', exp, now)).toBe(true);
     expect(await consumeImpersonationJti('jti-down', exp, now)).toBe(false);
   });
+
+  it('keeps used tokens as expiring keys, not as documents of an ever-growing collection', async () => {
+    const set = vi.spyOn(store.kv, 'set');
+    expect(await consumeImpersonationJti('jti-kv', exp, now)).toBe(true);
+    expect(set).toHaveBeenCalledWith(expect.stringMatching(/^imp-jti:[0-9a-f]{64}$/), '1', { ifAbsent: true, ttlMs: 420_000 });
+    expect(readdirSync(sharedDir)).toEqual([]);
+  });
 });

lib/impersonation/jwt.ts

@@ -158,6 +158,12 @@ const REPLAY_CACHE_MAX = 4096;
 class ReplayCache {
   private entries = new Map<string, number>(); // jti -> exp epoch seconds
 
+  /** Whether jti was recorded and could still verify. */
+  has(jti: string, now: number = Math.floor(Date.now() / 1000)): boolean {
+    this.prune(now);
+    return this.entries.has(jti);
+  }
+
   /** Returns true if jti was not previously seen and has been recorded. */
   consume(jti: string, exp: number, now: number = Math.floor(Date.now() / 1000)): boolean {
     this.prune(now);

lib/impersonation/replay-store.ts

@@ -12,17 +12,20 @@ function markerName(jti: string): string {
   return createHash('sha256').update(jti).digest('hex');
 }
 
-/**
- * The used jtis in the shared store (#1066): one empty document each, gone
- * once its token can no longer verify. Its local layout is today's marker
- * directory, though local mode keeps writing the markers itself (with
- * O_EXCL, which a shared state volume honours).
- */
-export function impersonationJtiCollection() {
-  return getStore().collection(STORE_DIR, {
-    local: { dir: () => getStatePath(STORE_DIR) },
-    ttlMs: IMPERSONATION_REPLAY_WINDOW_SEC * 1000,
-  });
+/** The used jti's key in the shared store (#1066); it expires with the token. */
+const sharedKey = (jti: string) => `imp-jti:${markerName(jti)}`;
+
+/** The use of a token couldn't be recorded: it is refused, and the route answers 503. */
+export class ImpersonationReplayUnavailableError extends Error {
+  constructor(options?: { cause?: unknown }) {
+    super('The impersonation replay store is unavailable', options);
+    this.name = 'ImpersonationReplayUnavailableError';
+  }
+}
+
+/** By name, as every bundle has its own copy of the class. */
+export function isImpersonationReplayUnavailable(error: unknown): boolean {
+  return (error as Error | null)?.name === 'ImpersonationReplayUnavailableError';
 }
 
 async function pruneExpired(dir: string, nowMs: number): Promise<void> {
@@ -38,8 +41,10 @@ async function pruneExpired(dir: string, nowMs: number): Promise<void> {
 }
 
 /**
- * Record an impersonation token's jti as used; false when it was used before
- * or the use can't be recorded.
+ * Record an impersonation token's jti as used; false when it was used before.
+ * Rejects with ImpersonationReplayUnavailableError when the use can't be
+ * recorded: the token is refused then, as a restart or another replica
+ * would otherwise accept it again.
  *
  * The in-process cache alone forgot every jti on restart and knew nothing of
  * other replicas, so a captured handoff link could be redeemed again within
@@ -47,24 +52,32 @@ async function pruneExpired(dir: string, nowMs: number): Promise<void> {
  * created with O_EXCL, which survives restarts and is shared by every
  * replica that mounts the same state volume. Markers are pruned once the
  * token they stand for can no longer verify. With a shared store the marker
- * is a document there, created only if missing, so exactly one replica
- * accepts the token.
+ * is a key there, set only if missing and expiring with the token, so
+ * exactly one replica accepts the token. The in-process cache learns a jti
+ * only once it was recorded.
  */
 export async function consumeImpersonationJti(
   jti: string,
   exp: number,
   nowMs: number = Date.now(),
 ): Promise<boolean> {
-  if (!impersonationReplayCache.consume(jti, exp, Math.floor(nowMs / 1000))) return false;
+  const nowSec = Math.floor(nowMs / 1000);
+  if (impersonationReplayCache.has(jti, nowSec)) return false;
+  if (!(await recordUse(jti, nowMs))) return false;
+  impersonationReplayCache.consume(jti, exp, nowSec);
+  return true;
+}
 
+/** True when this call recorded the jti, false when it was recorded before. */
+async function recordUse(jti: string, nowMs: number): Promise<boolean> {
   if (isSharedMode()) {
     try {
-      return (await impersonationJtiCollection().put(markerName(jti), '', { ifVersion: null })).ok;
+      return await getStore().kv.set(sharedKey(jti), '1', { ifAbsent: true, ttlMs: IMPERSONATION_REPLAY_WINDOW_SEC * 1000 });
     } catch (err) {
       logger.error('Impersonation replay store unavailable, refusing the token', {
         error: err instanceof Error ? err.message : 'Unknown',
       });
-      return false;
+      throw new ImpersonationReplayUnavailableError({ cause: err });
     }
   }
 
@@ -83,6 +96,6 @@ export async function consumeImpersonationJti(
     logger.error('Impersonation replay store unavailable, refusing the token', {
       error: err instanceof Error ? err.message : 'Unknown',
     });
-    return false;
+    throw new ImpersonationReplayUnavailableError({ cause: err });
   }
 }