WRNexusJS/packages/jwt/src/helpers.ts

import type { Context, Middleware } from "@wrnexus/core";
import {
  JwtError,
  signJwt,
  verifyJwt,
  type JwtClaims,
  type SignOptions,
  type VerifyOptions,
} from "./index.ts";

export interface AccessTokenClaims extends JwtClaims {
  sub: string;
  type: "access";
  scopes?: string[];
}

export interface RefreshTokenClaims extends JwtClaims {
  sub: string;
  type: "refresh";
  family?: string;
}

export function extractBearerToken(
  value: Headers | Request | Context | string | null | undefined,
): string | undefined {
  const header =
    typeof value === "string" || value == null
      ? (value ?? "")
      : "req" in value
        ? (value.req.headers.get("authorization") ?? "")
        : value instanceof Request
          ? (value.headers.get("authorization") ?? "")
          : (value.get("authorization") ?? "");
  return /^Bearer\s+(.+)$/i.exec(header)?.[1];
}

export async function tryVerifyJwt<T extends JwtClaims = JwtClaims>(
  token: string | undefined,
  secret: string,
  options: VerifyOptions = {},
): Promise<T | null> {
  if (!token) return null;
  try {
    return await verifyJwt<T>(token, secret, options);
  } catch {
    return null;
  }
}

export function assertJwtClaims<T extends JwtClaims>(
  claims: T,
  requirements: {
    subject?: boolean;
    type?: string;
    required?: string[];
  } = {},
): T {
  if (requirements.subject && !claims.sub) throw new JwtError("Token subject is required");
  if (requirements.type && claims.type !== requirements.type)
    throw new JwtError("Invalid token type");
  for (const name of requirements.required ?? []) {
    if (!(name in claims)) throw new JwtError(`Missing required claim: ${name}`);
  }
  return claims;
}

export function tokenScopes(claims: JwtClaims): string[] {
  const value = claims.scopes ?? claims.scope;
  if (Array.isArray(value))
    return value.filter((entry): entry is string => typeof entry === "string");
  if (typeof value === "string") return value.split(/\s+/).filter(Boolean);
  return [];
}

export function hasScopes(
  claims: JwtClaims,
  required: readonly string[],
  mode: "all" | "any" = "all",
): boolean {
  const scopes = new Set(tokenScopes(claims));
  return mode === "all"
    ? required.every((scope) => scopes.has(scope))
    : required.some((scope) => scopes.has(scope));
}

export function requireScopes(
  required: readonly string[],
  mode: "all" | "any" = "all",
): Middleware {
  return async (ctx, next) => {
    const claims = ctx.user && typeof ctx.user === "object" ? (ctx.user as JwtClaims) : {};
    if (!hasScopes(claims, required, mode)) {
      return Response.json({ ok: false, error: "Insufficient scope" }, { status: 403 });
    }
    return next();
  };
}

export function createAccessToken(
  subject: string,
  secret: string,
  options: Omit<SignOptions, "expiresIn"> & {
    expiresIn?: number;
    scopes?: string[];
    claims?: JwtClaims;
  } = {},
): Promise<string> {
  const { claims, scopes, expiresIn, ...signOptions } = options;
  return signJwt(
    { ...claims, sub: subject, type: "access", ...(scopes ? { scopes } : {}) },
    secret,
    { ...signOptions, expiresIn: expiresIn ?? 15 * 60 },
  );
}

export function createRefreshToken(
  subject: string,
  secret: string,
  options: Omit<SignOptions, "expiresIn"> & {
    expiresIn?: number;
    family?: string;
    claims?: JwtClaims;
  } = {},
): Promise<string> {
  const { claims, family, expiresIn, ...signOptions } = options;
  return signJwt(
    { ...claims, sub: subject, type: "refresh", ...(family ? { family } : {}) },
    secret,
    { ...signOptions, expiresIn: expiresIn ?? 30 * 24 * 60 * 60 },
  );
}

export async function verifyAccessToken(
  token: string,
  secret: string,
  options: VerifyOptions = {},
): Promise<AccessTokenClaims> {
  return assertJwtClaims(await verifyJwt<AccessTokenClaims>(token, secret, options), {
    subject: true,
    type: "access",
  });
}

export async function verifyRefreshToken(
  token: string,
  secret: string,
  options: VerifyOptions = {},
): Promise<RefreshTokenClaims> {
  return assertJwtClaims(await verifyJwt<RefreshTokenClaims>(token, secret, options), {
    subject: true,
    type: "refresh",
  });
}

