WRNexusJS/packages/oauth/src/advanced.ts

import { randomToken, type OAuthProvider, type OAuthTokens } from "./index.ts";
import { verifyJwtWithJwks, type JwtClaims, type RemoteJwks } from "@wrnexus/jwt";

export interface OAuthStateRecord {
  state: string;
  verifier: string;
  redirectUri: string;
  returnTo?: string;
  expiresAt: number;
}
export interface OAuthStateStore {
  set(record: OAuthStateRecord): Promise<void>;
  consume(state: string): Promise<OAuthStateRecord | null>;
}
export function memoryOAuthStateStore(now: () => number = Date.now): OAuthStateStore {
  const records = new Map<string, OAuthStateRecord>();
  return {
    async set(record) {
      records.set(record.state, record);
    },
    async consume(state) {
      const value = records.get(state);
      records.delete(state);
      if (!value || value.expiresAt <= now()) return null;
      return value;
    },
  };
}

export async function createOAuthState(
  store: OAuthStateStore,
  input: Omit<OAuthStateRecord, "state" | "expiresAt"> & { ttlMs?: number },
): Promise<OAuthStateRecord> {
  const record: OAuthStateRecord = {
    state: randomToken(),
    verifier: input.verifier,
    redirectUri: input.redirectUri,
    returnTo: input.returnTo,
    expiresAt: Date.now() + (input.ttlMs ?? 10 * 60 * 1000),
  };
  await store.set(record);
  return record;
}

export async function refreshOAuthTokens(
  provider: OAuthProvider,
  refreshToken: string,
  fetchImpl: typeof fetch = fetch,
): Promise<OAuthTokens> {
  const response = await fetchImpl(provider.tokenUrl, {
    method: "POST",
    headers: { "content-type": "application/x-www-form-urlencoded", accept: "application/json" },
    body: new URLSearchParams({
      grant_type: "refresh_token",
      refresh_token: refreshToken,
      client_id: provider.clientId,
      client_secret: provider.clientSecret,
    }),
  });
  if (!response.ok) throw new Error(`${provider.name} token refresh failed (${response.status})`);
  const tokens = (await response.json()) as OAuthTokens;
  if (!tokens.refresh_token) tokens.refresh_token = refreshToken;
  return tokens;
}

export interface OidcDiscovery {
  issuer: string;
  authorization_endpoint: string;
  token_endpoint: string;
  userinfo_endpoint?: string;
  jwks_uri: string;
  revocation_endpoint?: string;
}

function requireHttpsEndpoint(value: unknown, name: string): string {
  if (typeof value !== "string") throw new Error(`OIDC discovery is missing ${name}`);
  const url = new URL(value);
  if (url.protocol !== "https:" || url.username || url.password || url.hash) {
    throw new Error(`OIDC ${name} must be an HTTPS URL without credentials or a fragment`);
  }
  return value;
}

export async function discoverOidc(
  issuer: string,
  fetchImpl: typeof fetch = fetch,
): Promise<OidcDiscovery> {
  const base = issuer.replace(/\/$/, "");
  const response = await fetchImpl(`${base}/.well-known/openid-configuration`);
  if (!response.ok) throw new Error(`OIDC discovery failed (${response.status})`);
  const value = (await response.json()) as Partial<OidcDiscovery>;
  if (value.issuer !== base) throw new Error("OIDC issuer mismatch");
  return {
    issuer: base,
    authorization_endpoint: requireHttpsEndpoint(
      value.authorization_endpoint,
      "authorization_endpoint",
    ),
    token_endpoint: requireHttpsEndpoint(value.token_endpoint, "token_endpoint"),
    jwks_uri: requireHttpsEndpoint(value.jwks_uri, "jwks_uri"),
    userinfo_endpoint: value.userinfo_endpoint
      ? requireHttpsEndpoint(value.userinfo_endpoint, "userinfo_endpoint")
      : undefined,
    revocation_endpoint: value.revocation_endpoint
      ? requireHttpsEndpoint(value.revocation_endpoint, "revocation_endpoint")
      : undefined,
  };
}

export interface OidcIdTokenClaims extends JwtClaims {
  sub: string;
  iss: string;
  aud: string | string[];
  exp: number;
  iat: number;
  nonce?: string;
  azp?: string;
  at_hash?: string;
}

export interface VerifyOidcIdTokenOptions {
  issuer: string;
  clientId: string;
  jwks: RemoteJwks;
  nonce?: string;
  accessToken?: string;
  now?: number;
  clockTolerance?: number;
  maxAge?: number;
}

export function validateOidcClaims(
  claims: JwtClaims,
  options: Pick<VerifyOidcIdTokenOptions, "clientId" | "nonce">,
): asserts claims is OidcIdTokenClaims {
  if (typeof claims.sub !== "string" || !claims.sub) throw new Error("OIDC token has no subject");
  if (
    typeof claims.iss !== "string" ||
    (typeof claims.aud !== "string" && !Array.isArray(claims.aud)) ||
    typeof claims.exp !== "number" ||
    typeof claims.iat !== "number"
  ) {
    throw new Error("OIDC token is missing required claims");
  }
  const audiences = Array.isArray(claims.aud) ? claims.aud : [claims.aud];
  if (audiences.length > 1 && claims.azp !== options.clientId)
    throw new Error("OIDC token has invalid authorized party");
  if (claims.azp !== undefined && claims.azp !== options.clientId)
    throw new Error("OIDC token has invalid authorized party");
  if (options.nonce !== undefined && claims.nonce !== options.nonce)
    throw new Error("OIDC token has invalid nonce");
}

async function accessTokenHash(accessToken: string): Promise<string> {
  const digest = new Uint8Array(
    await crypto.subtle.digest("SHA-256", new TextEncoder().encode(accessToken)),
  ).slice(0, 16);
  let binary = "";
  for (const byte of digest) binary += String.fromCharCode(byte);
  return btoa(binary).replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/, "");
}

export async function verifyOidcIdToken(
  token: string,
  options: VerifyOidcIdTokenOptions,
): Promise<OidcIdTokenClaims> {
  const issuer = options.issuer.replace(/\/$/, "");
  const claims = await verifyJwtWithJwks(token, options.jwks, {
    issuer,
    audience: options.clientId,
    now: options.now,
    clockTolerance: options.clockTolerance,
    maxAge: options.maxAge,
  });
  validateOidcClaims(claims, options);
  if (options.accessToken !== undefined) {
    if (typeof claims.at_hash !== "string") throw new Error("OIDC token has no access-token hash");
    if ((await accessTokenHash(options.accessToken)) !== claims.at_hash)
      throw new Error("OIDC token has invalid access-token hash");
  }
  return claims;
}

export function validateOAuthReturnTo(
  value: string | undefined,
  origin: string,
  fallback = "/",
): string {
  if (!value) return fallback;
  try {
    const url = new URL(value, origin);
    return url.origin === new URL(origin).origin
      ? `${url.pathname}${url.search}${url.hash}`
      : fallback;
  } catch {
    return fallback;
  }
}