WRNexusJS/packages/jwt/src/jwks.ts

import { decodeJwt, type JwtClaims } from "./index.ts";
import { JwtError, type VerifyOptions } from "./index.ts";

export interface RemoteJwksOptions {
  fetch?: typeof fetch;
  cacheTtlMs?: number;
  maxKeys?: number;
  maxBytes?: number;
  now?: () => number;
}

export interface RemoteJwks {
  resolve(kid: string, alg: string): Promise<CryptoKey>;
  refresh(): Promise<void>;
  clear(): void;
  stats(): { fetches: number; hits: number; keys: number; expiresAt: number };
}

type JwksKey = JsonWebKey & { kid?: string; alg?: string; use?: string };
type StoredKey = { jwk: JwksKey; key?: Promise<CryptoKey> };

function decodeBase64Url(value: string): Uint8Array {
  const padding = value.length % 4 === 0 ? "" : "=".repeat(4 - (value.length % 4));
  const binary = atob(value.replace(/-/g, "+").replace(/_/g, "/") + padding);
  return Uint8Array.from(binary, (character) => character.charCodeAt(0));
}

export function createRemoteJwks(url: string, options: RemoteJwksOptions = {}): RemoteJwks {
  const endpoint = new URL(url);
  if (endpoint.protocol !== "https:") throw new TypeError("JWKS URL must use HTTPS");
  const fetchImpl = options.fetch ?? fetch;
  const cacheTtlMs = options.cacheTtlMs ?? 5 * 60_000;
  const maxKeys = options.maxKeys ?? 32;
  const maxBytes = options.maxBytes ?? 256 * 1024;
  if (!Number.isFinite(cacheTtlMs) || cacheTtlMs < 0)
    throw new RangeError("JWKS cacheTtlMs must be non-negative");
  if (!Number.isInteger(maxKeys) || maxKeys < 1)
    throw new RangeError("JWKS maxKeys must be positive");
  if (!Number.isInteger(maxBytes) || maxBytes < 1)
    throw new RangeError("JWKS maxBytes must be positive");
  const now = options.now ?? Date.now;
  const keys = new Map<string, StoredKey>();
  let expiresAt = 0;
  let etag: string | undefined;
  let refreshing: Promise<void> | undefined;
  let fetches = 0;
  let hits = 0;

  const refresh = async (): Promise<void> => {
    if (refreshing) return refreshing;
    refreshing = (async () => {
      fetches++;
      const response = await fetchImpl(endpoint, {
        headers: etag
          ? { accept: "application/json", "if-none-match": etag }
          : { accept: "application/json" },
      });
      if (response.status === 304) {
        expiresAt = now() + cacheTtlMs;
        return;
      }
      if (!response.ok) throw new JwtError(`JWKS fetch failed (${response.status})`);
      const declared = Number(response.headers.get("content-length"));
      if (Number.isFinite(declared) && declared > maxBytes)
        throw new JwtError("JWKS response is too large");
      const text = await response.text();
      if (new TextEncoder().encode(text).byteLength > maxBytes)
        throw new JwtError("JWKS response is too large");
      let document: { keys?: JwksKey[] };
      try {
        document = JSON.parse(text) as { keys?: JwksKey[] };
      } catch {
        throw new JwtError("JWKS response is invalid JSON");
      }
      if (!Array.isArray(document.keys) || document.keys.length > maxKeys)
        throw new JwtError("JWKS response has an invalid key set");
      const next = new Map<string, StoredKey>();
      for (const jwk of document.keys) {
        if (
          typeof jwk.kid !== "string" ||
          !jwk.kid ||
          jwk.kty !== "RSA" ||
          (jwk.use !== undefined && jwk.use !== "sig") ||
          (jwk.alg !== undefined && jwk.alg !== "RS256") ||
          typeof jwk.n !== "string" ||
          typeof jwk.e !== "string"
        ) {
          continue;
        }
        if (next.has(jwk.kid)) throw new JwtError(`JWKS contains duplicate kid: ${jwk.kid}`);
        next.set(jwk.kid, { jwk });
      }
      keys.clear();
      for (const [kid, key] of next) keys.set(kid, key);
      etag = response.headers.get("etag") ?? undefined;
      expiresAt = now() + cacheTtlMs;
    })().finally(() => {
      refreshing = undefined;
    });
    return refreshing;
  };

  return {
    async resolve(kid, alg) {
      if (!kid) throw new JwtError("Token has no key id");
      if (alg !== "RS256") throw new JwtError(`Unsupported token algorithm: ${alg}`);
      if (now() >= expiresAt) await refresh();
      let stored = keys.get(kid);
      if (!stored) {
        await refresh();
        stored = keys.get(kid);
      }
      if (!stored) throw new JwtError(`Unknown key id: ${kid}`);
      hits++;
      stored.key ??= crypto.subtle.importKey(
        "jwk",
        stored.jwk,
        { name: "RSASSA-PKCS1-v1_5", hash: "SHA-256" },
        false,
        ["verify"],
      );
      return stored.key;
    },
    refresh,
    clear() {
      keys.clear();
      expiresAt = 0;
      etag = undefined;
    },
    stats: () => ({ fetches, hits, keys: keys.size, expiresAt }),
  };
}

export async function verifyJwtWithJwks<T extends JwtClaims = JwtClaims>(
  token: string,
  jwks: RemoteJwks,
  options: VerifyOptions = {},
): Promise<T> {
  const parts = token.split(".");
  if (parts.length !== 3) throw new JwtError("Malformed token");
  const decoded = decodeJwt(token);
  const kid = typeof decoded.header.kid === "string" ? decoded.header.kid : "";
  const alg = typeof decoded.header.alg === "string" ? decoded.header.alg : "";
  const key = await jwks.resolve(kid, alg);
  const valid = await crypto.subtle.verify(
    "RSASSA-PKCS1-v1_5",
    key,
    decodeBase64Url(parts[2]!) as BufferSource,
    new TextEncoder().encode(`${parts[0]}.${parts[1]}`) as BufferSource,
  );
  if (!valid) throw new JwtError("Invalid signature");
  const claims = decoded.claims as T;
  const now = options.now ?? Math.floor(Date.now() / 1000);
  const tolerance = Math.max(0, options.clockTolerance ?? 0);
  if (typeof claims.exp === "number" && now - tolerance >= claims.exp)
    throw new JwtError("Token expired");
  if (typeof claims.nbf === "number" && now + tolerance < claims.nbf)
    throw new JwtError("Token not yet valid");
  if (options.issuer !== undefined && claims.iss !== options.issuer)
    throw new JwtError("Invalid issuer");
  if (options.audience !== undefined) {
    const expected = Array.isArray(options.audience) ? options.audience : [options.audience];
    const actual = Array.isArray(claims.aud)
      ? claims.aud
      : typeof claims.aud === "string"
        ? [claims.aud]
        : [];
    if (!expected.some((audience) => actual.includes(audience)))
      throw new JwtError("Invalid audience");
  }
  if (
    options.maxAge !== undefined &&
    typeof claims.iat === "number" &&
    now - claims.iat > options.maxAge + tolerance
  ) {
    throw new JwtError("Token is too old");
  }
  return claims;
}