feat(jwt): enhance JWT key rotation with cooldown and manual control features

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,
 			});
 		},
 		),

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();
+});
 });

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.",

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<Jwk> {
+): 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 };
 }