Improve type safety and compatibility with Bun (#2879)

* jwk: Improve type safety and compatibility with Bun
* improve type safety of jwk keys
* improve typing of verifyAccessToken
* update @types/http-errors
* Better report invalid content-encoding errors
* Mark jwk key fields as readonly

.changeset/early-tomatoes-cheer.md

@@ -0,0 +1,5 @@
+---
+"@atproto/jwk": patch
+---
+
+Mark jwk key fields as readonly

.changeset/eight-eyes-float.md

@@ -0,0 +1,5 @@
+---
+"@atproto/jwk-jose": patch
+---
+
+Improve compatibility with runtimes relying on webcrypto (by 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 (passthrough) unknown properties in JWT payload & headers

.changeset/thin-mice-lick.md

@@ -0,0 +1,5 @@
+---
+"@atproto/jwk": patch
+---
+
+Expose ValidationError to allow for proper error handling

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
@@ -35,53 +37,97 @@ export type Importable = string | KeyLike | Jwk
 
 export type { GenerateKeyPairOptions, GenerateKeyPairResult }
 
-export class JoseKey extends Key {
-  #keyObj?: KeyLike | Uint8Array
-
-  protected async getKey() {
+export class JoseKey<J extends Jwk = Jwk> extends Key<J> {
+  /**
+   * Some runtimes (e.g. Bun) require an `alg` second argument to be set when
+   * invoking `importJWK`. In order to be compatible with these runtimes, we
+   * provide the following method to ensure the `alg` is always set. We also
+   * take the opportunity to ensure that the `alg` is compatible with this key.
+   */
+  protected async getKeyObj(alg: string) {
+    if (!this.algorithms.includes(alg)) {
+      throw new JwkError(`Key cannot be used with algorithm "${alg}"`)
+    }
     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) {
+        throw new JwtCreateError('Missing "alg" in JWT header')
+      }
+
+      const keyObj = await this.getKeyObj(alg)
+      const jwtBuilder = new SignJWT(payload).setProtectedHeader({
+        ...header,
+        alg,
+        kid: this.kid,
+      })
+
+      const signedJwt = await jwtBuilder.sign(keyObj)
 
-    const keyObj = await this.getKey()
-    return new SignJWT(payload)
-      .setProtectedHeader({ ...header, kid: this.kid })
-      .sign(keyObj) as Promise<SignedJwt>
+      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/package.json

@@ -25,7 +25,8 @@
   },
   "dependencies": {
     "@atproto/jwk": "workspace:*",
-    "@atproto/jwk-jose": "workspace:*"
+    "@atproto/jwk-jose": "workspace:*",
+    "zod": "^3.23.8"
   },
   "devDependencies": {
     "typescript": "^5.6.3"

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

@@ -1,9 +1,20 @@
-import { Jwk, jwkSchema } from '@atproto/jwk'
+import { JwkError, jwkSchema } from '@atproto/jwk'
 import { GenerateKeyPairOptions, JoseKey } from '@atproto/jwk-jose'
+import z from 'zod'
 
 import { fromSubtleAlgorithm, isCryptoKeyPair } from './util.js'
 
-export class WebcryptoKey extends JoseKey {
+// Webcrypto keys are bound to a single algorithm
+export const jwkWithAlgSchema = z.intersection(
+  jwkSchema,
+  z.object({ alg: z.string() }),
+)
+
+export type JwkWithAlg = z.infer<typeof jwkWithAlgSchema>
+
+export class WebcryptoKey<
+  J extends JwkWithAlg = JwkWithAlg,
+> extends JoseKey<J> {
   // We need to override the static method generate from JoseKey because
   // the browser needs both the private and public keys
   static override async generate(
@@ -26,29 +37,35 @@ export class WebcryptoKey extends JoseKey {
     // > The "use" and "key_ops" JWK members SHOULD NOT be used together; [...]
     // > Applications should specify which of these members they use.
 
-    const { key_ops: _, ...jwk } = await crypto.subtle.exportKey(
+    const {
+      key_ops,
+      use,
+      alg = fromSubtleAlgorithm(cryptoKeyPair.privateKey.algorithm),
+      ...jwk
+    } = await crypto.subtle.exportKey(
       'jwk',
       cryptoKeyPair.privateKey.extractable
         ? cryptoKeyPair.privateKey
         : cryptoKeyPair.publicKey,
     )
 
-    const use = jwk.use ?? 'sig'
-    const alg =
-      jwk.alg ?? fromSubtleAlgorithm(cryptoKeyPair.privateKey.algorithm)
+    if (use && use !== 'sig') {
+      throw new TypeError(`Unsupported JWK use "${use}"`)
+    }
 
-    if (use !== 'sig') {
-      throw new TypeError('Unsupported JWK use')
+    if (key_ops && !key_ops.some((o) => o === 'sign' || o === 'verify')) {
+      // Make sure that "key_ops", if present, is compatible with "use"
+      throw new TypeError(`Invalid key_ops "${key_ops}" for "sig" use`)
     }
 
     return new WebcryptoKey(
-      jwkSchema.parse({ ...jwk, use, kid, alg }),
+      jwkWithAlgSchema.parse({ ...jwk, kid, alg, use: 'sig' }),
       cryptoKeyPair,
     )
   }
 
   constructor(
-    jwk: Jwk,
+    jwk: Readonly<J>,
     readonly cryptoKeyPair: CryptoKeyPair,
   ) {
     super(jwk)
@@ -58,12 +75,15 @@ export class WebcryptoKey extends JoseKey {
     return true
   }
 
-  get privateJwk(): Jwk | undefined {
+  get privateJwk(): Readonly<J> | undefined {
     if (super.isPrivate) return this.jwk
     throw new Error('Private Webcrypto Key not exportable')
   }
 
-  protected override async getKey() {
+  protected override async getKeyObj(alg: string) {
+    if (this.jwk.alg !== alg) {
+      throw new JwkError(`Key cannot be used with algorithm "${alg}"`)
+    }
     return this.cryptoKeyPair.privateKey
   }
 }

packages/oauth/jwk/src/index.ts

@@ -1,3 +1,8 @@
+// Since we expose zod schemas, let's expose ZodError (under a generic name) so
+// that dependents can catch schema parsing errors without requiring an explicit
+// dependency on zod, or risking a conflict in case of mismatching zob versions.
+export { ZodError as ValidationError } 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
 }