const COOKIE_NAME = /^[!#$%&'*+.^_`|~0-9A-Za-z-]+$/;

function cookieSource(value: Headers | Request | string | null | undefined): string {
  if (typeof value === "string" || value == null) return value ?? "";
  return value instanceof Request
    ? (value.headers.get("cookie") ?? "")
    : (value.get("cookie") ?? "");
}

export function readJwtCookie(
  value: Headers | Request | string | null | undefined,
  name = "__Host-wrn_token",
): string | undefined {
  if (!COOKIE_NAME.test(name)) throw new TypeError("Invalid JWT cookie name.");
  for (const part of cookieSource(value).split(";")) {
    const index = part.indexOf("=");
    if (index < 0) continue;
    const key = part.slice(0, index).trim();
    if (key !== name) continue;
    try {
      return decodeURIComponent(part.slice(index + 1).trim());
    } catch {
      return undefined;
    }
  }
  return undefined;
}

export function jwtCookie(
  token: string,
  options: {
    name?: string;
    maxAge?: number;
    secure?: boolean;
    sameSite?: "Strict" | "Lax" | "None";
    path?: string;
  } = {},
): string {
  const name = options.name ?? "__Host-wrn_token";
  const path = options.path ?? "/";
  const sameSite = options.sameSite ?? "Lax";
  const secure = options.secure !== false;
  if (!COOKIE_NAME.test(name)) throw new TypeError("Invalid JWT cookie name.");
  if (!path.startsWith("/") || /[;\r\n]/.test(path))
    throw new TypeError("Invalid JWT cookie path.");
  if (name.startsWith("__Host-") && (path !== "/" || !secure)) {
    throw new TypeError("__Host- JWT cookies require Path=/ and Secure.");
  }
  if (sameSite === "None" && !secure) {
    throw new TypeError("SameSite=None JWT cookies require Secure.");
  }
  const parts = [
    `${name}=${encodeURIComponent(token)}`,
    `Path=${path}`,
    "HttpOnly",
    `SameSite=${sameSite}`,
  ];
  if (secure) parts.push("Secure");
  if (options.maxAge !== undefined) {
    if (!Number.isFinite(options.maxAge)) throw new RangeError("JWT cookie maxAge must be finite.");
    parts.push(`Max-Age=${Math.max(0, Math.floor(options.maxAge))}`);
  }
  return parts.join("; ");
}

export function clearJwtCookie(
  options: Omit<Parameters<typeof jwtCookie>[1], "maxAge"> = {},
): string {
  return jwtCookie("", { ...options, maxAge: 0 });
}

export interface JwtTokenPair {
  accessToken: string;
  refreshToken: string;
  tokenType: "Bearer";
  expiresIn: number;
}

export async function createTokenPair(
  subject: string,
  input: {
    accessSecret: string;
    refreshSecret?: string;
    accessExpiresIn?: number;
    refreshExpiresIn?: number;
    scopes?: string[];
    family?: string;
    accessOptions?: Omit<SignOptions, "expiresIn">;
    refreshOptions?: Omit<SignOptions, "expiresIn">;
  },
): Promise<JwtTokenPair> {
  const expiresIn = input.accessExpiresIn ?? 15 * 60;
  const [accessToken, refreshToken] = await Promise.all([
    createAccessToken(subject, input.accessSecret, {
      ...input.accessOptions,
      expiresIn,
      scopes: input.scopes,
    }),
    createRefreshToken(subject, input.refreshSecret ?? input.accessSecret, {
      ...input.refreshOptions,
      expiresIn: input.refreshExpiresIn ?? 30 * 24 * 60 * 60,
      family: input.family,
    }),
  ]);
  return { accessToken, refreshToken, tokenType: "Bearer", expiresIn };
}

export function jwtResponse(
  accessToken: string,
  input: { refreshToken?: string; expiresIn?: number; tokenType?: string; scope?: string[] } = {},
): Response {
  return Response.json(
    {
      accessToken,
      tokenType: input.tokenType ?? "Bearer",
      expiresIn: input.expiresIn ?? 900,
      ...(input.refreshToken ? { refreshToken: input.refreshToken } : {}),
      ...(input.scope ? { scope: input.scope.join(" ") } : {}),
    },
    { headers: { "cache-control": "no-store", pragma: "no-cache" } },
  );
}