Fix ID token and session token JWK conflcit in KV store

src/auth.ts

@@ -4,7 +4,7 @@ import type { EmulatorEnv } from './emulator';
 import { useEmulator } from './emulator';
 import type { ErrorInfo } from './errors';
 import { AppErrorCodes, AuthClientErrorCode, FirebaseAppError, FirebaseAuthError } from './errors';
-import type { KeyStorer } from './key-store';
+import { scopedKeyStorer, type KeyStorer } from './key-store';
 import type { FirebaseIdToken, FirebaseTokenVerifier } from './token-verifier';
 import { createIdTokenVerifier, createSessionCookieVerifier } from './token-verifier';
 import type { UserRecord } from './user-record';
@@ -17,8 +17,8 @@ export class BaseAuth {
   private readonly _authApiClient?: AuthApiClient;
 
   constructor(projectId: string, keyStore: KeyStorer, credential?: Credential) {
-    this.idTokenVerifier = createIdTokenVerifier(projectId, keyStore);
-    this.sessionCookieVerifier = createSessionCookieVerifier(projectId, keyStore);
+    this.idTokenVerifier = createIdTokenVerifier(projectId, scopedKeyStorer(keyStore, 'id-token'));
+    this.sessionCookieVerifier = createSessionCookieVerifier(projectId, scopedKeyStorer(keyStore, 'session-cookie'));
 
     if (credential) {
       this._authApiClient = new AuthApiClient(projectId, credential);

src/index.ts

@@ -42,6 +42,10 @@ export class WorkersKVStoreSingle extends WorkersKVStore {
     super(cacheKey, cfKVNamespace);
   }
 
+  public scoped(scope: string): WorkersKVStoreSingle {
+    return WorkersKVStoreSingle.getOrInitialize(`${this.cacheKey}:${scope}`, this.cfKVNamespace);
+  }
+
   static getOrInitialize(cacheKey: string, cfKVNamespace: KVNamespace): WorkersKVStoreSingle {
     if (!WorkersKVStoreSingle.instance) {
       WorkersKVStoreSingle.instance = new Map<string, WorkersKVStoreSingle>();

src/key-store.ts

@@ -1,17 +1,32 @@
 export interface KeyStorer {
   get<ExpectedValue = unknown>(): Promise<ExpectedValue | null>;
   put(value: string, expirationTtl: number): Promise<void>;
+  scoped?(scope: string): KeyStorer;
+}
+
+export function scopedKeyStorer(keyStorer: KeyStorer, scope: string): KeyStorer {
+  if (typeof keyStorer.scoped === 'function') {
+    return keyStorer.scoped(scope);
+  }
+  return keyStorer;
 }
 
 /**
  * Class to get or store fetched public keys from a client certificates URL.
  */
 export class WorkersKVStore implements KeyStorer {
   constructor(
-    private readonly cacheKey: string,
-    private readonly cfKVNamespace: KVNamespace
+    protected readonly cacheKey: string,
+    protected readonly cfKVNamespace: KVNamespace
   ) {}
 
+  /**
+   * Returns a view of this store that caches public keys under a derived KV key.
+   */
+  public scoped(scope: string): WorkersKVStore {
+    return new WorkersKVStore(`${this.cacheKey}:${scope}`, this.cfKVNamespace);
+  }
+
   public async get<ExpectedValue = unknown>(): Promise<ExpectedValue | null> {
     return await this.cfKVNamespace.get<ExpectedValue>(this.cacheKey, 'json');
   }

tests/auth.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from 'vitest';
+import { describe, expect, it, vi } from 'vitest';
 import { ApiSettings } from '../src/api-requests';
 import { BaseAuth } from '../src/auth';
 import { AuthApiClient } from '../src/auth-api-requests';
@@ -24,6 +24,28 @@ const sessionCookieUids = [
   generateRandomString(20),
 ];
 
+describe('BaseAuth key store scoping', () => {
+  it('scopes key store separately for id-token and session-cookie verifiers', () => {
+    const scoped = vi.fn(
+      (_scope: string): KeyStorer => ({
+        get: vi.fn(),
+        put: vi.fn(),
+      })
+    );
+    const keyStorer: KeyStorer = {
+      get: vi.fn(),
+      put: vi.fn(),
+      scoped,
+    };
+
+    new BaseAuth(projectId, keyStorer);
+
+    expect(scoped).toHaveBeenCalledTimes(2);
+    expect(scoped).toHaveBeenCalledWith('id-token');
+    expect(scoped).toHaveBeenCalledWith('session-cookie');
+  });
+});
+
 describe('createSessionCookie()', () => {
   const expiresIn = 24 * 60 * 60 * 1000;
 

tests/key-store.test.ts

@@ -0,0 +1,43 @@
+import { Miniflare } from 'miniflare';
+import { describe, expect, it } from 'vitest';
+import { WorkersKVStoreSingle } from '../src/index';
+import { WorkersKVStore } from '../src/key-store';
+
+const nullScript = 'export default { fetch: () => new Response(null, { status: 404 }) };';
+const mf = new Miniflare({
+  modules: true,
+  script: nullScript,
+  kvNamespaces: ['TEST_NAMESPACE'],
+});
+
+describe('WorkersKVStore.scoped', () => {
+  it('suffixes the KV key with the scope', async () => {
+    const TEST_NAMESPACE = await mf.getKVNamespace('TEST_NAMESPACE');
+    const baseKey = 'scoped-cache-key';
+    const keyStore = new WorkersKVStore(baseKey, TEST_NAMESPACE);
+
+    const idTokenStore = keyStore.scoped('id-token');
+    const sessionCookieStore = keyStore.scoped('session-cookie');
+
+    await idTokenStore.put(JSON.stringify([{ kid: 'id-kid' }]), 3600);
+    await sessionCookieStore.put(JSON.stringify([{ kid: 'session-kid' }]), 3600);
+
+    expect(await TEST_NAMESPACE.get(`${baseKey}:id-token`, 'json')).toEqual([{ kid: 'id-kid' }]);
+    expect(await TEST_NAMESPACE.get(`${baseKey}:session-cookie`, 'json')).toEqual([{ kid: 'session-kid' }]);
+    expect(await TEST_NAMESPACE.get(baseKey)).toBeNull();
+  });
+});
+
+describe('WorkersKVStoreSingle.scoped', () => {
+  it('returns a singleton per scoped cache key', async () => {
+    const TEST_NAMESPACE = await mf.getKVNamespace('TEST_NAMESPACE');
+    const keyStore = WorkersKVStoreSingle.getOrInitialize('base-key', TEST_NAMESPACE);
+
+    const idTokenStore = keyStore.scoped('id-token');
+    const idTokenStoreAgain = keyStore.scoped('id-token');
+    const sessionCookieStore = keyStore.scoped('session-cookie');
+
+    expect(idTokenStore).toBe(idTokenStoreAgain);
+    expect(sessionCookieStore).not.toBe(idTokenStore);
+  });
+});