Improve type safety and compatibility with Bun

.changeset/eight-eyes-float.md

@@ -0,0 +1,5 @@
+---
+"@atproto/jwk-jose": patch
+---
+
+Explicit JOSE's importJWK() "alg" argument

.changeset/eighty-planes-sleep.md

@@ -0,0 +1,6 @@
+---
+"@atproto/jwk-jose": patch
+"@atproto/jwk": patch
+---
+
+Remove unsafe type casting during JWT verification

.changeset/strong-mice-talk.md

@@ -0,0 +1,5 @@
+---
+"@atproto/jwk": patch
+---
+
+Allow unknown properties in JWT payload & headers

.changeset/thin-mice-lick.md

@@ -0,0 +1,5 @@
+---
+"@atproto/jwk": patch
+---
+
+Expose ZodError to allow for proper schema validation error catching

packages/oauth/jwk-jose/src/jose-key.ts

@@ -1,4 +1,19 @@
-import { JwtVerifyError } from '@atproto/jwk'
+import {
+  Jwk,
+  JwkError,
+  JwtCreateError,
+  JwtHeader,
+  JwtPayload,
+  JwtVerifyError,
+  Key,
+  RequiredKey,
+  SignedJwt,
+  VerifyOptions,
+  VerifyResult,
+  jwkValidator,
+  jwtHeaderSchema,
+  jwtPayloadSchema,
+} from '@atproto/jwk'
 import {
   SignJWT,
   errors,
@@ -14,19 +29,6 @@ import {
   type KeyLike,
 } from 'jose'
 
-import {
-  Jwk,
-  JwkError,
-  JwtCreateError,
-  JwtHeader,
-  JwtPayload,
-  Key,
-  SignedJwt,
-  VerifyOptions,
-  VerifyPayload,
-  VerifyResult,
-  jwkValidator,
-} from '@atproto/jwk'
 import { either } from './util'
 
 const { JOSEError } = errors
@@ -36,52 +38,89 @@ export type Importable = string | KeyLike | Jwk
 export type { GenerateKeyPairOptions, GenerateKeyPairResult }
 
 export class JoseKey extends Key {
-  #keyObj?: KeyLike | Uint8Array
-
-  protected async getKey() {
+  protected async getKeyObj(alg: string) {
     try {
-      return (this.#keyObj ||= await importJWK(this.jwk as JWK))
+      return await importJWK(this.jwk as JWK, alg)
     } catch (cause) {
       throw new JwkError('Failed to import JWK', undefined, { cause })
     }
   }
 
-  async createJwt(header: JwtHeader, payload: JwtPayload) {
-    if (header.kid && header.kid !== this.kid) {
-      throw new JwtCreateError(
-        `Invalid "kid" (${header.kid}) used to sign with key "${this.kid}"`,
-      )
-    }
+  async createJwt(header: JwtHeader, payload: JwtPayload): Promise<SignedJwt> {
+    try {
+      const { kid } = header
+      if (kid && kid !== this.kid) {
+        throw new JwtCreateError(
+          `Invalid "kid" (${kid}) used to sign with key "${this.kid}"`,
+        )
+      }
 
-    if (!header.alg || !this.algorithms.includes(header.alg)) {
-      throw new JwtCreateError(
-        `Invalid "alg" (${header.alg}) used to sign with key "${this.kid}"`,
-      )
-    }
+      const { alg } = header
+      if (!alg || !this.algorithms.includes(alg)) {
+        throw new JwtCreateError(
+          `Invalid "alg" (${alg}) used to sign with key "${this.kid}"`,
+        )
+      }
 
-    const keyObj = await this.getKey()
-    return new SignJWT(payload)
-      .setProtectedHeader({ ...header, kid: this.kid })
-      .sign(keyObj) as Promise<SignedJwt>
+      const keyObj = await this.getKeyObj(alg)
+      const jwtBuilder = new SignJWT(payload).setProtectedHeader({
+        ...header,
+        alg,
+        kid: this.kid,
+      })
+
+      const signedJwt = await jwtBuilder.sign(keyObj)
+
+      return signedJwt as SignedJwt
+    } catch (cause) {
+      if (cause instanceof JOSEError) {
+        throw new JwtCreateError(cause.message, cause.code, { cause })
+      } else {
+        throw JwtCreateError.from(cause)
+      }
+    }
   }
 
-  async verifyJwt<
-    P extends VerifyPayload = JwtPayload,
-    C extends string = string,
-  >(token: SignedJwt, options?: VerifyOptions<C>): Promise<VerifyResult<P, C>> {
+  async verifyJwt<C extends string = never>(
+    token: SignedJwt,
+    options?: VerifyOptions<C>,
+  ): Promise<VerifyResult<C>> {
     try {
-      const keyObj = await this.getKey()
-      const result = await jwtVerify(token, keyObj, {
-        ...options,
-        algorithms: this.algorithms,
-      } as JWTVerifyOptions)
-
-      return result as VerifyResult<P, C>
-    } catch (error) {
-      if (error instanceof JOSEError) {
-        throw new JwtVerifyError(error.message, error.code, { cause: error })
+      const result = await jwtVerify(
+        token,
+        async ({ alg }) => this.getKeyObj(alg),
+        { ...options, algorithms: this.algorithms } as JWTVerifyOptions,
+      )
+
+      // @NOTE if all tokens are signed exclusively through createJwt(), then
+      // there should be no need to parse the payload and headers here. But
+      // since the JWT could have been signed with the same key from somewhere
+      // else, let's parse it to ensure the integrity (and type safety) of the
+      // data.
+      const headerParsed = jwtHeaderSchema.safeParse(result.protectedHeader)
+      if (!headerParsed.success) {
+        throw new JwtVerifyError('Invalid JWT header', undefined, {
+          cause: headerParsed.error,
+        })
+      }
+
+      const payloadParsed = jwtPayloadSchema.safeParse(result.payload)
+      if (!payloadParsed.success) {
+        throw new JwtVerifyError('Invalid JWT payload', undefined, {
+          cause: payloadParsed.error,
+        })
+      }
+
+      return {
+        protectedHeader: headerParsed.data,
+        // "requiredClaims" enforced by jwtVerify()
+        payload: payloadParsed.data as RequiredKey<JwtPayload, C>,
+      }
+    } catch (cause) {
+      if (cause instanceof JOSEError) {
+        throw new JwtVerifyError(cause.message, cause.code, { cause })
       } else {
-        throw JwtVerifyError.from(error)
+        throw JwtVerifyError.from(cause)
       }
     }
   }

packages/oauth/jwk-webcrypto/src/webcrypto-key.ts

@@ -63,7 +63,11 @@ export class WebcryptoKey extends JoseKey {
     throw new Error('Private Webcrypto Key not exportable')
   }
 
-  protected override async getKey() {
+  protected override async getKeyObj(alg: string) {
+    // TODO: compare with this.cryptoKeyPair.privateKey.algorithm.name ?
+    if (this.jwk.alg !== alg) {
+      throw new TypeError('Invalid algorithm')
+    }
     return this.cryptoKeyPair.privateKey
   }
 }

packages/oauth/jwk/src/index.ts

@@ -1,3 +1,7 @@
+// Since we expose zod schemas, let's expose ZodError so that dependents can
+// catch schema parsing errors.
+export { ZodError } from 'zod'
+
 export * from './alg.js'
 export * from './errors.js'
 export * from './jwk.js'

packages/oauth/jwk/src/jwt-verify.ts

@@ -1,7 +1,7 @@
 import { JwtHeader, JwtPayload } from './jwt.js'
 import { RequiredKey } from './util.js'
 
-export type VerifyOptions<C extends string = string> = {
+export type VerifyOptions<C extends string = never> = {
   audience?: string | readonly string[]
   /** in seconds */
   clockTolerance?: number
@@ -14,9 +14,7 @@ export type VerifyOptions<C extends string = string> = {
   requiredClaims?: readonly C[]
 }
 
-export type VerifyPayload = Record<string, unknown>
-
-export type VerifyResult<P extends VerifyPayload, C extends string> = {
-  payload: RequiredKey<P & JwtPayload, C>
+export type VerifyResult<C extends string = never> = {
+  payload: RequiredKey<JwtPayload, C>
   protectedHeader: JwtHeader
 }