codey.lol/src/utils/jwt.ts

import jwt from 'jsonwebtoken';
import crypto from 'crypto';
import { API_URL } from '../config';

// JWKS endpoint for fetching RSA public keys
const JWKS_URL = `${API_URL}/auth/jwks`;

// Cache for JWKS keys
interface JwkKey {
  kty: string;
  use?: string;
  kid: string;
  alg?: string;
  n: string;
  e: string;
}

interface JwksResponse {
  keys: JwkKey[];
}

interface CachedKeys {
  keys: Map<string, string>; // kid -> PEM public key
  fetchedAt: number;
}

export interface JwtPayload {
  sub?: string;
  exp?: number;
  iat?: number;
  roles?: string[];
  [key: string]: unknown;
}

// Cache TTL: 5 minutes (keys are refreshed on cache miss or verification failure)
const CACHE_TTL_MS = 5 * 60 * 1000;

let cachedKeys: CachedKeys | null = null;

/**
 * Convert a JWK RSA public key to PEM format
 */
function jwkToPem(jwk: JwkKey): string {
  if (jwk.kty !== 'RSA') {
    throw new Error(`Unsupported key type: ${jwk.kty}`);
  }

  // Decode base64url-encoded modulus and exponent
  const n = Buffer.from(jwk.n, 'base64url');
  const e = Buffer.from(jwk.e, 'base64url');

  // Build the RSA public key using Node's crypto
  const publicKey = crypto.createPublicKey({
    key: {
      kty: 'RSA',
      n: jwk.n,
      e: jwk.e,
    },
    format: 'jwk',
  });

  return publicKey.export({ type: 'spki', format: 'pem' }) as string;
}

/**
 * Fetch JWKS from the API and cache the keys
 */
async function fetchJwks(): Promise<Map<string, string>> {
  try {
    const response = await fetch(JWKS_URL, {
      headers: { 'Accept': 'application/json' },
    });

    if (!response.ok) {
      throw new Error(`JWKS fetch failed: ${response.status} ${response.statusText}`);
    }

    const jwks: JwksResponse = await response.json();
    const keys = new Map<string, string>();

    for (const jwk of jwks.keys) {
      if (jwk.kty === 'RSA' && jwk.kid) {
        try {
          const pem = jwkToPem(jwk);
          keys.set(jwk.kid, pem);
        } catch (err) {
          console.warn(`Failed to convert JWK ${jwk.kid} to PEM:`, (err as Error).message);
        }
      }
    }

    if (keys.size === 0) {
      throw new Error('No valid RSA keys found in JWKS');
    }

    cachedKeys = {
      keys,
      fetchedAt: Date.now(),
    };

    console.log(`[JWT] Fetched ${keys.size} RSA public key(s) from JWKS`);
    return keys;

  } catch (error) {
    console.error('[JWT] Failed to fetch JWKS:', (error as Error).message);
    throw error;
  }
}

/**
 * Get cached keys or fetch fresh ones
 */
async function getKeys(forceRefresh = false): Promise<Map<string, string>> {
  const now = Date.now();
  
  if (!forceRefresh && cachedKeys && (now - cachedKeys.fetchedAt) < CACHE_TTL_MS) {
    return cachedKeys.keys;
  }

  return fetchJwks();
}

/**
 * Verify a JWT token using RS256 with keys from JWKS
 */
export async function verifyToken(token: string | null | undefined): Promise<JwtPayload | null> {
  if (!token) {
    return null;
  }

  try {
    const decoded = jwt.decode(token, { complete: true });
    if (!decoded?.header?.kid) {
      throw new Error('No kid in token header');
    }

    const kid = decoded.header.kid;
    let keys = await getKeys();
    let publicKey = keys.get(kid);

    // If key not found, try refreshing JWKS (handles key rotation)
    if (!publicKey) {
      console.log(`[JWT] Key ${kid} not in cache, refreshing JWKS...`);
      keys = await getKeys(true);
      publicKey = keys.get(kid);

      if (!publicKey) {
        throw new Error(`Unknown kid: ${kid}`);
      }
    }

    // Verify using RS256
    const payload = jwt.verify(token, publicKey, { algorithms: ['RS256'] }) as JwtPayload;
    return payload;

  } catch (error) {
    console.error('JWT verification failed:', (error as Error).message);
    return null;
  }
}

/**
 * Synchronous verification - uses cached keys only (for middleware compatibility)
 * Falls back to null if cache is empty; caller should handle async verification
 */
export function verifyTokenSync(token: string | null | undefined): JwtPayload | null {
  if (!token) {
    return null;
  }

  if (!cachedKeys || cachedKeys.keys.size === 0) {
    // No cached keys - caller needs to use async version
    return null;
  }

  try {
    const decoded = jwt.decode(token, { complete: true });
    if (!decoded?.header?.kid) {
      throw new Error('No kid in token header');
    }

    const kid = decoded.header.kid;
    const publicKey = cachedKeys.keys.get(kid);

    if (!publicKey) {
      // Key not found - might need refresh, return null to signal async verification needed
      return null;
    }

    // Verify using RS256
    const payload = jwt.verify(token, publicKey, { algorithms: ['RS256'] }) as JwtPayload;
    return payload;

  } catch (error) {
    console.error('JWT verification failed:', (error as Error).message);
    return null;
  }
}

/**
 * Pre-warm the JWKS cache at startup
 */
export async function initializeJwks(): Promise<void> {
  try {
    await fetchJwks();
  } catch (error) {
    console.error('[JWT] Failed to pre-warm JWKS cache:', (error as Error).message);
    // Don't throw - allow lazy initialization on first verification
  }
}