Design a JWT validator

import type { KeyObject } from "node:crypto";

export type Alg = "RS256" | "ES256" | "EdDSA";

export interface Jwk { kid: string; alg: Alg; use?: "sig" | "enc"; key: KeyObject }
/** Fetches and parses the issuer's JWKS. Parse with createPublicKey({ key, format: "jwk" });
 *  drop keys without kid or alg and any key carrying private members (d). */
export interface KeySource { fetch(issuer: string): Promise<readonly Jwk[]> }
export interface KeyLookup { get(kid: string): Promise<Jwk> }      // throws JwtError
export interface Clock { nowSec(): number }

export interface ValidatorConfig {
  issuer: string;
  audience: string;
  algs: readonly Alg[];        // pinned: the header may only pick from this list
  skewSec: number;             // 0..300
  maxAgeSec?: number;          // if set, iat is required and must be recent
  requiredClaims?: readonly string[];
}

export type JwtErrorCode =
  | "malformed" | "unsupported_alg" | "unknown_kid" | "keys_unavailable" | "bad_signature"
  | "expired" | "not_yet_valid" | "issued_in_future" | "too_old"
  | "wrong_issuer" | "wrong_audience" | "missing_claim";

export class JwtError extends Error {
  constructor(readonly code: JwtErrorCode, message: string) { super(message); this.name = "JwtError"; }
}

export interface Claims { iss: string; sub: string; aud: string | string[]; exp: number; nbf?: number; iat?: number; jti?: string; [k: string]: unknown }

export type Result = { ok: true; claims: Claims; kid: string } | { ok: false; error: JwtError };
export interface JwtValidator { validate(token: string): Promise<Result> }

import { verify as cryptoVerify, type KeyObject } from "node:crypto";

const B64URL = /^[A-Za-z0-9_-]+$/;
const KEY_TYPE: Record<Alg, string> = { RS256: "rsa", ES256: "ec", EdDSA: "ed25519" };

const isObject = (v: unknown): v is Record<string, unknown> => typeof v === "object" && v !== null && !Array.isArray(v);
const isNumericDate = (v: unknown): v is number => typeof v === "number" && Number.isFinite(v) && v >= 0;

function decodeJson(part: string): unknown {
  if (!B64URL.test(part)) throw new JwtError("malformed", "segment is not base64url");
  try { return JSON.parse(Buffer.from(part, "base64url").toString("utf8")); }
  catch { throw new JwtError("malformed", "segment is not JSON"); }
}

function verifySignature(alg: Alg, key: KeyObject, data: Buffer, sig: Buffer): boolean {
  if (key.asymmetricKeyType !== KEY_TYPE[alg]) return false;      // an RSA key can never verify an EdDSA token
  try {
    switch (alg) {
      case "RS256": return cryptoVerify("sha256", data, key, sig);
      case "ES256": return sig.length === 64 && cryptoVerify("sha256", data, { key, dsaEncoding: "ieee-p1363" }, sig);
      case "EdDSA": return cryptoVerify(null, data, key, sig);
    }
  } catch { return false; }
}

export class Validator implements JwtValidator {
  constructor(private readonly cfg: ValidatorConfig, private readonly keys: KeyLookup, private readonly clock: Clock) {
    if (!cfg.issuer || !cfg.audience) throw new Error("issuer and audience are required");
    if (cfg.algs.length === 0) throw new Error("pin at least one alg");
    if (!Number.isInteger(cfg.skewSec) || cfg.skewSec < 0 || cfg.skewSec > 300) throw new Error("skewSec must be an integer in [0, 300]");
    if (cfg.maxAgeSec !== undefined && !(cfg.maxAgeSec > 0)) throw new Error("maxAgeSec must be positive");
  }

  async validate(token: string): Promise<Result> {
    try { return { ok: true, ...(await this.check(token)) }; }
    catch (e) {
      if (e instanceof JwtError) return { ok: false, error: e };
      throw e;                                                   // a bug, not a bad token: do not mask it as a 401
    }
  }

