learning_ai_common_plat/packages/auth/src/jwt.ts

/**
 * JWT utilities — configurable issuer, expiry, and algorithm.
 * Supports HS256 (symmetric, default) and RS256 (asymmetric) via jose.
 *
 * RS256 mode (Phase 4C SmartAuth):
 * - Sign with RSA private key (PEM)
 * - Verify with RSA public key (PEM) or remote JWKS URL
 * - Dual verification: tries RS256 first, falls back to HS256 during migration
 */

import {
  SignJWT,
  jwtVerify,
  importPKCS8,
  importSPKI,
  createRemoteJWKSet,
  type CryptoKey as JoseCryptoKey,
} from 'jose';
import type { JwtUtils, JwtUtilsOptions, TokenPayload } from './types.js';

function getHmacSecret(): Uint8Array {
  const secret = process.env.JWT_SECRET;
  if (!secret) throw new Error('JWT_SECRET must be set');
  return new TextEncoder().encode(secret);
}

/**
 * Create a JWT utility set with the given issuer and expiry configuration.
 *
 * @example
 * ```ts
 * // HS256 (default, backward-compatible)
 * const jwt = createJwtUtils({ issuer: "bytelyst-platform" });
 *
 * // RS256 (SmartAuth Phase 4C)
 * const jwt = createJwtUtils({
 *   issuer: "bytelyst-platform",
 *   algorithm: "RS256",
 *   rsaPrivateKey: process.env.JWT_PRIVATE_KEY,
 *   rsaPublicKey: process.env.JWT_PUBLIC_KEY,
 * });
 *
 * // RS256 verify-only (product backends — no private key)
 * const jwt = createJwtUtils({
 *   issuer: "bytelyst-platform",
 *   algorithm: "RS256",
 *   jwksUrl: "https://api.bytelyst.com/auth/.well-known/jwks.json",
 * });
 * ```
 */
export function createJwtUtils(options: JwtUtilsOptions): JwtUtils {
  const {
    issuer,
    accessTokenExpiry = '1h',
    refreshTokenExpiry = '30d',
    algorithm = 'HS256',
    rsaPrivateKey,
    rsaPublicKey,
    jwksUrl,
  } = options;

  // ── Key caches ────────────────────────────────────

  let _rsaPrivateKeyObj: JoseCryptoKey | null = null;
  let _rsaPublicKeyObj: JoseCryptoKey | null = null;
  let _jwksKeySet: ReturnType<typeof createRemoteJWKSet> | null = null;

  async function getRsaPrivateKey(): Promise<JoseCryptoKey> {
    if (_rsaPrivateKeyObj) return _rsaPrivateKeyObj;
    if (!rsaPrivateKey) throw new Error('rsaPrivateKey is required for RS256 signing');
    _rsaPrivateKeyObj = (await importPKCS8(rsaPrivateKey, 'RS256')) as JoseCryptoKey;
    return _rsaPrivateKeyObj;
  }

  async function getRsaPublicKey(): Promise<JoseCryptoKey> {
    if (_rsaPublicKeyObj) return _rsaPublicKeyObj;
    if (!rsaPublicKey) throw new Error('rsaPublicKey is required for RS256 local verification');
    _rsaPublicKeyObj = (await importSPKI(rsaPublicKey, 'RS256')) as JoseCryptoKey;
    return _rsaPublicKeyObj;
  }

  function getJwksKeySet(): ReturnType<typeof createRemoteJWKSet> {
    if (_jwksKeySet) return _jwksKeySet;
    if (!jwksUrl) throw new Error('jwksUrl is required for remote JWKS verification');
    _jwksKeySet = createRemoteJWKSet(new URL(jwksUrl));
    return _jwksKeySet;
  }

  // ── Signing ───────────────────────────────────────

  async function sign(claims: Record<string, unknown>, expiry: string): Promise<string> {
    if (algorithm === 'RS256') {
      const key = await getRsaPrivateKey();
      return new SignJWT(claims)
        .setProtectedHeader({ alg: 'RS256' })
        .setIssuedAt()
        .setExpirationTime(expiry)
        .setIssuer(issuer)
        .sign(key);
    }
    return new SignJWT(claims)
      .setProtectedHeader({ alg: 'HS256' })
      .setIssuedAt()
      .setExpirationTime(expiry)
      .setIssuer(issuer)
      .sign(getHmacSecret());
  }

  // ── Verification (dual: RS256 first, HS256 fallback) ──

  async function verifyWithRS256(token: string): Promise<TokenPayload | null> {
    try {
      if (jwksUrl) {
        const keySet = getJwksKeySet();
        const { payload } = await jwtVerify(token, keySet, { issuer });
        return payload as unknown as TokenPayload;
      }
      if (rsaPublicKey) {
        const key = await getRsaPublicKey();
        const { payload } = await jwtVerify(token, key, { issuer });
        return payload as unknown as TokenPayload;
      }
      return null;
    } catch {
      return null;
    }
  }

  async function verifyWithHS256(token: string): Promise<TokenPayload | null> {
    try {
      const secret = getHmacSecret();
      const { payload } = await jwtVerify(token, secret, { issuer });
      return payload as unknown as TokenPayload;
    } catch {
      return null;
    }
  }

  return {
    async createAccessToken(payload) {
      return sign(
        {
          ...payload,
          productId: payload.productId || issuer,
          type: 'access',
        },
        accessTokenExpiry
      );
    },

    async createRefreshToken(payload) {
      return sign(
        {
          sub: payload.sub,
          productId: payload.productId || issuer,
          type: 'refresh',
        },
        refreshTokenExpiry
      );
    },

    async verifyToken(token: string) {
      // Dual verification: try RS256 first (if configured), then HS256 fallback
      if (algorithm === 'RS256' || jwksUrl || rsaPublicKey) {
        const rs256Result = await verifyWithRS256(token);
        if (rs256Result) return rs256Result;
      }
      // HS256 fallback (safe during migration; removed after full RS256 rollout)
      try {
        return await verifyWithHS256(token);
      } catch {
        return null;
      }
    },
  };
}