apophis-fastify/src/extensions/jwt.ts

/**
 * JWT Extension for APOPHIS
 *
 * Provides APOSTL predicates for JWT claim inspection, header access,
 * and signature validation. Supports extracting JWT from Authorization
 * header, response body, or custom headers.
 *
 * Example:
 * ```typescript
 * import { jwtExtension } from '@apophis/fastify/extensions/jwt'
 *
 * await fastify.register(apophis, {
 *   extensions: [
 *     jwtExtension({
 *       jwks: 'https://auth.example.com/.well-known/jwks.json',
 *       verify: true,
 *     })
 *   ]
 * })
 * ```
 *
 * APOSTL formulas:
 * ```apostl
 * jwt_claims(this).sub != null
 * jwt_claims(this).exp > jwt_claims(this).iat
 * jwt_header(this).alg == "RS256"
 * jwt_valid(this) == true
 * jwt_format(this) == "compact"
 * ```
 */

import type { ApophisExtension, PredicateContext } from '../extension/types.js'
import { createHmac, createPublicKey, createVerify } from 'node:crypto'

// ============================================================================
// Types
// ============================================================================

export interface JwtExtensionConfig {
  /** JWKS endpoint URL for signature verification */
  jwks?: string
  /** Inline key material: kid -> PEM or JWK */
  keys?: Record<string, string>
  /** Where to extract JWT from. Default: 'authorization' */
  extractFrom?: 'authorization' | 'body' | 'header' | 'query'
  /** Custom header name when extractFrom is 'header' */
  headerName?: string
  /** Whether to verify signatures. Default: true */
  verify?: boolean
}

interface DecodedJwt {
  header: Record<string, unknown>
  payload: Record<string, unknown>
  signature: string
}

// ============================================================================
// Base64URL Decoding (no verification)
// ============================================================================

function base64UrlDecode(str: string): string {
  const base64 = str.replace(/-/g, '+').replace(/_/g, '/')
  const padLen = (4 - (base64.length % 4)) % 4
  const padded = base64 + '='.repeat(padLen)
  return Buffer.from(padded, 'base64').toString('utf8')
}

function base64UrlToBuffer(str: string): Buffer {
  const base64 = str.replace(/-/g, '+').replace(/_/g, '/')
  const padLen = (4 - (base64.length % 4)) % 4
  const padded = base64 + '='.repeat(padLen)
  return Buffer.from(padded, 'base64')
}

function parseCompactJwt(token: string): DecodedJwt | null {
  const parts = token.split('.')
  if (parts.length !== 3) return null

  try {
    const header = JSON.parse(base64UrlDecode(parts[0]!))
    const payload = JSON.parse(base64UrlDecode(parts[1]!))
    return { header, payload, signature: parts[2]! }
  } catch {
    return null
  }
}

// ============================================================================
// JWT Extraction
// ============================================================================

function extractJwt(ctx: PredicateContext, config: JwtExtensionConfig): string | null {
  const extractFrom = config.extractFrom ?? 'authorization'

  switch (extractFrom) {
    case 'authorization': {
      const auth = ctx.evalContext.request.headers['authorization']
      if (typeof auth === 'string') {
        const match = auth.match(/^Bearer\s+(.+)$/i)
        if (match && match[1]) {
          return match[1].trim()
        }
      }
      return null
    }
    case 'body': {
      const body = ctx.evalContext.response.body
      if (body && typeof body === 'object') {
        const token = (body as Record<string, unknown>)['access_token']
          ?? (body as Record<string, unknown>)['token']
          ?? (body as Record<string, unknown>)['jwt']
        if (typeof token === 'string') return token
      }
      return null
    }
    case 'header': {
      const headerName = config.headerName ?? 'x-jwt'
      const value = ctx.evalContext.request.headers[headerName.toLowerCase()]
      if (typeof value === 'string') return value
      return null
    }
    case 'query': {
      const query = ctx.evalContext.request.query
      const token = query['token'] ?? query['jwt'] ?? query['access_token']
      if (typeof token === 'string') return token
      return null
    }
    default:
      return null
  }
}

// ============================================================================
// Extension State
// ============================================================================

interface JwtState {
  seenJtis: Set<string>
  decodedCache: Map<string, DecodedJwt>
  jwksCache?: Record<string, unknown>
}

// ============================================================================
// Predicate Resolvers
// ============================================================================

function resolveClaims(ctx: PredicateContext, config: JwtExtensionConfig, state: JwtState): unknown {
  const token = extractJwt(ctx, config)
  if (!token) return null

  const cached = state.decodedCache.get(token)
  if (cached) return cached.payload

  const decoded = parseCompactJwt(token)
  if (!decoded) return null

  state.decodedCache.set(token, decoded)

  // Track JTI for replay detection
  const jti = decoded.payload['jti']
  if (typeof jti === 'string') {
    state.seenJtis.add(jti)
  }

  return decoded.payload
}

