fix(auth): keep one live key per kind, and report a key's own algorithm

Two problems found in review, both in the key store.

Nothing stopped a kind from having two live keys, and the bootstrap path
walks straight into it: two workers starting against an empty table both
find no key and both insert one. From then on each signs and encrypts with
its own. That is not the harmless split the comment here claimed — the
issuer reaches for a single key rather than the published set when it
decrypts a session cookie and when it verifies an access token, so a cookie
written by one worker is unreadable to the other and a token minted by one
is rejected by the other. It stays silent until someone cannot sign in.

A partial unique index over the kind, where the key has not been retired,
makes the second insert a dropped write instead. Both workers then read the
table again and use the key that won, which is all that matters. The
conflict clause stops naming a target: both indexes on the table mean the
same thing at this call site, that the row already exists in some form.

Creating a key is now attempted once rather than retried, because a store
declining the write is an expected answer and spinning on it would hang the
request instead of failing it.

Separately, a key pair reported the algorithm the issuer currently uses
rather than the one stored on the key it was built from, so a retained key
would advertise the wrong algorithm in a token header and in the JWKS after
a rotation — which defeats keeping it. The material was already being
imported with the stored value; only what was handed back disagreed.

Retiring a key and creating its replacement now have to happen together, so
that a kind never has two live keys and never has none.

packages/auth/src/key.ts

