add key rotation to jwks endpoint

core/mod-auth/src/oauth_jwks.test.ts

@@ -0,0 +1,135 @@
+import { expect, test, vi, beforeEach, afterEach } from "vitest";
+import { JWKSStore } from "./oauth_jwks.ts";
+import type { KeyService } from "@noctf/server-core/services/key";
+
+const mockStaticKey = Buffer.alloc(32, "a");
+let mockKeyService: KeyService;
+
+beforeEach(() => {
+  mockKeyService = {
+    deriveKey: vi.fn().mockReturnValue(mockStaticKey),
+  } as unknown as KeyService;
+
+  vi.useFakeTimers();
+});
+
+afterEach(() => {
+  vi.useRealTimers();
+});
+
+test("generates and stores public keys correctly", async () => {
+  const mockTime = 1781438400000;
+  vi.setSystemTime(mockTime);
+
+  const store = new JWKSStore(mockKeyService, 86400000);
+  const validKeys = await store.getValidKeys();
+
+  expect(validKeys).toHaveLength(2);
+  expect(validKeys[0].not_before).toBe(1781481600000);
+  expect(validKeys[1].not_before).toBe(1781395200000);
+
+  await expect(mockKeyService.deriveKey).toHaveBeenCalledWith(
+    "auth:jwk:1781481600000",
+  );
+  await expect(mockKeyService.deriveKey).toHaveBeenCalledWith(
+    "auth:jwk:1781395200000",
+  );
+
+  for (const key of validKeys) {
+    expect(key.pub).toMatchObject({
+      kty: "OKP",
+      crv: "Ed25519",
+      alg: "EdDSA",
+      use: "sig",
+    });
+    expect(typeof key.pub.x).toBe("string");
+    expect(typeof key.pub.kid).toBe("string");
+  }
+});
+
+test("returns the correct active signing key based on time", async () => {
+  vi.setSystemTime(864000000);
+  const store1 = new JWKSStore(mockKeyService, 86400000);
+  const signingKey1 = await store1.getSigningKey();
+  expect(signingKey1.not_before).toBe(864000000);
+  const validKeys1 = await store1.getValidKeys();
+  expect(validKeys1[0].not_before).toBe(864000000);
+  expect(validKeys1[1].not_before).toBe(777600000);
+
+  vi.setSystemTime(854000000);
+  const store2 = new JWKSStore(mockKeyService, 86400000);
+  const signingKey2 = await store2.getSigningKey();
+  expect(signingKey2.not_before).toBe(777600000);
+  const validKeys2 = await store2.getValidKeys();
+  expect(validKeys2[0].not_before).toBe(864000000);
+  expect(validKeys2[1].not_before).toBe(777600000);
+});
+
+test("caches keys and does not recreate them within the same epoch", async () => {
+  vi.setSystemTime(864000000);
+  const store = new JWKSStore(mockKeyService, 86400000);
+
+  await store.getSigningKey();
+  await store.getValidKeys();
+  await store.getSigningKey();
+
+  await expect(mockKeyService.deriveKey).toHaveBeenCalledTimes(2);
+});
+
+test("rotates keys and generates new ones when time advances to next epoch", async () => {
+  vi.setSystemTime(864000000);
+  const store = new JWKSStore(mockKeyService, 86400000);
+
+  await store.getSigningKey();
+  await expect(mockKeyService.deriveKey).toHaveBeenCalledTimes(2);
+
+  vi.setSystemTime(993600000);
+
+  const signingKey = await store.getSigningKey();
+  expect(signingKey.not_before).toBe(950400000);
+
+  const validKeys = await store.getValidKeys();
+  expect(validKeys[0].not_before).toBe(1036800000);
+  expect(validKeys[1].not_before).toBe(950400000);
+
+  await expect(mockKeyService.deriveKey).toHaveBeenCalledTimes(4);
+});
+
+test("properly handles key generation and signing key selection around the half-epoch boundary", async () => {
+  const epoch = 86400000;
+  const halfEpoch = Math.floor(epoch / 2);
+  const baseEpoch = 864000000; // 10 days in ms
+
+  // 1. Just before the half-epoch boundary
+  const timeBefore = baseEpoch + halfEpoch - 1; // 907199999
+  vi.setSystemTime(timeBefore);
+  const storeBefore = new JWKSStore(mockKeyService, epoch);
+
+  const signingKeyBefore = await storeBefore.getSigningKey();
+  expect(signingKeyBefore.not_before).toBe(baseEpoch);
+
+  const validKeysBefore = await storeBefore.getValidKeys();
+  expect(validKeysBefore).toHaveLength(2);
+  expect(validKeysBefore[0].not_before).toBe(baseEpoch);
+  expect(validKeysBefore[1].not_before).toBe(baseEpoch - epoch);
+
+  // 2. Exactly at the half-epoch boundary
+  const timeAt = baseEpoch + halfEpoch; // 907200000
+  vi.setSystemTime(timeAt);
+  const storeAt = new JWKSStore(mockKeyService, epoch);
+
+  const signingKeyAt = await storeAt.getSigningKey();
+  // The signing key should still be the base epoch key
+  expect(signingKeyAt.not_before).toBe(baseEpoch);
+
+  const validKeysAt = await storeAt.getValidKeys();
+  expect(validKeysAt).toHaveLength(2);
+  expect(validKeysAt[0].not_before).toBe(baseEpoch + epoch);
+  expect(validKeysAt[1].not_before).toBe(baseEpoch);
+});
+
+test("max token validity is half of key validity", async () => {
+  const store = new JWKSStore(mockKeyService, 86400000);
+  expect(store.getMaxMessageValidityMs()).toEqual(43200000);
+});
+

