From 06d6b7e0eed12b8b1acaf2b2a69cf73b5668c3ea Mon Sep 17 00:00:00 2001
From: Daniel Rivers <me@danielrivers.com>
Date: Thu, 6 Jun 2024 19:36:02 +0100
Subject: [PATCH] feat: add jwks caching

---
 lib/main.test.ts | 30 +++++++++++++++++++++++++++++-
 lib/main.ts      | 17 +++++++++++++++--
 lib/utils.ts     | 10 ++++++++++
 package.json     |  3 ++-
 4 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 lib/utils.ts

diff --git a/lib/main.test.ts b/lib/main.test.ts
index ee62e22..41c8902 100644
--- a/lib/main.test.ts
+++ b/lib/main.test.ts
@@ -1,7 +1,8 @@
 // setToken.test.ts
-import { describe, it, expect } from "vitest";
+import { describe, it, expect, vi } from "vitest";
 import { createHmac } from "crypto";
 import { validateToken } from "./main";
+import * as utils from "./utils";
 
 function base64UrlEncode(str: string) {
   return Buffer.from(str)
@@ -81,6 +82,33 @@ describe("Validate token", () => {
     });
   });
 
+  it("use cache when JWKS has already been fetched", async () => {
+    utils.JWKSCache.clear();
+
+    const getJWKSMock = vi.spyOn(utils, "getJWKS");
+
+    const token =
+      "eyJhbGciOiJSUzI1NiIsImtpZCI6IjRjOmZhOjllOmQ2OjQ3OjIzOmI3OjM5OmM3OjhmOjk3OjI4OjQ1OmExOjg0OjM1IiwidHlwIjoiSldUIn0.eyJkYXRhIjp7InVzZXIiOnsiZW1haWwiOiJkYW5pZWxAa2luZGUuY29tIiwiZmlyc3RfbmFtZSI6IkRhbmllbCIsImlkIjoia3BfNjViMjhkNzFiYmExNGZhMzgwZDU2ZDJkOGQzNTAzZGEiLCJpc19wYXNzd29yZF9yZXNldF9yZXF1ZXN0ZWQiOmZhbHNlLCJpc19zdXNwZW5kZWQiOmZhbHNlLCJsYXN0X25hbWUiOiJSaXZlcnMiLCJvcmdhbml6YXRpb25zIjpbeyJjb2RlIjoib3JnXzU5MGQ3ZjFhODZhIiwicGVybWlzc2lvbnMiOm51bGwsInJvbGVzIjpudWxsfV0sInBob25lIjpudWxsLCJ1c2VybmFtZSI6ImRhbmllbCJ9fSwiZXZlbnRfaWQiOiJldmVudF8wMThmMzMxYTMxNzhmN2ZlZjI4NGI5NWZlNjc3MDM4NCIsInNvdXJjZSI6ImFkbWluIiwidGltZXN0YW1wIjoiMjAyNC0wNS0wMVQxNzo0MTo0NS41OTIxNDUrMTA6MDAiLCJ0eXBlIjoidXNlci51cGRhdGVkIn0.hAxfcxDNnzN8_U7sovti71NElh5pqVe6UEFKgVD1ZygVJUdEhmjYQOOSr6Aixj2ySs_hujZBvCRWeqG6jNPYbHRiV5kx0XaL6g3cW1DCoqpTpkxXtjf18HNYHCJmsUqMiSwfYpmVcI7kaIDfd0XwhWWH5gRdjAAMDneEwMKANklTzR_g_kIl5cVW5eVWntC4rFsSjRVvGSNb-OMsy2GJLWXUF8fc8Qru56VkJImeOE6ZOMi6wBhtx7HhOZEcEFgQjRvHeoQKdVmEE3BRUnO_LXTMMSjvP_kyfrS4JMaGWHc6mc8k1hZo_maASLSuXMF8882LZnr96cJFMHj8irRAug";
+
+    const tokenPayload = {
+      token,
+      domain: "https://danielkinde.kinde.com",
+    };
+
+    await validateToken(tokenPayload);
+    await validateToken(tokenPayload);
+
+    // Provide an invalid token to ensure that the token re-recieved from the server
+    await validateToken({
+      ...tokenPayload,
+      token:
+        "eyJhbGciOiJSUzI1NiIsImtpZCI6IjRjOmZhOjllOmQ2OjQ3OjIzOmI3OjM5OmM3OjhmOjk3OjI4OjQ1OmExOjg0OjM1IiwidHlwIjoiSldUIn0.eyJkYXRhIjp7InVzZXIiOnsiZW1haWwiOiJtZSt3ZWJvb2tAZGFuaWVscml2ZXJzLmNvbSIsImZpcnN0X25hbWUiOiJhYSIsImlkIjoia3BfYmM2YjI4MTczZDZkNGRmYWI1NjU3NTg4NWIwMjE0YjEiLCJpc19wYXNzd29yZF9yZXNldF9yZXF1ZXN0ZWQiOmZhbHNlLCJpc19zdXNwZW5kZWQiOmZhbHNlLCJsYXN0X25hbWUiOiJhaGEiLCJvcmdhbml6YXRpb25zIjpbeyJjb2RlIjoib3JnXzU5MGQ3ZjFhODZhIiwicGVybWlzc2lvbnMiOm51bGwsInJvbGVzIjpudWxsfV0sInBob25lIjpudWxsLCJ1c2VybmFtZSI6bnVsbH19LCJldmVudF9pZCI6ImV2ZW50XzAxOGYyYTllNTkyZWNjZjUyMzI5MTgzYTQ1Y2QxOTU2Iiwic291cmNlIjoiYWRtaW4iLCJ0aW1lc3RhbXAiOiIyMDI0LTA0LTMwVDAyOjA5OjMxLjY0OTE2MisxMDowMCIsInR5cGUiOiJ1c2VyLnVwZGF0ZWQifQ.YIFd21Ek7R_hfpfEpAcwW5ebaDSDsT7TMYF5HTbg70CfWw36IDqKqQWKR6T1_vP0lI5s0xJlDptbjykvWfSm44fkz0LgjCWQhM_ENzTZiAa89pa2X1prjKH4vyS7lTqSCNXvCeYiAaFZSlr2X3s2aztASB4jGBDETziGCh_klNh4Gun3AcbkWOXz_QPm3YGNqgc3hYSBsLdOQbCQ_BxS2Wc60D3NAShVaodPrtOLC1bvY1vn_HucZHT9l-KuTKgY1st6D4er2K6DuHZaFBMMdvTaFQX5zN8OZltxeiucja4sg2vbtexryMdSdHY3y5Cz70dKWW6Ph2kHucK6xScQoQs",
+    });
+    await validateToken(tokenPayload);
+
+    expect(getJWKSMock).toHaveBeenCalledTimes(2);
+  });
+
   it("token is present but not valid", async () => {
     const token = jwtSign({
       header: { alg: "HS256", typ: "JWT" },
diff --git a/lib/main.ts b/lib/main.ts
index b0ec293..6644072 100644
--- a/lib/main.ts
+++ b/lib/main.ts
@@ -1,4 +1,5 @@
-import { createRemoteJWKSet, jwtVerify } from "jose";
+import { jwtVerify } from "jose";
+import { JWKSCache, getJWKS } from "./utils";
 
 export type jwtValidationResponse = {
   valid: boolean;
@@ -8,13 +9,25 @@ export type jwtValidationResponse = {
 async function verifyJwt(
   token: string,
   domain: string,
+  forceJWKSFetch?: boolean,
 ): Promise<jwtValidationResponse> {
-  const JWKS = createRemoteJWKSet(new URL(`${domain}/.well-known/jwks.json`));
+  if (forceJWKSFetch) {
+    JWKSCache.delete(domain);
+  }
+  let JWKS = JWKSCache.get(domain);
+
+  if (!JWKS) {
+    JWKS = getJWKS(domain);
+    JWKSCache.set(domain, JWKS);
+  }
 
   try {
     await jwtVerify(token, JWKS);
     return { valid: true, message: "Token is valid" };
   } catch (error) {
+    if (!forceJWKSFetch) {
+      return verifyJwt(token, domain, true);
+    }
     return {
       valid: false,
       message: error instanceof Error ? error.message : "Unknown Error",
diff --git a/lib/utils.ts b/lib/utils.ts
new file mode 100644
index 0000000..f1cb65f
--- /dev/null
+++ b/lib/utils.ts
@@ -0,0 +1,10 @@
+import { type JWK, createRemoteJWKSet } from "jose";
+
+export const JWKSCache: Map<string, JWK.KeyStore> = new Map<
+  string,
+  JWK.KeyStore
+>();
+
+export function getJWKS(domain: string) {
+  return createRemoteJWKSet(new URL(`${domain}/.well-known/jwks.json`));
+}
diff --git a/package.json b/package.json
index a193416..6105483 100644
--- a/package.json
+++ b/package.json
@@ -32,5 +32,6 @@
   },
   "dependencies": {
     "jose": "^5.2.4"
-  }
+  },
+  "packageManager": "pnpm@8.14.1+sha256.2df78e65d433d7693b9d3fbdaf431b2d96bb4f96a2ffecd51a50efe16e50a6a8"
 }