@@ -42,6 +42,23 @@ export interface StoredKey {
 export interface KeyStore {
 	/** Every key of a kind, expired ones included. Order does not matter. */
 	list(kind: KeyKind): Promise<StoredKey[]>;
+
+	/**
+	 * Add a key, unless the kind already has a live one.
+	 *
+	 * A store that lets two live keys of one kind exist at the same time
+	 * breaks things that are hard to see. Two issuers starting against an empty
+	 * store both find nothing, both generate a key, and from then on each signs
+	 * and encrypts with its own: a session cookie written by one is
+	 * undecryptable to the other, and an access token minted by one is rejected
+	 * by the other as invalid — because both reach for a single key rather than
+	 * trying the published set. Losing the second write is the whole point, so
+	 * this is not an error and reports nothing; the caller reads the list again
+	 * and uses whichever key survived.
+	 *
+	 * Retiring a key and creating its replacement therefore have to happen
+	 * together, so that the kind never has two live keys and never has none.
+	 */
 	create(kind: KeyKind, key: StoredKey): Promise<void>;
 }
 
@@ -56,6 +73,13 @@ function prefix(kind: KeyKind): string {
  * The default, and what every deployment used before `keyStore` existed — the
  * keys are read and written under exactly the prefixes they always were, so an
  * issuer that does not pass a store keeps finding the keys it already had.
+ *
+ * It cannot honour the one-live-key rule `create` asks for: through get and
+ * set there is no way to make "write unless one exists" a single operation.
+ * Two issuers bootstrapping against an empty store at the same moment will
+ * therefore end up with a key each, with the consequences described above. A
+ * store that can express a conditional write does not have this problem, and
+ * is what a deployment running more than one instance wants.
  */
 export function StorageKeyStore(storage: StorageAdapter): KeyStore {
 	return {

packages/auth/src/keys.ts

@@ -37,7 +37,7 @@ async function toKeyPair(kind: KeyKind, stored: StoredKey): Promise<KeyPair> {
 	if (kind === 'signing') jwk.use = 'sig';
 	return {
 		id: stored.id,
-		alg: alg[kind],
+		alg: stored.alg,
 		created: new Date(stored.created),
 		expired: stored.expired ? new Date(stored.expired) : undefined,
 		public: publicKey,
@@ -49,17 +49,30 @@ async function toKeyPair(kind: KeyKind, stored: StoredKey): Promise<KeyPair> {
 /**
  * Every key of a kind, newest first, creating one if none is usable.
  *
- * Expired keys are returned alongside live ones and sorted after them: the
- * first entry is what signs, and the rest are what still verifies. Retiring a
- * key therefore does not invalidate the tokens it signed — they age out on
- * their own — which is the only way a rotation is not also a mass sign-out.
+ * Expired keys are returned alongside live ones and sorted after them, so the
+ * first entry is the one that signs. Publishing the rest is what lets a
+ * verifier reading the JWKS still check a token signed before a rotation.
+ *
+ * A store may refuse the write, and is expected to when another issuer has
+ * already created a key for this kind — see the note on {@link KeyStore.create}
+ * about why two live keys of one kind is not a state this can be left in. So
+ * the created key is never returned directly: the list is read again, and
+ * whichever key the store actually kept is the one everybody uses.
  */
-async function keysOf(store: KeyStore, kind: KeyKind): Promise<KeyPair[]> {
+async function keysOf(store: KeyStore, kind: KeyKind, bootstrapped = false): Promise<KeyPair[]> {
 	const stored = await store.list(kind);
 	const results = await Promise.all(stored.map((k) => toKeyPair(kind, k)));
 	results.sort((a, b) => b.created.getTime() - a.created.getTime());
 	if (results.some((item) => !item.expired)) return results;
 
+	// One attempt, and then an error rather than another try. A store that
+	// accepts neither the write nor another issuer's would otherwise spin here
+	// forever, and a request that hangs is a worse way to learn about it than
+	// a request that fails.
+	if (bootstrapped) {
+		throw new Error(`Unable to create a ${kind} key: the store reports none after writing one.`);
+	}
+
 	const key = await generateKeyPair(alg[kind], { extractable: true });
 	const created: StoredKey = {
 		id: crypto.randomUUID(),
@@ -69,9 +82,7 @@ async function keysOf(store: KeyStore, kind: KeyKind): Promise<KeyPair[]> {
 		alg: alg[kind]
 	};
 	await store.create(kind, created);
-	// Read back rather than returning what was just built, so that two issuers
-	// starting at once converge on whichever key the store actually kept.
-	return keysOf(store, kind);
+	return keysOf(store, kind, true);
 }
 
 export function signingKeys(store: KeyStore): Promise<KeyPair[]> {

packages/core/migrations/0011_auth_state_in_postgres.sql

@@ -14,11 +14,22 @@
 -- `update ... where time_used is null returning *`, so exactly one caller is
 -- told it went first.
 --
--- `auth_key` is different: nothing races for it. It is here because it is the
--- one record whose loss ends every session at once, and a cache is a place
--- things are allowed to be evicted from. Keys are retired by setting
--- `expired_at`, never deleted, so the tokens they signed stay verifiable until
--- they expire on their own.
+-- `auth_key` is here for a different reason: it is the one record whose loss
+-- ends every session at once, and a cache is a place things are allowed to be
+-- evicted from. It gets a conditional write too, though. At most one key of a
+-- kind may be live, which `auth_key_one_live_per_kind` enforces rather than
+-- leaving to convention — without it, two workers starting against an empty
+-- table both find no key, both insert one, and each then signs and encrypts
+-- with its own. A session cookie written by one is undecryptable to the other
+-- and a token minted by one is rejected by the other, because both reach for a
+-- single key rather than trying the whole published set. The split is silent
+-- until someone cannot sign in. With the index the second insert is dropped and
+-- both workers use the key that won.
+--
+-- Keys are retired by setting `expired_at`, never deleted, so a verifier
+-- reading the published JWKS can still check a token signed before a rotation.
+-- Retiring one and creating its replacement have to happen together, so the
+-- kind never has two live keys and never has none.
 --
 -- Both credential tables store a hash and never the credential. An
 -- authorization code travels in a query string and a refresh token resumes a
@@ -85,4 +96,5 @@ CREATE UNIQUE INDEX "authorization_code_hash_unique" ON "authorization_code" USI
 CREATE UNIQUE INDEX "refresh_token_hash_unique" ON "refresh_token" USING btree ("token_hash");--> statement-breakpoint
 CREATE INDEX "refresh_token_subject_idx" ON "refresh_token" USING btree ("subject");--> statement-breakpoint
 CREATE UNIQUE INDEX "auth_key_key_id_unique" ON "auth_key" USING btree ("key_id");--> statement-breakpoint
+CREATE UNIQUE INDEX "auth_key_one_live_per_kind" ON "auth_key" USING btree ("kind") WHERE "auth_key"."expired_at" is null;--> statement-breakpoint
 CREATE UNIQUE INDEX "auth_kv_key_unique" ON "auth_kv" USING btree ("key");

packages/core/migrations/meta/0011_snapshot.json

@@ -1,5 +1,5 @@
 {
-  "id": "e2a5e726-9f05-4bf0-9bd9-27cb7501f26e",
+  "id": "6fdda454-c7ee-42c1-931b-fa595385aac0",
   "prevId": "8915fb5b-8f0d-4f6f-a808-27b19bb4604a",
   "version": "7",
   "dialect": "postgresql",
@@ -528,6 +528,22 @@
           "concurrently": false,
           "method": "btree",
           "with": {}
+        },
+        "auth_key_one_live_per_kind": {
+          "name": "auth_key_one_live_per_kind",
+          "columns": [
+            {
+              "expression": "kind",
+              "isExpression": false,
+              "asc": true,
+              "nulls": "last"
+            }
+          ],
+          "isUnique": true,
+          "where": "\"auth_key\".\"expired_at\" is null",
+          "concurrently": false,
+          "method": "btree",
+          "with": {}
         }
       },
       "foreignKeys": {},

packages/core/migrations/meta/_journal.json

@@ -82,7 +82,7 @@
     {
       "idx": 11,
       "version": "7",
-      "when": 1788605264671,
+      "when": 1788607804606,
       "tag": "0011_auth_state_in_postgres",
       "breakpoints": true
     }

packages/core/src/auth/signing-key.sql.ts

@@ -1,3 +1,4 @@
+import { sql } from 'drizzle-orm';
 import { pgEnum, pgTable, text, uniqueIndex } from 'drizzle-orm/pg-core';
 
 import { id, timestamps, utc } from '../db/types.js';
@@ -14,9 +15,17 @@ export const AuthKeyKindEnum = pgEnum('auth_key_kind', ['signing', 'encryption']
  * not query patterns, just that this is the row nobody can afford to have
  * quietly evicted.
  *
- * Retiring a key sets `expired_at` instead of deleting it. The tokens it
- * signed stay verifiable until they age out on their own, which is what keeps
- * a rotation from also being a mass sign-out.
+ * Retiring a key sets `expired_at` instead of deleting it, so a verifier
+ * reading the published JWKS can still check a token signed before a rotation.
+ *
+ * At most one key of a kind may be live at a time, and the partial unique index
+ * is what enforces it rather than a convention. Without it, two workers
+ * starting against an empty table both find no key, both insert one, and each
+ * then signs and encrypts with its own — so a session cookie written by one is
+ * undecryptable to the other, and an access token minted by one is rejected by
+ * the other. Both reach for a single key rather than trying the whole set, so
+ * the split is silent until someone cannot sign in. With the index the second
+ * insert is dropped, both read the table again, and both use the key that won.
  */
 export const AuthKeyTable = pgTable(
 	'auth_key',
@@ -36,5 +45,10 @@ export const AuthKeyTable = pgTable(
 		/** Null while the key is still in use. Set when it is retired. */
 		expiredAt: utc('expired_at')
 	},
-	(t) => [uniqueIndex('auth_key_key_id_unique').on(t.keyId)]
+	(t) => [
+		uniqueIndex('auth_key_key_id_unique').on(t.keyId),
+		uniqueIndex('auth_key_one_live_per_kind')
+			.on(t.kind)
+			.where(sql`${t.expiredAt} is null`)
+	]
 );

packages/core/src/auth/signing-key.test.ts

@@ -71,4 +71,36 @@ describe('PostgresKeyStore', () => {
 
 		expect(await store.list('signing')).toHaveLength(1);
 	});
+
+	test('a kind cannot end up with two live keys', async () => {
+		await store.create('signing', key());
+		// Dropped rather than refused: the caller reads the list again and uses
+		// whichever key is there, so losing this write is the intended outcome.
+		await store.create('signing', key());
+
+		expect(await store.list('signing')).toHaveLength(1);
+	});
+
+	/**
+	 * The bootstrap race, which is what the constraint is for. Two workers
+	 * starting against an empty table would otherwise end up with a key each,
+	 * and from then on neither can read what the other wrote.
+	 */
+	test('simultaneous bootstraps converge on one key', async () => {
+		await Promise.all(Array.from({ length: 5 }, () => store.create('signing', key())));
+
+		expect(await store.list('signing')).toHaveLength(1);
+	});
+
+	test('a replacement is allowed once the previous key is retired', async () => {
+		const first = key();
+		await store.create('signing', first);
+		await sql`update auth_key set expired_at = now() where key_id = ${first.id}`;
+
+		await store.create('signing', key());
+
+		const all = await store.list('signing');
+		expect(all).toHaveLength(2);
+		expect(all.filter((k) => !k.expired)).toHaveLength(1);
+	});
 });

packages/core/src/auth/signing-key.ts

@@ -21,11 +21,15 @@ function toStored(row: Row): StoredKey {
 /**
  * The issuer's signing and encryption keys, in Postgres.
  *
- * Two issuers starting against an empty table both generate a key and both
- * insert one, and that is fine: each is valid, both are published in the JWKS,
- * so a token signed by either verifies against either issuer. The conflict
- * clause is not for that race — a key id is a fresh UUID and cannot collide
- * with another issuer's — it is so that a retried write is not an error.
+ * `create` drops its write when the kind already has a live key, which is what
+ * the interface asks for and what keeps two workers bootstrapping at the same
+ * moment from ending up with a key each. Which of them wins does not matter.
+ * That they end up agreeing does — a key each means cookies one worker writes
+ * are unreadable to the other, and tokens one mints are rejected by the other.
+ *
+ * The conflict clause names no target on purpose: both unique indexes on this
+ * table mean the same thing here, that the row we wanted already exists in some
+ * form, and the answer to either is to keep what is there and read it back.
  */
 export function PostgresKeyStore(): KeyStore {
 	return {
@@ -52,7 +56,7 @@ export function PostgresKeyStore(): KeyStore {
 						privateKey: key.privateKey,
 						expiredAt: key.expired ? new Date(key.expired) : null
 					})
-					.onConflictDoNothing({ target: AuthKeyTable.keyId });
+					.onConflictDoNothing();
 			});
 		}
 	};