chore: accept public key hash in kid header (#649)

* chore: accept public key ID in kid header

* fix build

Author: arcoraven

src/db/keypair/get.ts

@@ -2,14 +2,17 @@ import { Keypairs } from "@prisma/client";
 import { createHash } from "crypto";
 import { prisma } from "../client";
 
-export const getKeypairByPublicKey = async ({
-  publicKey,
-}: {
-  publicKey: string;
-}): Promise<Keypairs | null> => {
-  const hash = createHash("sha256").update(publicKey).digest("hex");
-
+export const getKeypairByHash = async (
+  hash: string,
+): Promise<Keypairs | null> => {
   return prisma.keypairs.findUnique({
     where: { hash },
   });
 };
+
+export const getKeypairByPublicKey = async (
+  publicKey: string,
+): Promise<Keypairs | null> => {
+  const hash = createHash("sha256").update(publicKey).digest("hex");
+  return getKeypairByHash(hash);
+};

src/server/middleware/auth.ts

@@ -165,18 +165,24 @@ export const onRequest = async ({
       const payload = decoded.payload as JwtPayload;
       const header = decoded.header;
 
-      // Get the public key from the `iss` payload field or `kid` header field.
-      const publicKey = payload.iss ?? header.kid;
+      // Get the public key from the `iss` payload field.
+      const publicKey = payload.iss;
       if (publicKey) {
         const authWallet = await getAuthWallet();
         if (publicKey === (await authWallet.getAddress())) {
           return await handleAccessToken(jwt, req, getUser);
         } else if (publicKey === THIRDWEB_DASHBOARD_ISSUER) {
           return await handleDashboardAuth(jwt);
         } else {
-          return await handleKeypairAuth(jwt, req, publicKey);
+          return await handleKeypairAuth({ jwt, req, publicKey });
         }
       }
+
+      // Get the public key hash from the `kid` header.
+      const publicKeyHash = header.kid;
+      if (publicKeyHash) {
+        return await handleKeypairAuth({ jwt, req, publicKeyHash });
+      }
     }
   }
 
@@ -291,19 +297,22 @@ const handleWebsocketAuth = async (
  * @param publicKey string
  * @returns AuthResponse
  */
-const handleKeypairAuth = async (
-  jwt: string,
-  req: FastifyRequest,
-  publicKey: string,
-): Promise<AuthResponse> => {
+const handleKeypairAuth = async (args: {
+  jwt: string;
+  req: FastifyRequest;
+  publicKey?: string;
+  publicKeyHash?: string;
+}): Promise<AuthResponse> => {
   // The keypair auth feature must be explicitly enabled.
   if (!env.ENABLE_KEYPAIR_AUTH) {
     return { isAuthed: false };
   }
 
+  const { jwt, req, publicKey, publicKeyHash } = args;
+
   let error: string | undefined;
   try {
-    const keypair = await getKeypair({ publicKey });
+    const keypair = await getKeypair({ publicKey, publicKeyHash });
     if (!keypair) {
       error = "The provided public key is incorrect or not added to Engine.";
       throw error;

src/utils/cache/keypair.ts

@@ -1,20 +1,39 @@
 import type { Keypairs } from "@prisma/client";
-import { getKeypairByPublicKey } from "../../db/keypair/get";
+import { getKeypairByHash, getKeypairByPublicKey } from "../../db/keypair/get";
 
 // Cache a public key to the Keypair object, or null if not found.
 export const keypairCache = new Map<string, Keypairs | null>();
 
-export const getKeypair = async ({
-  publicKey,
-}: {
-  publicKey: string;
+/**
+ * Get a keypair by public key or hash.
+ */
+export const getKeypair = async (args: {
+  publicKey?: string;
+  publicKeyHash?: string;
 }): Promise<Keypairs | null> => {
-  const cached = keypairCache.get(publicKey);
+  const { publicKey, publicKeyHash } = args;
+
+  const key = publicKey
+    ? `public-key:${args.publicKey}`
+    : publicKeyHash
+    ? `public-key-hash:${args.publicKeyHash}`
+    : null;
+
+  if (!key) {
+    throw new Error('Must provide "publicKey" or "publicKeyHash".');
+  }
+
+  const cached = keypairCache.get(key);
   if (cached) {
     return cached;
   }
 
-  const keypair = await getKeypairByPublicKey({ publicKey });
-  keypairCache.set(publicKey, keypair);
+  const keypair = publicKey
+    ? await getKeypairByPublicKey(publicKey)
+    : publicKeyHash
+    ? await getKeypairByHash(publicKeyHash)
+    : null;
+
+  keypairCache.set(key, keypair);
   return keypair;
 };