auth-portal/apps/api/src/lib/oidc/keys.ts

import {
  exportJWK,
  exportPKCS8,
  generateKeyPair,
  importJWK,
  importPKCS8,
  jwtVerify,
  SignJWT,
  type JWK,
} from 'jose'
import { createHash, randomUUID } from 'node:crypto'
import type { FastifyInstance } from 'fastify'
import {
  getActiveOidcSigningKey,
  insertOidcSigningKey,
  listOidcSigningKeys,
} from '@authportal/db'
import { oidcIssuerFromConfig, type AppConfig } from '../../config.js'

export type OidcKeyMaterial = {
  kid: string
  privateKey: CryptoKey
  publicJwk: JWK
}

let cached: OidcKeyMaterial | null = null
// importJWK — асинхронный разбор JWK; без кэша он выполнялся на каждый /oauth/userinfo.
const importedPublicKeys = new Map<string, CryptoKey>()

function publicJwkFromPrivateExport(jwk: JWK, kid: string): JWK {
  return {
    kty: jwk.kty,
    n: jwk.n,
    e: jwk.e,
    alg: 'RS256',
    use: 'sig',
    kid,
  }
}

async function materialFromPem(
  kid: string,
  privatePem: string,
  publicJwkJson?: string,
): Promise<OidcKeyMaterial> {
  // jose defaults extractable=false for private keys → exportJWK throws
  const privateKey = await importPKCS8(privatePem, 'RS256', {
    extractable: true,
  })

  if (publicJwkJson) {
    try {
      const stored = JSON.parse(publicJwkJson) as JWK
      return {
        kid,
        privateKey,
        publicJwk: publicJwkFromPrivateExport({ ...stored, kid }, kid),
      }
    } catch {
      // fall through to derive from private key
    }
  }

  const full = await exportJWK(privateKey)
  return {
    kid,
    privateKey,
    publicJwk: publicJwkFromPrivateExport(full, kid),
  }
}

/**
 * Ensure an active RSA signing key exists (env PEM or DB / generate).
 */
export async function ensureOidcSigningKey(
  app: FastifyInstance,
): Promise<OidcKeyMaterial> {
  if (cached) return cached

  const envPem = app.config.oidcRsaPrivateKey?.trim()
  if (envPem) {
    const kid = createHash('sha256').update(envPem).digest('hex').slice(0, 16)
    cached = await materialFromPem(kid, envPem)
    return cached
  }

  const existing = getActiveOidcSigningKey(app.db)
  if (existing) {
    cached = await materialFromPem(
      existing.kid,
      existing.privatePem,
      existing.publicJwkJson,
    )
    return cached
  }

  const { privateKey, publicKey } = await generateKeyPair('RS256', {
    extractable: true,
  })
  const kid = randomUUID().replace(/-/g, '').slice(0, 16)
  const privatePem = await exportPKCS8(privateKey)
  const publicJwk = await exportJWK(publicKey)
  publicJwk.alg = 'RS256'
  publicJwk.use = 'sig'
  publicJwk.kid = kid
  insertOidcSigningKey(app.db, {
    kid,
    privatePem,
    publicJwkJson: JSON.stringify(publicJwk),
  })
  cached = { kid, privateKey, publicJwk }
  app.log.info({ kid }, 'OIDC RS256 signing key generated')
  return cached
}

/** Ключ подписи загружен — приложение готово выдавать токены. */
export function isOidcKeyReady(): boolean {
  return cached !== null
}

/** Reset cache (tests). */
export function resetOidcKeyCache(): void {
  cached = null
  importedPublicKeys.clear()
  jwksCache = new WeakMap()
}

export function buildOidcDiscovery(config: AppConfig) {
  const issuer = oidcIssuerFromConfig(config)
  return {
    issuer,
    authorization_endpoint: `${issuer}/oauth/authorize`,
    token_endpoint: `${issuer}/oauth/token`,
    userinfo_endpoint: `${issuer}/oauth/userinfo`,
    jwks_uri: `${issuer}/.well-known/jwks.json`,
    response_types_supported: ['code'],
    subject_types_supported: ['public'],
    id_token_signing_alg_values_supported: ['RS256'],
    scopes_supported: ['openid', 'profile', 'email', 'groups'],
    token_endpoint_auth_methods_supported: [
      'client_secret_post',
      'client_secret_basic',
    ],
    claims_supported: [
      'sub',
      'iss',
      'aud',
      'exp',
      'iat',
      'email',
      'email_verified',
      'name',
      'preferred_username',
      'groups',
      'roles',
    ],
    grant_types_supported: ['authorization_code'],
    code_challenge_methods_supported: ['S256', 'plain'],
  }
}

export async function buildJwks(
  app: FastifyInstance,
): Promise<{ keys: JWK[] }> {
  await ensureOidcSigningKey(app)
  const hit = jwksCache.get(app.db)
  if (hit && hit.kid === cached?.kid) return hit.jwks
  const jwks = computeJwks(app)
  jwksCache.set(app.db, { kid: cached?.kid ?? '', jwks })
  return jwks
}

/** Мемоизация JWKS: набор ключей меняется только при выпуске нового (меняется kid). */
let jwksCache = new WeakMap<object, { kid: string; jwks: { keys: JWK[] } }>()

function computeJwks(app: FastifyInstance): { keys: JWK[] } {
  const keys = listOidcSigningKeys(app.db)
    .map((row): JWK | null => {
      try {
        const jwk = JSON.parse(row.publicJwkJson) as JWK
        return { ...jwk, kid: row.kid, alg: 'RS256', use: 'sig' }
      } catch {
        return null
      }
    })
    .filter((k): k is JWK => k !== null)

  if (cached && !keys.some((k) => k.kid === cached!.kid)) {
    keys.unshift(cached.publicJwk)
  }
  return { keys }
}

/** Принудительный сброс JWKS (после ручных правок ключей в БД). */
export function invalidateJwks(app: FastifyInstance): void {
  jwksCache.delete(app.db)
}

export async function signOidcJwt(
  app: FastifyInstance,
  claims: Record<string, unknown>,
  expiresInSeconds: number,
): Promise<string> {
  const key = await ensureOidcSigningKey(app)
  const issuer = oidcIssuerFromConfig(app.config)
  return new SignJWT(claims)
    .setProtectedHeader({ alg: 'RS256', kid: key.kid, typ: 'JWT' })
    .setIssuer(issuer)
    .setIssuedAt()
    .setExpirationTime(`${expiresInSeconds}s`)
    .sign(key.privateKey)
}

/** Verify OIDC access token (RS256) and return payload. */
export async function verifyOidcAccessToken(
  app: FastifyInstance,
  token: string,
): Promise<Record<string, unknown>> {
  const key = await ensureOidcSigningKey(app)
  const issuer = oidcIssuerFromConfig(app.config)
  let publicKey: CryptoKey | undefined = importedPublicKeys.get(key.kid)
  if (!publicKey) {
    // Для RS256 + RSA JWK importJWK всегда возвращает CryptoKey.
    publicKey = (await importJWK(key.publicJwk, 'RS256')) as CryptoKey
    importedPublicKeys.set(key.kid, publicKey)
  }
  const { payload } = await jwtVerify(token, publicKey, {
    issuer,
    algorithms: ['RS256'],
  })
  return payload as Record<string, unknown>
}