From 7c5d0c2cf3fe82c35fec020853daed8a14e20eb5 Mon Sep 17 00:00:00 2001
From: Bereket Engida <Bekacru@gmail.com>
Date: Sat, 29 Nov 2025 22:47:02 -0800
Subject: [PATCH] refactor(jwt): improve key rotation logic and test cases for
 cooldown handling

---
 .../src/plugins/jwt/rotation.test.ts            | 12 ++++++------
 packages/better-auth/src/plugins/jwt/utils.ts   | 17 +++++------------
 2 files changed, 11 insertions(+), 18 deletions(-)

diff --git a/packages/better-auth/src/plugins/jwt/rotation.test.ts b/packages/better-auth/src/plugins/jwt/rotation.test.ts
index ded9094a24..56404772e0 100644
--- a/packages/better-auth/src/plugins/jwt/rotation.test.ts
+++ b/packages/better-auth/src/plugins/jwt/rotation.test.ts
@@ -35,8 +35,8 @@ describe("jwt rotation", async () => {
 		expect(storage.length).toBe(1);
 		const firstKey = storage[0];
 
-		// Advance time past rotation interval
-		vi.advanceTimersByTime(1100);
+		// Advance time past rotation interval (and cooldown)
+		vi.advanceTimersByTime(5 * 60 * 1000 + 100);
 
 		// Second key creation (should rotate)
 		await auth.api.signJWT({
@@ -83,8 +83,8 @@ describe("jwt rotation", async () => {
 		// Advance time past rotation interval but within grace period
 		vi.advanceTimersByTime(1100);
 
-		// Trigger rotation by signing
-		await auth.api.signJWT({ body: { payload: { sub: "user1" } } });
+		// Trigger rotation (override cooldown to allow rotation)
+		await auth.api.rotateKey({ body: { cooldown: 500 } });
 		expect(storage.length).toBe(2);
 
 		// Check JWKS endpoint
@@ -131,8 +131,8 @@ it("should respect cooldown period and prevent rapid rotation", async () => {
 	await auth.api.signJWT({ body: { payload: { sub: "user1" } } });
 	expect(storage.length).toBe(1);
 
-	// Advance time past rotation interval
-	vi.advanceTimersByTime(1100);
+	// Advance time past rotation interval (and cooldown)
+	vi.advanceTimersByTime(5 * 60 * 1000 + 100);
 
 	// Trigger rotation
 	await auth.api.signJWT({ body: { payload: { sub: "user1" } } });
diff --git a/packages/better-auth/src/plugins/jwt/utils.ts b/packages/better-auth/src/plugins/jwt/utils.ts
index f36c59f39e..697b45a52f 100644
--- a/packages/better-auth/src/plugins/jwt/utils.ts
+++ b/packages/better-auth/src/plugins/jwt/utils.ts
@@ -177,11 +177,11 @@ export async function createJwk(
 
 /**
  * Rotates the JWT signing key if needed.
- * 
+ *
  * Rotation occurs when:
  * - No key exists, OR
  * - Latest key is expired AND was created more than the cooldown period ago
- * 
+ *
  * @param ctx - The endpoint context
  * @param options - JWT plugin options
  * @param config - Rotation configuration
@@ -210,22 +210,15 @@ export async function rotateJwk(
 
 	const ROTATION_COOLDOWN = config?.cooldown ?? 5 * 60 * 1000; // 5 minutes default
 	const now = Date.now();
-		
+
 	let shouldRotate = false;
-	
+
 	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;
-		}
+		shouldRotate = now - latestKey.createdAt.getTime() >= ROTATION_COOLDOWN;
 	}
 
 	if (shouldRotate) {