  private async check(token: string): Promise<{ claims: Claims; kid: string }> {
    if (token.length > 8192) throw new JwtError("malformed", "token too large");
    const parts = token.split(".");
    if (parts.length !== 3) throw new JwtError("malformed", "expected three segments");
    const [h64, p64, s64] = parts;

    const header = decodeJson(h64);
    if (!isObject(header)) throw new JwtError("malformed", "header is not an object");
    const alg = header["alg"];
    if (typeof alg !== "string" || !this.cfg.algs.some((a) => a === alg)) throw new JwtError("unsupported_alg", "alg not allowed: " + String(alg));
    if ("crit" in header) throw new JwtError("malformed", "crit extensions are not supported");
    const kid = header["kid"];
    if (typeof kid !== "string" || kid.length === 0 || kid.length > 128) throw new JwtError("malformed", "kid missing or invalid");

    const jwk = await this.keys.get(kid);
    if (jwk.alg !== alg) throw new JwtError("unsupported_alg", "key " + kid + " is not for " + alg);
    if (!B64URL.test(s64)) throw new JwtError("malformed", "signature is not base64url");
    const ok = verifySignature(jwk.alg, jwk.key, Buffer.from(h64 + "." + p64), Buffer.from(s64, "base64url"));
    if (!ok) throw new JwtError("bad_signature", "signature does not verify");

    return { claims: this.checkClaims(decodeJson(p64)), kid };   // claims are read only after the signature holds
  }

  private checkClaims(p: unknown): Claims {
    if (!isObject(p)) throw new JwtError("malformed", "payload is not an object");
    const now = this.clock.nowSec();
    const skew = this.cfg.skewSec;

    if (p["iss"] !== this.cfg.issuer) throw new JwtError("wrong_issuer", "unexpected issuer");
    if (typeof p["sub"] !== "string" || p["sub"] === "") throw new JwtError("missing_claim", "sub");
    const aud = p["aud"];
    const auds = typeof aud === "string" ? [aud] : Array.isArray(aud) && aud.every((a) => typeof a === "string") ? aud : null;
    if (!auds) throw new JwtError("missing_claim", "aud");
    if (!auds.includes(this.cfg.audience)) throw new JwtError("wrong_audience", "token is for another audience");

    const exp = p["exp"];
    if (!isNumericDate(exp)) throw new JwtError("missing_claim", "exp");
    if (now - skew >= exp) throw new JwtError("expired", "expired");

    const nbf = p["nbf"];
    if (nbf !== undefined) {
      if (!isNumericDate(nbf)) throw new JwtError("malformed", "nbf is not a NumericDate");
      if (now + skew < nbf) throw new JwtError("not_yet_valid", "not valid yet");
    }
    const iat = p["iat"];
    if (iat !== undefined) {
      if (!isNumericDate(iat)) throw new JwtError("malformed", "iat is not a NumericDate");
      if (iat > now + skew) throw new JwtError("issued_in_future", "iat in the future");
      if (this.cfg.maxAgeSec !== undefined && now - iat > this.cfg.maxAgeSec + skew) throw new JwtError("too_old", "token older than maxAgeSec");
    } else if (this.cfg.maxAgeSec !== undefined) {
      throw new JwtError("missing_claim", "iat");
    }
    for (const c of this.cfg.requiredClaims ?? []) if (!(c in p)) throw new JwtError("missing_claim", c);
    return p as Claims;
  }
}

export class JwksCache implements KeyLookup {
  private entry?: { keys: Map<string, Jwk>; fetchedAt: number };
  private inflight?: Promise<void>;
  private lastAttemptAt = -Infinity;

  constructor(private readonly src: KeySource, private readonly issuer: string, private readonly clock: Clock,
              private readonly ttlSec = 600, private readonly minRefreshSec = 30, private readonly maxStaleSec = 86_400) {
    if (!(minRefreshSec > 0 && ttlSec >= minRefreshSec && maxStaleSec >= ttlSec)) throw new Error("need 0 < minRefresh <= ttl <= maxStale");
  }