function resolveHeader(ctx: PredicateContext, config: JwtExtensionConfig, state: JwtState): unknown {
  const token = extractJwt(ctx, config)
  if (!token) return null

  const cached = state.decodedCache.get(token)
  if (cached) return cached.header

  const decoded = parseCompactJwt(token)
  if (!decoded) return null

  state.decodedCache.set(token, decoded)
  return decoded.header
}

function getVerificationKey(config: JwtExtensionConfig, kid: string | undefined): string | undefined {
  if (!config.keys) return undefined
  if (kid && config.keys[kid]) return config.keys[kid]
  return config.keys.default
}

function verifyJwt(token: string, decoded: DecodedJwt, config: JwtExtensionConfig): boolean {
  const alg = typeof decoded.header.alg === 'string' ? decoded.header.alg : undefined
  if (!alg) return false

  // Reject unsupported algorithms early to prevent alg confusion
  const supportedAlgs = ['HS256', 'RS256', 'ES256']
  if (!supportedAlgs.includes(alg)) return false

  const [encodedHeader, encodedPayload, encodedSignature] = token.split('.')
  if (!encodedHeader || !encodedPayload || !encodedSignature) return false

  const signingInput = `${encodedHeader}.${encodedPayload}`
  const signature = base64UrlToBuffer(encodedSignature)
  const kid = typeof decoded.header.kid === 'string' ? decoded.header.kid : undefined
  const key = getVerificationKey(config, kid)
  if (!key) return false

  if (alg === 'HS256') {
    const expected = createHmac('sha256', key).update(signingInput).digest()
    return expected.length === signature.length && expected.equals(signature)
  }

  if (alg === 'RS256') {
    const verifier = createVerify('RSA-SHA256')
    verifier.update(signingInput)
    verifier.end()
    return verifier.verify(createPublicKey(key), signature)
  }

  if (alg === 'ES256') {
    const verifier = createVerify('SHA256')
    verifier.update(signingInput)
    verifier.end()
    return verifier.verify(createPublicKey(key), signature)
  }

  return false
}

function resolveValid(ctx: PredicateContext, config: JwtExtensionConfig, state: JwtState): unknown {
  const token = extractJwt(ctx, config)
  if (!token) return false

  const cached = state.decodedCache.get(token)
  const decoded = cached ?? parseCompactJwt(token)
  if (!decoded) return false
  if (!cached) {
    state.decodedCache.set(token, decoded)
  }

  if (config.verify === false) {
    return true
  }

  return verifyJwt(token, decoded, config)
}

function resolveFormat(ctx: PredicateContext, config: JwtExtensionConfig): unknown {
  const token = extractJwt(ctx, config)
  if (!token) return null

  const parts = token.split('.')
  if (parts.length === 3) return 'compact'
  if (parts.length === 1) {
    try {
      JSON.parse(base64UrlDecode(parts[0]!))
      return 'json'
    } catch {
      return null
    }
  }
  return null
}

// ============================================================================
// Extension Factory
// ============================================================================

/**
 * Create a JWT extension for APOPHIS.
 *
 * Provides predicates for inspecting JWT claims and headers without
 * requiring the application under test to expose them.
 *
 * @param config - JWT extension configuration
 * @returns ApophisExtension ready for registration
 *
 * @example
 * ```typescript
 * const jwt = jwtExtension({
 *   jwks: 'https://auth.example.com/.well-known/jwks.json',
 *   verify: true,
 * })
 *
 * await fastify.register(apophis, { extensions: [jwt] })
 * ```
 */
export function jwtExtension(config: JwtExtensionConfig = {}): ApophisExtension {
  return {
    name: 'jwt',
    productionSafe: true,

    headers: ['jwt_claims', 'jwt_header', 'jwt_valid', 'jwt_format'],

    onSuiteStart: () => {
      const state: Record<string, unknown> = {
        seenJtis: new Set<string>(),
        decodedCache: new Map<string, DecodedJwt>(),
      }
      return state
    },

    predicates: {
      jwt_claims: (ctx) => {
        const state = ctx.extensionState as unknown as JwtState
        const value = resolveClaims(ctx, config, state)
        return { value, success: true }
      },

      jwt_header: (ctx) => {
        const state = ctx.extensionState as unknown as JwtState
        const value = resolveHeader(ctx, config, state)
        return { value, success: true }
      },

      jwt_valid: (ctx) => {
        const state = ctx.extensionState as unknown as JwtState
        const value = resolveValid(ctx, config, state)
        return { value, success: true }
      },

      jwt_format: (ctx) => {
        const value = resolveFormat(ctx, config)
        return { value, success: true }
      },
    },
  }
}

// Type is already exported at declaration site