NianAIGC/lib/server/auth/jwt.ts

import { createPublicKey, createVerify } from "node:crypto";
import type { JsonWebKey as CryptoJsonWebKey, KeyObject } from "node:crypto";
import { getAuthRuntimeConfig, type AuthRuntimeConfig } from "@/lib/auth/config";
import type { AuthSession, AuthUser } from "@/lib/auth/session";

export type AuthTokenClaims = {
  iss?: string;
  sub?: string;
  aud?: string | string[];
  exp?: number;
  iat?: number;
  nbf?: number;
  jti?: string;
  scope?: string | string[];
  client_id?: string;
  clientId?: string;
  user_id?: string | number;
  username?: string;
  tenant_id?: string | number;
  dept_id?: string | number;
  authorities?: string[] | string;
  [claim: string]: unknown;
};

type JwtHeader = {
  alg?: string;
  kid?: string;
  typ?: string;
  [key: string]: unknown;
};

type Jwks = {
  keys?: JwksKey[];
};

type JwksKey = CryptoJsonWebKey & {
  kid?: string;
  alg?: string;
  use?: string;
};

let jwksCache: {
  url: string;
  fetchedAt: number;
  keys: JwksKey[];
} | null = null;

export class JwtVerificationError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "JwtVerificationError";
  }
}

export async function verifyAuthJwt(token: string, config = getAuthRuntimeConfig()): Promise<AuthTokenClaims> {
  if (!config.jwksUrl) throw new JwtVerificationError("JWKS URL is not configured.");
  const parts = token.split(".");
  if (parts.length !== 3) throw new JwtVerificationError("Invalid JWT format.");
  const header = parseJwtPart<JwtHeader>(parts[0]);
  const claims = parseJwtPart<AuthTokenClaims>(parts[1]);

  if (header.alg !== "RS256") throw new JwtVerificationError("Unsupported JWT algorithm.");
  const publicKey = await getPublicKeyForHeader(header, config.jwksUrl);
  const verifier = createVerify("RSA-SHA256");
  verifier.update(`${parts[0]}.${parts[1]}`);
  verifier.end();
  if (!verifier.verify(publicKey, base64UrlToBuffer(parts[2]))) {
    throw new JwtVerificationError("JWT signature verification failed.");
  }

  validateClaims(claims, config);
  return claims;
}

export function createSessionFromClaims(
  claims: AuthTokenClaims,
  config: AuthRuntimeConfig,
  tokenResponseExpiresIn?: number
): AuthSession {
  const now = Math.floor(Date.now() / 1000);
  const jwtExpiry = numberClaim(claims.exp);
  const responseExpiry = tokenResponseExpiresIn ? now + tokenResponseExpiresIn : undefined;
  const expiresAt = Math.min(jwtExpiry || responseExpiry || now, responseExpiry || jwtExpiry || now);
  return {
    version: 1,
    user: userFromClaims(claims, config),
    issuedAt: now,
    expiresAt
  };
}

export function userFromClaims(claims: AuthTokenClaims, config = getAuthRuntimeConfig()): AuthUser {
  const clientId = stringClaim(claims.client_id) || stringClaim(claims.clientId) || config.clientId;
  const subject = stringClaim(claims.sub) || stringClaim(claims.username) || stringClaim(claims.user_id) || "unknown";
  const principalId = stringClaim(claims.user_id) || subject;
  const username = stringClaim(claims.username) || stringClaim(claims.sub);
  return {
    id: `auth:${sanitizeOwnerPart(clientId)}:${sanitizeOwnerPart(principalId)}`,
    subject,
    username,
    displayName: username || `用户 ${principalId}`,
    clientId,
    tenantId: stringClaim(claims.tenant_id),
    authorities: stringListClaim(claims.authorities),
    scope: stringListClaim(claims.scope)
  };
}

export function clearJwksCacheForTests() {
  jwksCache = null;
}

async function getPublicKeyForHeader(header: JwtHeader, jwksUrl: string): Promise<KeyObject> {
  const keys = await fetchJwksKeys(jwksUrl);
  const key = keys.find((item) => {
    if (item.kty !== "RSA") return false;
    if (!header.kid) return true;
    return item.kid === header.kid;
  });
  if (!key) throw new JwtVerificationError("JWT key id was not found in JWKS.");
  return createPublicKey({ key, format: "jwk" });
}

async function fetchJwksKeys(jwksUrl: string): Promise<JwksKey[]> {
  const now = Date.now();
  if (jwksCache?.url === jwksUrl && now - jwksCache.fetchedAt < 5 * 60 * 1000) return jwksCache.keys;
  const response = await fetch(jwksUrl, { cache: "no-store" });
  if (!response.ok) throw new JwtVerificationError(`JWKS request failed: ${response.status}`);
  const payload = await response.json() as Jwks;
  const keys = Array.isArray(payload.keys) ? payload.keys : [];
  jwksCache = { url: jwksUrl, fetchedAt: now, keys };
  return keys;
}

function validateClaims(claims: AuthTokenClaims, config: AuthRuntimeConfig) {
  const now = Math.floor(Date.now() / 1000);
  const skew = config.clockSkewSeconds;
  const exp = numberClaim(claims.exp);
  if (!exp || exp <= now - skew) throw new JwtVerificationError("JWT has expired.");
  const nbf = numberClaim(claims.nbf);
  if (nbf && nbf > now + skew) throw new JwtVerificationError("JWT is not active yet.");
  const iat = numberClaim(claims.iat);
  if (iat && iat > now + skew) throw new JwtVerificationError("JWT issued-at is in the future.");
  if (claims.iss !== config.issuer) throw new JwtVerificationError("JWT issuer is not trusted.");
  const clientId = stringClaim(claims.client_id) || stringClaim(claims.clientId);
  if (clientId && clientId !== config.clientId) throw new JwtVerificationError("JWT client id is not allowed.");
  const requiredScopes = config.scope.split(/\s+/).filter(Boolean);
  if (requiredScopes.length) {
    const tokenScopes = new Set(stringListClaim(claims.scope));
    if (tokenScopes.size > 0) {
      for (const scope of requiredScopes) {
        if (!tokenScopes.has(scope)) throw new JwtVerificationError("JWT scope is not allowed.");
      }
    }
  }
}

function parseJwtPart<T>(value: string): T {
  try {
    return JSON.parse(base64UrlToBuffer(value).toString("utf8")) as T;
  } catch {
    throw new JwtVerificationError("Invalid JWT JSON.");
  }
}

function base64UrlToBuffer(value: string): Buffer {
  return Buffer.from(value.replace(/-/g, "+").replace(/_/g, "/"), "base64");
}

function numberClaim(value: unknown): number | undefined {
  if (typeof value === "number" && Number.isFinite(value)) return value;
  if (typeof value === "string" && value.trim()) {
    const parsed = Number(value);
    if (Number.isFinite(parsed)) return parsed;
  }
  return undefined;
}

function stringClaim(value: unknown): string | undefined {
  if (typeof value === "string" && value.trim()) return value.trim();
  if (typeof value === "number" && Number.isFinite(value)) return String(value);
  return undefined;
}

function stringListClaim(value: unknown): string[] {
  if (Array.isArray(value)) return value.map(stringClaim).filter((item): item is string => Boolean(item));
  if (typeof value === "string") return value.split(/\s+/).map((item) => item.trim()).filter(Boolean);
  return [];
}

function sanitizeOwnerPart(value: string): string {
  return value.replace(/[^A-Za-z0-9_.:@-]+/g, "_").slice(0, 96) || "unknown";
}