From 7fd0b503837b378ff9a8764290938014a2a5a11a Mon Sep 17 00:00:00 2001
From: Joachim Van Herwegen <joachimvh@gmail.com>
Date: Wed, 8 Feb 2023 14:16:24 +0100
Subject: [PATCH] fix: Store internal JWK as JWKS to be backwards compatible

---
 .../configuration/CachedJwkGenerator.ts        | 18 ++++++++++--------
 .../configuration/CachedJwkGenerator.test.ts   |  3 ++-
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/identity/configuration/CachedJwkGenerator.ts b/src/identity/configuration/CachedJwkGenerator.ts
index 8250d74cc..058fc0cc2 100644
--- a/src/identity/configuration/CachedJwkGenerator.ts
+++ b/src/identity/configuration/CachedJwkGenerator.ts
@@ -1,7 +1,7 @@
 import { createPublicKey } from 'crypto';
 import type { KeyObject } from 'crypto';
 import { exportJWK, generateKeyPair, importJWK } from 'jose';
-import type { AsymmetricSigningAlgorithm } from 'oidc-provider';
+import type { JWKS, AsymmetricSigningAlgorithm } from 'oidc-provider';
 import type { KeyValueStorage } from '../../storage/keyvalue/KeyValueStorage';
 import type { AlgJwk, JwkGenerator } from './JwkGenerator';
 
@@ -17,12 +17,12 @@ export class CachedJwkGenerator implements JwkGenerator {
   public readonly alg: AsymmetricSigningAlgorithm;
 
   private readonly key: string;
-  private readonly storage: KeyValueStorage<string, AlgJwk>;
+  private readonly storage: KeyValueStorage<string, JWKS>;
 
   private privateJwk?: AlgJwk;
   private publicJwk?: AlgJwk;
 
-  public constructor(alg: AsymmetricSigningAlgorithm, storageKey: string, storage: KeyValueStorage<string, AlgJwk>) {
+  public constructor(alg: AsymmetricSigningAlgorithm, storageKey: string, storage: KeyValueStorage<string, JWKS>) {
     this.alg = alg;
     this.key = storageKey;
     this.storage = storage;
@@ -33,10 +33,12 @@ export class CachedJwkGenerator implements JwkGenerator {
       return this.privateJwk;
     }
 
-    const jwk = await this.storage.get(this.key);
-    if (jwk) {
-      this.privateJwk = jwk;
-      return jwk;
+    // We store in JWKS format for backwards compatibility reasons.
+    // If we want to just store the key instead we will need some way to do the migration.
+    const jwks = await this.storage.get(this.key);
+    if (jwks) {
+      this.privateJwk = jwks.keys[0] as AlgJwk;
+      return this.privateJwk;
     }
 
     const { privateKey } = await generateKeyPair(this.alg);
@@ -45,7 +47,7 @@ export class CachedJwkGenerator implements JwkGenerator {
     const privateJwk = { ...await exportJWK(privateKey) } as AlgJwk;
     privateJwk.alg = this.alg;
 
-    await this.storage.set(this.key, privateJwk);
+    await this.storage.set(this.key, { keys: [ privateJwk ]});
     this.privateJwk = privateJwk;
     return privateJwk;
   }
diff --git a/test/unit/identity/configuration/CachedJwkGenerator.test.ts b/test/unit/identity/configuration/CachedJwkGenerator.test.ts
index f997adda1..5587bda40 100644
--- a/test/unit/identity/configuration/CachedJwkGenerator.test.ts
+++ b/test/unit/identity/configuration/CachedJwkGenerator.test.ts
@@ -1,5 +1,6 @@
 import { generateKeyPair, importJWK, jwtVerify, SignJWT } from 'jose';
 import * as jose from 'jose';
+import type { JWKS } from 'oidc-provider';
 import { CachedJwkGenerator } from '../../../../src/identity/configuration/CachedJwkGenerator';
 import type { AlgJwk } from '../../../../src/identity/configuration/JwkGenerator';
 import type { KeyValueStorage } from '../../../../src/storage/keyvalue/KeyValueStorage';
@@ -8,7 +9,7 @@ describe('A CachedJwkGenerator', (): void => {
   const alg = 'ES256';
   const storageKey = 'jwks';
   let storageMap: Map<string, AlgJwk>;
-  let storage: jest.Mocked<KeyValueStorage<string, AlgJwk>>;
+  let storage: jest.Mocked<KeyValueStorage<string, JWKS>>;
   let generator: CachedJwkGenerator;
 
   beforeEach(async(): Promise<void> => {