core/mod-auth/src/oauth_jwks.ts

@@ -3,22 +3,34 @@ import type { KeyService } from "@noctf/server-core/services/key";
 import { CryptoKey, importJWK, JWK } from "jose";
 import { createHash } from "node:crypto";
 
-const EPOCH_SECONDS = 86400;
+type Jwk = { pub: JWK; secret: CryptoKey; not_before: number };
 
 export class JWKSStore {
-  private key: { pub: JWK; secret: CryptoKey };
+  private keys: Jwk[] = [];
+  private readonly halfEpoch: number;
 
-  constructor(private readonly keyService: KeyService) {}
+  constructor(
+    private readonly keyService: KeyService,
+    private epoch: number,
+  ) {
+    this.halfEpoch = Math.floor(epoch / 2);
+  }
 
-  async getKey() {
-    if (!this.key) {
-      this.key = await this.generateKey();
-    }
-    return this.key;
+  async getSigningKey(): Promise<Jwk> {
+    const keys = await this.ensureKeys();
+    return keys.find((x) => x.not_before <= Date.now())!;
   }
 
-  private async generateKey() {
-    const key = this.keyService.deriveKey("auth:jwk");
+  async getValidKeys() {
+    return await this.ensureKeys();
+  }
+
+  getMaxMessageValidityMs() {
+    return this.halfEpoch;
+  }
+
+  private async generateKey(not_before: number) {
+    const key = this.keyService.deriveKey(`auth:jwk:${not_before}`);
     const secret = {
       kty: "OKP",
       crv: "Ed25519",
@@ -39,6 +51,21 @@ export class JWKSStore {
       alg: secret.alg,
       use: secret.use,
     };
-    return { pub, secret: (await importJWK(secret)) as CryptoKey };
+    return { pub, secret: (await importJWK(secret)) as CryptoKey, not_before };
+  }
+
+  private async ensureKeys() {
+    const time = Date.now() - this.halfEpoch;
+    const e1 = Math.floor(time / this.epoch) * this.epoch;
+    const e2 = e1 + this.epoch;
+    if (
+      this.keys.length &&
+      this.keys[0].not_before === e2 &&
+      this.keys[1].not_before === e1
+    ) {
+      return this.keys;
+    }
+    this.keys = await Promise.all([e2, e1].map((e) => this.generateKey(e)));
+    return this.keys;
   }
 }

core/mod-auth/src/oauth_routes.ts

@@ -43,7 +43,8 @@ export default async function (fastify: FastifyInstance) {
     identityService,
     tokenService,
   );
-  const jwksStore = new JWKSStore(keyService);
+  // Tokens can be valid for half of this
+  const jwksStore = new JWKSStore(keyService, 7 * 86_400_000);
 
   identityService.register(provider);
   fastify.register(fastifyFormbody);
@@ -53,7 +54,7 @@ export default async function (fastify: FastifyInstance) {
   const signIdToken = async (clientId: string, userId: number) => {
     const membership = await teamService.getMembershipForUser(userId);
     const user = await userService.get(userId);
-    const key = await jwksStore.getKey();
+    const key = await jwksStore.getSigningKey();
 
     return await new SignJWT({
       "noctf.dev/team_id": membership?.team_id.toString(),
@@ -85,8 +86,9 @@ export default async function (fastify: FastifyInstance) {
     };
   });
 
-  route(fastify, GetOAuthJWKS, {}, async () => {
-    return { keys: [(await jwksStore.getKey()).pub] };
+  route(fastify, GetOAuthJWKS, {}, async (_request, reply) => {
+    reply.header('cache-control', 'max-age=3600, public');
+    return { keys: (await jwksStore.getValidKeys()).map(({ pub }) => pub) };
   });
 
   route(fastify, InitOAuth, {}, async (request) => {