  async get(kid: string): Promise<Jwk> {
    const now = this.clock.nowSec();
    const due = !this.entry || now - this.entry.fetchedAt >= this.ttlSec;
    if (due && this.mayRefresh(now)) await this.refresh();
    if (!this.entry || now - this.entry.fetchedAt >= this.maxStaleSec) throw new JwtError("keys_unavailable", "no fresh enough JWKS");

    let k = this.entry.keys.get(kid);
    if (!k && this.mayRefresh(now)) {                    // maybe rotated in; forged kids are rate limited here
      await this.refresh();
      k = this.entry.keys.get(kid);
    }
    if (!k) throw new JwtError("unknown_kid", "no key for kid " + kid);
    return k;
  }

  /** Join a fetch already in flight (concurrent first requests must wait, not fail), or start one if the rate limit allows. */
  private mayRefresh(now: number) { return this.inflight !== undefined || now - this.lastAttemptAt >= this.minRefreshSec; }

  /** Single flight: concurrent callers share one fetch. Never rejects. */
  private refresh(): Promise<void> {
    this.inflight ??= this.load().finally(() => { this.inflight = undefined; });
    return this.inflight;
  }

  private async load(): Promise<void> {
    const now = this.clock.nowSec();
    this.lastAttemptAt = now;
    try {
      const keys = await this.src.fetch(this.issuer);
      this.entry = { keys: new Map(keys.filter((k) => k.use !== "enc").map((k) => [k.kid, k])), fetchedAt: now };
    } catch {
      // keep the previous entry; get() refuses it once it is older than maxStaleSec. Emit a metric here.
    }
  }
}

import { describe, it, expect, vi } from "vitest";
import { generateKeyPairSync, sign } from "node:crypto";

const { publicKey, privateKey } = generateKeyPairSync("ed25519");
const b64 = (o: object) => Buffer.from(JSON.stringify(o)).toString("base64url");
const mint = (payload: object, header: object = { alg: "EdDSA", kid: "k1", typ: "JWT" }) => {
  const input = b64(header) + "." + b64(payload);
  return input + "." + sign(null, Buffer.from(input), privateKey).toString("base64url");
};

const T0 = 1_800_000_000;
let now = T0;
const clock = { nowSec: () => now };
const ISS = "https://idp.example";
const good = { iss: ISS, sub: "u1", aud: "academy-api", iat: T0, exp: T0 + 300 };
const code = (r: Result) => (r.ok ? "ok" : r.error.code);

function setup() {
  const fetchKeys = vi.fn(async () => [{ kid: "k1", alg: "EdDSA" as const, key: publicKey }]);
  const keys = new JwksCache({ fetch: fetchKeys }, ISS, clock);
  const v = new Validator({ issuer: ISS, audience: "academy-api", algs: ["EdDSA"], skewSec: 60 }, keys, clock);
  return { v, fetchKeys };
}

describe("Validator", () => {
  it("accepts a well-formed token", async () => {
    now = T0;
    expect(code(await setup().v.validate(mint(good)))).toBe("ok");
  });
  it("rejects alg none even with a valid payload", async () => {
    now = T0;
    const unsigned = b64({ alg: "none", kid: "k1" }) + "." + b64(good) + ".";
    expect(code(await setup().v.validate(unsigned))).toBe("unsupported_alg");
  });
  it("honours skew at the exp boundary", async () => {
    const { v } = setup();
    now = good.exp + 59;
    expect(code(await v.validate(mint(good)))).toBe("ok");
    now = good.exp + 60;
    expect(code(await v.validate(mint(good)))).toBe("expired");
  });
  it("rejects a token for another audience", async () => {
    now = T0;
    expect(code(await setup().v.validate(mint({ ...good, aud: ["billing"] })))).toBe("wrong_audience");
  });
  it("a burst of unknown kids costs one fetch", async () => {
    now = T0;
    const { v, fetchKeys } = setup();
    const forged = mint(good, { alg: "EdDSA", kid: "nope" });
    const results = await Promise.all(Array.from({ length: 50 }, () => v.validate(forged)));
    expect(results.every((r) => code(r) === "unknown_kid")).toBe(true);
    expect(fetchKeys).toHaveBeenCalledTimes(1);
  });
});