From 2082671c36339e6fbdb5f928634e040ec5a2f80c Mon Sep 17 00:00:00 2001 From: Bereket Engida Date: Mon, 24 Nov 2025 15:07:26 -0800 Subject: [PATCH] feat(jwt): enhance JWT key rotation with cooldown and manual control features --- packages/better-auth/src/plugins/jwt/index.ts | 11 +- .../src/plugins/jwt/rotation.test.ts | 253 +++++++++++++++++- packages/better-auth/src/plugins/jwt/sign.ts | 5 +- packages/better-auth/src/plugins/jwt/utils.ts | 39 ++- 4 files changed, 287 insertions(+), 21 deletions(-) diff --git a/packages/better-auth/src/plugins/jwt/index.ts b/packages/better-auth/src/plugins/jwt/index.ts index e844f4989a..68852a0386 100644 --- a/packages/better-auth/src/plugins/jwt/index.ts +++ b/packages/better-auth/src/plugins/jwt/index.ts @@ -138,7 +138,7 @@ export const jwt = (options?: JwtOptions | undefined) => { if (!options?.jwks?.disableAutomaticRotation) { - await rotateJwk(ctx, options); + await rotateJwk(ctx, options); // Trigger rotation if needed } let keySets = await adapter.getAllKeys(ctx); @@ -331,14 +331,11 @@ export const jwt = (options?: JwtOptions | undefined) => { .optional(), }, async (ctx) => { - const key = await rotateJwk(ctx, options, ctx.body); - - // Key was just created if it's less than 10 seconds old - const rotated = Date.now() - key.createdAt.getTime() < 10000; + const result = await rotateJwk(ctx, options, ctx.body); return ctx.json({ - rotated, - keyId: key.id, + rotated: result.rotated, + keyId: result.key.id, }); }, ), diff --git a/packages/better-auth/src/plugins/jwt/rotation.test.ts b/packages/better-auth/src/plugins/jwt/rotation.test.ts index 842c15eaa0..ded9094a24 100644 --- a/packages/better-auth/src/plugins/jwt/rotation.test.ts +++ b/packages/better-auth/src/plugins/jwt/rotation.test.ts @@ -98,6 +98,257 @@ describe("jwt rotation", async () => { expect(jwksAfterGrace.keys.length).toBe(1); // First key should be gone expect(jwksAfterGrace.keys[0]?.kid).toBe(storage[1]!.id); - vi.useRealTimers(); + vi.useRealTimers(); +}); + +it("should respect cooldown period and prevent rapid rotation", async () => { + vi.useFakeTimers(); + const storage: Jwk[] = []; + const rotationInterval = 1; // 1 second + + const { auth } = await getTestInstance({ + plugins: [ + jwt({ + jwks: { + rotationInterval, + }, + adapter: { + getJwks: async () => storage, + createJwk: async (data) => { + const key = { + ...data, + id: crypto.randomUUID(), + }; + storage.push(key); + return key; + }, + }, + }), + ], }); + + // Create first key + await auth.api.signJWT({ body: { payload: { sub: "user1" } } }); + expect(storage.length).toBe(1); + + // Advance time past rotation interval + vi.advanceTimersByTime(1100); + + // Trigger rotation + await auth.api.signJWT({ body: { payload: { sub: "user1" } } }); + expect(storage.length).toBe(2); + + // Advance time by 1 second (less than 5 minute cooldown) + vi.advanceTimersByTime(1000); + + // Try to rotate again - should be blocked by cooldown + await auth.api.signJWT({ body: { payload: { sub: "user1" } } }); + expect(storage.length).toBe(2); // No new key created + + // Advance time past cooldown (5 minutes + 1 second) + vi.advanceTimersByTime(5 * 60 * 1000 + 1000); + + // Now rotation should work + await auth.api.signJWT({ body: { payload: { sub: "user1" } } }); + expect(storage.length).toBe(3); + + vi.useRealTimers(); +}); + +it("should not rotate when automatic rotation is disabled", async () => { + vi.useFakeTimers(); + const storage: Jwk[] = []; + + const { auth } = await getTestInstance({ + plugins: [ + jwt({ + jwks: { + rotationInterval: 1, // 1 second + disableAutomaticRotation: true, + }, + adapter: { + getJwks: async () => storage, + createJwk: async (data) => { + const key = { + ...data, + id: crypto.randomUUID(), + }; + storage.push(key); + return key; + }, + }, + }), + ], + }); + + // Create initial key manually + await auth.api.rotateKey({ body: { force: true } }); + expect(storage.length).toBe(1); + + // Advance time past rotation interval + vi.advanceTimersByTime(1100); + + // Try to sign - should use existing key, not rotate + await auth.api.signJWT({ body: { payload: { sub: "user1" } } }); + expect(storage.length).toBe(1); // No rotation + + // Check JWKS endpoint - should also not rotate + await auth.api.getJwks(); + expect(storage.length).toBe(1); + + vi.useRealTimers(); +}); + +it("should allow manual rotation via rotateKey endpoint", async () => { + vi.useFakeTimers(); + const storage: Jwk[] = []; + + const { auth } = await getTestInstance({ + plugins: [ + jwt({ + jwks: { + disableAutomaticRotation: true, + }, + adapter: { + getJwks: async () => storage, + createJwk: async (data) => { + const key = { + ...data, + id: crypto.randomUUID(), + }; + storage.push(key); + return key; + }, + }, + }), + ], + }); + + // Create initial key manually + const initialResult = await auth.api.rotateKey({ body: { force: true } }); + expect(storage.length).toBe(1); + const firstKeyId = initialResult.keyId; + + // Manually trigger rotation with force + const result = await auth.api.rotateKey({ + body: { force: true }, + }); + + expect(result.rotated).toBe(true); + expect(result.keyId).not.toBe(firstKeyId); + expect(storage.length).toBe(2); + + vi.useRealTimers(); +}); + +it("should respect custom cooldown in manual rotation", async () => { + vi.useFakeTimers(); + const storage: Jwk[] = []; + + const { auth } = await getTestInstance({ + plugins: [ + jwt({ + adapter: { + getJwks: async () => storage, + createJwk: async (data) => { + const key = { + ...data, + id: crypto.randomUUID(), + }; + storage.push(key); + return key; + }, + }, + }), + ], + }); + + // Create initial key + await auth.api.signJWT({ body: { payload: { sub: "user1" } } }); + expect(storage.length).toBe(1); + + // Advance time slightly to differentiate keys + vi.advanceTimersByTime(100); + + // Force first rotation + const firstRotate = await auth.api.rotateKey({ + body: { force: true }, + }); + expect(firstRotate.rotated).toBe(true); + expect(storage.length).toBe(2); + + // Advance time slightly (less than cooldown) + vi.advanceTimersByTime(100); + + // Try to rotate again with custom cooldown - should not rotate (within cooldown) + const secondRotate = await auth.api.rotateKey({ + body: { cooldown: 10000 }, // 10 seconds + }); + expect(secondRotate.rotated).toBe(false); + expect(secondRotate.keyId).toBe(firstRotate.keyId); + expect(storage.length).toBe(2); + + // Advance past custom cooldown + vi.advanceTimersByTime(10100); + + // Now rotation should work + const thirdRotate = await auth.api.rotateKey({ + body: { cooldown: 10000 }, + }); + expect(thirdRotate.rotated).toBe(true); + expect(thirdRotate.keyId).not.toBe(firstRotate.keyId); + expect(storage.length).toBe(3); + + vi.useRealTimers(); +}); + +it("should force rotation bypassing cooldown", async () => { + vi.useFakeTimers(); + const storage: Jwk[] = []; + + const { auth } = await getTestInstance({ + plugins: [ + jwt({ + adapter: { + getJwks: async () => storage, + createJwk: async (data) => { + const key = { + ...data, + id: crypto.randomUUID(), + }; + storage.push(key); + return key; + }, + }, + }), + ], + }); + + // Create initial key + await auth.api.signJWT({ body: { payload: { sub: "user1" } } }); + expect(storage.length).toBe(1); + + // Force rotate multiple times rapidly + const firstRotate = await auth.api.rotateKey({ + body: { force: true }, + }); + expect(firstRotate.rotated).toBe(true); + + const secondRotate = await auth.api.rotateKey({ + body: { force: true }, + }); + expect(secondRotate.rotated).toBe(true); + expect(secondRotate.keyId).not.toBe(firstRotate.keyId); + + const thirdRotate = await auth.api.rotateKey({ + body: { force: true }, + }); + expect(thirdRotate.rotated).toBe(true); + expect(thirdRotate.keyId).not.toBe(secondRotate.keyId); + + // Should have created 4 keys total (1 initial + 3 forced rotations) + expect(storage.length).toBe(4); + + vi.useRealTimers(); +}); }); diff --git a/packages/better-auth/src/plugins/jwt/sign.ts b/packages/better-auth/src/plugins/jwt/sign.ts index 4d7b2c99a7..9a47bc3c66 100644 --- a/packages/better-auth/src/plugins/jwt/sign.ts +++ b/packages/better-auth/src/plugins/jwt/sign.ts @@ -56,10 +56,11 @@ export async function signJWT( if (options?.jwks?.disableAutomaticRotation) { const adapter = getJwksAdapter(ctx.context.adapter, options); key = await adapter.getLatestKey(ctx); - } else { - key = await rotateJwk(ctx, options); + const result = await rotateJwk(ctx, options); + key = result.key; } + if (!key) { throw new BetterAuthError( "No signing key found. Automatic rotation is disabled. Please create a key manually using the /rotate-jwk endpoint or rotateJwk() function.", diff --git a/packages/better-auth/src/plugins/jwt/utils.ts b/packages/better-auth/src/plugins/jwt/utils.ts index 0283eae50a..f36c59f39e 100644 --- a/packages/better-auth/src/plugins/jwt/utils.ts +++ b/packages/better-auth/src/plugins/jwt/utils.ts @@ -185,7 +185,7 @@ export async function createJwk( * @param ctx - The endpoint context * @param options - JWT plugin options * @param config - Rotation configuration - * @returns The current valid key (either existing or newly created) + * @returns Object containing the current key and whether rotation occurred */ export async function rotateJwk( ctx: GenericEndpointContext, @@ -204,22 +204,39 @@ export async function rotateJwk( */ force?: boolean; }, -): Promise { +): Promise<{ key: Jwk; rotated: boolean }> { const adapter = getJwksAdapter(ctx.context.adapter, options); const latestKey = await adapter.getLatestKey(ctx); const ROTATION_COOLDOWN = config?.cooldown ?? 5 * 60 * 1000; // 5 minutes default const now = Date.now(); + + let shouldRotate = false; - const shouldRotate = config?.force || - !latestKey || - (latestKey.expiresAt && - latestKey.expiresAt < new Date() && - latestKey.createdAt.getTime() < now - ROTATION_COOLDOWN); - - if (shouldRotate) { - return await createJwk(ctx, options); + if (config?.force) { + shouldRotate = true; + } else if (!latestKey) { + shouldRotate = true; + } else if (latestKey.expiresAt && latestKey.expiresAt < new Date()) { + const wasCreatedAfterExpiration = latestKey.createdAt.getTime() > latestKey.expiresAt.getTime(); + + if (wasCreatedAfterExpiration) { + const keyAge = now - latestKey.createdAt.getTime(); + shouldRotate = keyAge >= ROTATION_COOLDOWN; + } else { + shouldRotate = true; + } } - return latestKey; + if (shouldRotate) { + const key = await createJwk(ctx, options); + return { key, rotated: true }; + } + + if (!latestKey) { + const key = await createJwk(ctx, options); + return { key, rotated: true }; + } + + return { key: latestKey, rotated: false }; }