feat(jwt): add key rotation (#6147)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Paola Estefanía de Campos <84341268+Paola3stefania@users.noreply.github.com>

docs/content/docs/plugins/jwt.mdx

@@ -338,6 +338,12 @@ Table Name: `jwks`
       type: "Date", 
       description: "Timestamp of when the web key was created" 
     },
+    { 
+      name: "expiresAt", 
+      type: "Date", 
+      description: "Timestamp of when the web key expires",
+      isOptional: true
+    },
   ]}
   />
 
@@ -404,6 +410,24 @@ jwt({
 })
 ```
 
+### Key Rotation
+
+You can enable key rotation by setting the `rotationInterval` option. This will automatically rotate the key pair at the specified interval.
+
+The default value is `undefined` (disabled).
+
+```ts title="auth.ts"
+jwt({
+  jwks: {
+    rotationInterval: 60 * 60 * 24 * 30, // 30 days
+    gracePeriod: 60 * 60 * 24 * 30 // 30 days
+  }
+})
+```
+
+- `rotationInterval`: The interval in seconds to rotate the key pair.
+- `gracePeriod`: The period in seconds to keep the old key pair valid after rotation. This is useful to allow clients to verify tokens signed by the old key pair. The default value is 30 days.
+
 ### Modify JWT payload
 
 By default the entire user object is added to the JWT payload. You can modify the payload by providing a function to the `definePayload` option.

packages/better-auth/src/plugins/jwt/adapter.ts

@@ -19,19 +19,12 @@ export const getJwksAdapter = (
 			});
 		},
 		getLatestKey: async (ctx: GenericEndpointContext) => {
-			if (options?.adapter?.getLatestKey) {
-				return await options.adapter.getLatestKey(ctx);
+			if (options?.adapter?.getJwks) {
+				const keys = await options.adapter.getJwks(ctx);
+				return keys?.sort(
+					(a, b) => b.createdAt.getTime() - a.createdAt.getTime(),
+				)[0];
 			}
-			const key = await adapter.findMany<Jwk>({
-				model: "jwks",
-				sortBy: {
-					field: "createdAt",
-					direction: "desc",
-				},
-				limit: 1,
-			});
-
-			return key[0];
 		},
 		createJwk: async (ctx: GenericEndpointContext, webKey: Omit<Jwk, "id">) => {
 			if (options?.adapter?.createJwk) {

packages/better-auth/src/plugins/jwt/index.ts

@@ -134,13 +134,13 @@ export const jwt = (options?: JwtOptions | undefined) => {
 						throw new APIError("NOT_FOUND");
 					}
 
-					const adapter = getJwksAdapter(ctx.context.adapter);
+					const adapter = getJwksAdapter(ctx.context.adapter, options);
 
 					let keySets = await adapter.getAllKeys(ctx);
 
 					if (!keySets || keySets?.length === 0) {
-						const key = await createJwk(ctx, options);
-						keySets = [key];
+						await createJwk(ctx, options);
+						keySets = await adapter.getAllKeys(ctx);
 					}
 
 					if (!keySets?.length) {
@@ -148,6 +148,19 @@ export const jwt = (options?: JwtOptions | undefined) => {
 							"No key sets found. Make sure you have a key in your database.",
 						);
 					}
+
+					const now = Date.now();
+					const DEFAULT_GRACE_PERIOD = 60 * 60 * 24 * 30;
+					const gracePeriod =
+						(options?.jwks?.gracePeriod ?? DEFAULT_GRACE_PERIOD) * 1000;
+
+					const keys = keySets.filter((key) => {
+						if (!key.expiresAt) {
+							return true;
+						}
+						return key.expiresAt.getTime() + gracePeriod > now;
+					});
+
 					const keyPairConfig = options?.jwks?.keyPairConfig;
 					const defaultCrv = keyPairConfig
 						? "crv" in keyPairConfig
@@ -155,7 +168,7 @@ export const jwt = (options?: JwtOptions | undefined) => {
 							: undefined
 						: undefined;
 					return ctx.json({
-						keys: keySets.map((keySet) => {
+						keys: keys.map((keySet) => {
 							return {
 								alg: keySet.alg ?? options?.jwks?.keyPairConfig?.alg ?? "EdDSA",
 								crv: keySet.crv ?? defaultCrv,

packages/better-auth/src/plugins/jwt/jwt.test.ts

@@ -363,6 +363,12 @@ describe.each([
 		const jwks = await auth.api.getJwks();
 		const localJwks = createLocalJWKSet(jwks);
 		const decoded = await jwtVerify(jwt?.token!, localJwks);
+		
+		// Verify the kid from the JWT exists in the JWKS
+		const kidFromJwt = decoded.protectedHeader.kid;
+		const keyExists = jwks.keys.some((key) => key.kid === kidFromJwt);
+		expect(keyExists).toBe(true);
+		
 		expect(decoded).toMatchObject({
 			payload: {
 				iss: "https://example.com",
@@ -374,7 +380,7 @@ describe.each([
 			},
 			protectedHeader: {
 				alg: keyPairConfig.alg,
-				kid: jwks.keys[0]!.kid,
+				kid: expect.any(String),
 			},
 		});
 	});
@@ -713,9 +719,6 @@ describe("jwt - custom adapter", async () => {
 						getJwks: async () => {
 							return storage;
 						},
-						getLatestKey: async () => {
-							return storage[0] ?? null;
-						},
 						createJwk: async (data) => {
 							const key = {
 								...data,

packages/better-auth/src/plugins/jwt/rotation.test.ts

@@ -0,0 +1,103 @@
+import { describe, expect, it, vi } from "vitest";
+import { getTestInstance } from "../../test-utils/test-instance";
+import { jwt } from ".";
+import type { Jwk } from "./types";
+
+describe("jwt rotation", async () => {
+	it("should rotate keys when expired", async () => {
+		vi.useFakeTimers();
+		const storage: Jwk[] = [];
+		const { auth } = await getTestInstance({
+			plugins: [
+				jwt({
+					jwks: {
+						rotationInterval: 1, // 1 second
+					},
+					adapter: {
+						getJwks: async () => storage,
+						createJwk: async (data) => {
+							const key = {
+								...data,
+								id: crypto.randomUUID(),
+							};
+							storage.push(key);
+							return key;
+						},
+					},
+				}),
+			],
+		});
+
+		// First key creation
+		await auth.api.signJWT({
+			body: { payload: { sub: "user1" } },
+		});
+		expect(storage.length).toBe(1);
+		const firstKey = storage[0];
+
+		// Advance time past rotation interval
+		vi.advanceTimersByTime(1100);
+
+		// Second key creation (should rotate)
+		await auth.api.signJWT({
+			body: { payload: { sub: "user1" } },
+		});
+		expect(storage.length).toBe(2);
+		const secondKey = storage[1];
+		expect(secondKey!.id).not.toBe(firstKey!.id);
+
+		vi.useRealTimers();
+	});
+
+	it("should return keys within grace period", async () => {
+		vi.useFakeTimers();
+		const storage: Jwk[] = [];
+		const rotationInterval = 1; // 1 second
+		const gracePeriod = 1; // 1 second
+
+		const { auth } = await getTestInstance({
+			plugins: [
+				jwt({
+					jwks: {
+						rotationInterval,
+						gracePeriod,
+					},
+					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" } } });
+
+		// Advance time past rotation interval but within grace period
+		vi.advanceTimersByTime(1100);
+
+		// Trigger rotation by signing
+		await auth.api.signJWT({ body: { payload: { sub: "user1" } } });
+		expect(storage.length).toBe(2);
+
+		// Check JWKS endpoint
+		const jwks = await auth.api.getJwks();
+		expect(jwks.keys.length).toBe(2); // Both keys should be present
+
+		// Advance time past grace period
+		vi.advanceTimersByTime(1000);
+
+		const jwksAfterGrace = await auth.api.getJwks();
+		expect(jwksAfterGrace.keys.length).toBe(1); // First key should be gone
+		expect(jwksAfterGrace.keys[0]?.kid).toBe(storage[1]!.id);
+
+		vi.useRealTimers();
+	});
+});

packages/better-auth/src/plugins/jwt/schema.ts

@@ -15,6 +15,10 @@ export const schema = {
 				type: "date",
 				required: true,
 			},
+			expiresAt: {
+				type: "date",
+				required: false,
+			},
 		},
 	},
 } satisfies BetterAuthPluginDBSchema;

packages/better-auth/src/plugins/jwt/sign.ts

@@ -54,16 +54,12 @@ export async function signJWT(
 
 	const adapter = getJwksAdapter(ctx.context.adapter, options);
 	let key = await adapter.getLatestKey(ctx);
-	if (!key) {
+	if (!key || (key.expiresAt && key.expiresAt < new Date())) {
 		key = await createJwk(ctx, options);
 	}
 	const privateKeyEncryptionEnabled =
 		!options?.jwks?.disablePrivateKeyEncryption;
 
-	if (key === undefined) {
-		key = await createJwk(ctx, options);
-	}
-
 	let privateWebKey = privateKeyEncryptionEnabled
 		? await symmetricDecrypt({
 				key: ctx.context.secret,

packages/better-auth/src/plugins/jwt/types.ts

@@ -30,6 +30,18 @@ export interface JwtOptions {
 				 * @default false
 				 */
 				disablePrivateKeyEncryption?: boolean;
+				/**
+				 * The key rotation interval in seconds.
+				 *
+				 * @default undefined (disabled)
+				 */
+				rotationInterval?: number;
+				/**
+				 * The grace period in seconds.
+				 *
+				 * @default 2592000 (30 days)
+				 */
+				gracePeriod?: number;
 		  }
 		| undefined;
 
@@ -130,18 +142,6 @@ export interface JwtOptions {
 		getJwks?: (
 			ctx: GenericEndpointContext,
 		) => Promise<Jwk[] | null | undefined>;
-		/**
-		 * A custom function to get the latest key from the database or
-		 * other source
-		 *
-		 * This will override the default getLatestKey from the database
-		 *
-		 * @param ctx - The context of the request
-		 * @returns The latest key
-		 */
-		getLatestKey?: (
-			ctx: GenericEndpointContext,
-		) => Promise<Jwk | null | undefined>;
 		/**
 		 * A custom function to create a new key in the database or
 		 * other source
@@ -194,6 +194,7 @@ export interface Jwk {
 	publicKey: string;
 	privateKey: string;
 	createdAt: Date;
+	expiresAt?: Date;
 	alg?: JWSAlgorithms | undefined;
 	crv?: ("Ed25519" | "P-256" | "P-521") | undefined;
 }

packages/better-auth/src/plugins/jwt/utils.ts

@@ -160,6 +160,13 @@ export async function createJwk(
 				)
 			: stringifiedPrivateWebKey,
 		createdAt: new Date(),
+		...(options?.jwks?.rotationInterval
+			? {
+					expiresAt: new Date(
+						Date.now() + options.jwks.rotationInterval * 1000,
+					),
+				}
+			: {}),
 	};
 
 	const adapter = getJwksAdapter(ctx.context.adapter, options);