learning_ai_common_plat/services/platform-service/src/lib/api-key-auth.ts

import crypto from 'node:crypto';
import bcrypt from 'bcryptjs';
import type { FastifyInstance, FastifyRequest } from 'fastify';
import { ForbiddenError, TooManyRequestsError, UnauthorizedError } from './errors.js';
import { getCollection } from './datastore.js';
import * as rateLimitStore from '../modules/ratelimit/store.js';
import type { RateLimitRule } from '../modules/ratelimit/types.js';
import * as auditRepo from '../modules/audit/repository.js';
import type { AuditDoc } from '../modules/audit/types.js';

interface ApiTokenLookupDoc {
  id: string;
  productId: string;
  userId: string;
  userName: string;
  tokenType: 'user_api' | 'product_api' | 'service_api';
  environment: 'dev' | 'staging' | 'prod';
  prefix: string;
  tokenHash: string;
  status: 'active' | 'revoked' | 'expired';
  scopes: string[];
  expiresAt: string;
  lastUsed: string | null;
}

export interface ApiKeyAuthPayload {
  tokenId: string;
  productId: string;
  userId: string;
  userName: string;
  tokenType: 'product_api' | 'service_api';
  environment: 'dev' | 'staging' | 'prod';
  scopes: string[];
  prefix: string;
}

interface AccessOptions {
  allowJwt?: boolean;
  jwtRoles?: string[];
  apiKeyScopes?: string[];
  apiKeyTokenTypes?: ApiKeyAuthPayload['tokenType'][];
  rateLimitKey?: string;
}

interface AccessActor {
  actorId: string;
  productId: string;
  source: 'jwt' | 'api_key';
}

interface ApiKeyRateLimitConfig {
  [key: string]: RateLimitRule;
}

declare module 'fastify' {
  interface FastifyRequest {
    apiKeyAuth?: ApiKeyAuthPayload;
  }
}

function tokenCollection() {
  return getCollection<ApiTokenLookupDoc>('api_tokens', '/id');
}

function resolveAuditProductId(req: FastifyRequest): string {
  const headerProductId = req.headers['x-product-id'];
  if (req.apiKeyAuth?.productId) return req.apiKeyAuth.productId;
  if (typeof headerProductId === 'string' && headerProductId.length > 0) return headerProductId;
  return process.env.DEFAULT_PRODUCT_ID ?? 'lysnrai';
}

function writeApiKeyAudit(
  req: FastifyRequest,
  action: 'api_key.access_rejected' | 'api_key.rate_limited',
  details: Record<string, unknown>
): void {
  const doc: AuditDoc = {
    id: `aud_${crypto.randomUUID()}`,
    productId: resolveAuditProductId(req),
    userId: req.apiKeyAuth?.userId ?? 'api_key_unknown',
    action,
    category: 'security',
    details,
    ipAddress: req.ip,
    userAgent:
      typeof req.headers['user-agent'] === 'string' ? req.headers['user-agent'] : undefined,
    createdAt: new Date().toISOString(),
  };

  auditRepo.create(doc).catch(err => {
    req.log.error({ err, action }, '[api-key] audit write failed');
  });
}

function logApiKeyWarning(
  req: FastifyRequest,
  event: string,
  details: Record<string, unknown> = {}
): void {
  writeApiKeyAudit(
    req,
    event.includes('rate_limited') ? 'api_key.rate_limited' : 'api_key.access_rejected',
    {
      event,
      requestId: req.id,
      productIdHeader: req.headers['x-product-id'],
      apiKeyPrefix: req.apiKeyAuth?.prefix,
      apiKeyTokenId: req.apiKeyAuth?.tokenId,
      apiKeyProductId: req.apiKeyAuth?.productId,
      ...details,
    }
  );

  req.log.warn(
    {
      event,
      requestId: req.id,
      productIdHeader: req.headers['x-product-id'],
      apiKeyPrefix: req.apiKeyAuth?.prefix,
      apiKeyTokenId: req.apiKeyAuth?.tokenId,
      apiKeyProductId: req.apiKeyAuth?.productId,
      ...details,
    },
    '[api-key] access rejected'
  );
}

function extractApiKey(req: FastifyRequest): string | null {
  const header = req.headers['x-api-key'];
  if (typeof header === 'string' && header.trim().length > 0) {
    return header.trim();
  }

  const auth = req.headers.authorization;
  if (!auth) return null;

  if (auth.startsWith('ApiKey ')) {
    return auth.slice('ApiKey '.length).trim();
  }

  if (auth.startsWith('Bearer wai_')) {
    return auth.slice('Bearer '.length).trim();
  }

  return null;
}

function loadApiKeyRateLimitConfig(): ApiKeyRateLimitConfig {
  const defaults: ApiKeyRateLimitConfig = {
    'jobs:read': { maxRequests: 60, windowSeconds: 60 },
    'jobs:write': { maxRequests: 15, windowSeconds: 60 },
    'exports:read': { maxRequests: 30, windowSeconds: 60 },
    'exports:write': { maxRequests: 10, windowSeconds: 60 },
    'maintenance:read': { maxRequests: 30, windowSeconds: 60 },
    'maintenance:write': { maxRequests: 10, windowSeconds: 60 },
    'ip-rules:read': { maxRequests: 30, windowSeconds: 60 },
    'ip-rules:write': { maxRequests: 10, windowSeconds: 60 },
    'webhooks:read': { maxRequests: 60, windowSeconds: 60 },
    'webhooks:write': { maxRequests: 15, windowSeconds: 60 },
  };

  const raw = process.env.API_KEY_RATE_LIMIT_CONFIG_JSON;
  if (!raw) return defaults;

  try {
    const parsed = JSON.parse(raw) as ApiKeyRateLimitConfig;
    return { ...defaults, ...parsed };
  } catch {
    return defaults;
  }
}

function loadApiKeyProductRateLimitConfig(
  baseConfig: ApiKeyRateLimitConfig
): ApiKeyRateLimitConfig {
  const raw = process.env.API_KEY_PRODUCT_RATE_LIMIT_CONFIG_JSON;
  if (raw) {
    try {
      return JSON.parse(raw) as ApiKeyRateLimitConfig;
    } catch {
      /* fall back to derived defaults */
    }
  }

  return Object.fromEntries(
    Object.entries(baseConfig).map(([key, rule]) => [
      key,
      {
        maxRequests: Math.max(rule.maxRequests * 5, rule.maxRequests),
        windowSeconds: rule.windowSeconds,
      },
    ])
  );
}

function tokenHasScope(grantedScopes: string[], requiredScope: string): boolean {
  if (grantedScopes.includes('*')) return true;

  return grantedScopes.some(scope => {
    if (scope === requiredScope) return true;
    if (scope.endsWith('*')) {
      const prefix = scope.slice(0, -1);
      return requiredScope.startsWith(prefix);
    }
    return false;
  });
}

function getRuntimeEnvironment(): 'dev' | 'staging' | 'prod' {
  const explicit = process.env.PLATFORM_RUNTIME_ENV;
  if (explicit === 'dev' || explicit === 'staging' || explicit === 'prod') {
    return explicit;
  }

  switch (process.env.NODE_ENV) {
    case 'production':
      return 'prod';
    case 'test':
      return 'dev';
    default:
      return 'dev';
  }
}

function ensureApiKeyScopes(
  req: FastifyRequest,
  requiredScopes: string[] = [],
  allowedTokenTypes?: ApiKeyAuthPayload['tokenType'][]
): ApiKeyAuthPayload {
  const apiKey = req.apiKeyAuth;
  if (!apiKey) {
    logApiKeyWarning(req, 'missing_api_key_context');
    throw new UnauthorizedError('API key required');
  }

  const productIdHeader = req.headers['x-product-id'];
  if (typeof productIdHeader === 'string' && productIdHeader.length > 0) {
    if (productIdHeader !== apiKey.productId) {
      logApiKeyWarning(req, 'product_mismatch', {
        requestedProductId: productIdHeader,
      });
      throw new ForbiddenError('API key is not valid for the requested product');
    }
  }

  if (
    requiredScopes.length > 0 &&
    !requiredScopes.every(scope => tokenHasScope(apiKey.scopes, scope))
  ) {
    logApiKeyWarning(req, 'missing_scope', {
      requiredScopes,
      grantedScopes: apiKey.scopes,
    });
    throw new ForbiddenError('API key missing required scopes');
  }

  if (allowedTokenTypes && allowedTokenTypes.length > 0) {
    if (!allowedTokenTypes.includes(apiKey.tokenType)) {
      logApiKeyWarning(req, 'token_type_not_permitted', {
        allowedTokenTypes,
        tokenType: apiKey.tokenType,
      });
      throw new ForbiddenError('API key token type is not permitted for this route');
    }
  }

  return apiKey;
}

async function enforceApiKeyRateLimit(req: FastifyRequest, rateLimitKey?: string): Promise<void> {
  if (!rateLimitKey || !req.apiKeyAuth) return;

  const apiKeyRateLimitConfig = loadApiKeyRateLimitConfig();
  const keyRule = apiKeyRateLimitConfig[rateLimitKey];
  if (!keyRule) return;

  const compositeKey = `api-key:${req.apiKeyAuth.productId}:${req.apiKeyAuth.tokenId}:${rateLimitKey}`;
  const keyResult = await rateLimitStore.checkAndRecord(compositeKey, keyRule);
  if (!keyResult.allowed) {
    logApiKeyWarning(req, 'token_rate_limited', {
      rateLimitKey,
      retryAfterMs: keyResult.retryAfterMs,
    });
    throw new TooManyRequestsError('API key rate limit exceeded', {
      retryAfter: Math.ceil((keyResult.retryAfterMs ?? 0) / 1000),
    });
  }

  const productRule = loadApiKeyProductRateLimitConfig(apiKeyRateLimitConfig)[rateLimitKey];
  if (!productRule) return;

  const productCompositeKey = `api-key-product:${req.apiKeyAuth.productId}:${rateLimitKey}`;
  const productResult = await rateLimitStore.checkAndRecord(productCompositeKey, productRule);
  if (!productResult.allowed) {
    logApiKeyWarning(req, 'product_rate_limited', {
      rateLimitKey,
      retryAfterMs: productResult.retryAfterMs,
    });
    throw new TooManyRequestsError('Product API key rate limit exceeded', {
      retryAfter: Math.ceil((productResult.retryAfterMs ?? 0) / 1000),
    });
  }
}

export async function requireJwtOrApiKey(
  req: FastifyRequest,
  { allowJwt = false, jwtRoles, apiKeyScopes, apiKeyTokenTypes, rateLimitKey }: AccessOptions = {}
): Promise<AccessActor> {
  const jwt = req.jwtPayload;
  if (jwt?.sub) {
    if (jwtRoles && jwtRoles.length > 0) {
      if (!jwt.role || !jwtRoles.includes(jwt.role)) {
        throw new ForbiddenError('Admin access required');
      }
    } else if (!allowJwt) {
      throw new ForbiddenError('JWT access is not permitted for this route');
    }

    return {
      actorId: jwt.sub,
      productId: jwt.productId ?? process.env.DEFAULT_PRODUCT_ID ?? 'lysnrai',
      source: 'jwt',
    };
  }

  const apiKey = ensureApiKeyScopes(req, apiKeyScopes, apiKeyTokenTypes);
  await enforceApiKeyRateLimit(req, rateLimitKey);

  return {
    actorId: apiKey.userId,
    productId: apiKey.productId,
    source: 'api_key',
  };
}

export async function registerOptionalApiKeyContext(app: FastifyInstance): Promise<void> {
  app.addHook('onRequest', async req => {
    const rawKey = extractApiKey(req);
    if (!rawKey) return;

    const prefix = rawKey.slice(0, 12);
    const filter: Record<string, unknown> = {
      prefix,
      status: 'active',
      expiresAt: { $gte: new Date().toISOString() },
    };

    const candidates = await tokenCollection().findMany({
      filter,
      limit: 10,
    });

    let rejectionReason: string | null = null;

    for (const candidate of candidates) {
      const ok = await bcrypt.compare(rawKey, candidate.tokenHash);
      if (!ok) continue;

      if (candidate.tokenType === 'user_api') {
        rejectionReason = 'user_api_not_allowed';
        continue;
      }

      if (candidate.environment !== getRuntimeEnvironment()) {
        rejectionReason = 'environment_mismatch';
        continue;
      }

      req.apiKeyAuth = {
        tokenId: candidate.id,
        productId: candidate.productId,
        userId: candidate.userId,
        userName: candidate.userName,
        tokenType: candidate.tokenType,
        environment: candidate.environment,
        scopes: candidate.scopes,
        prefix: candidate.prefix,
      };

      tokenCollection()
        .update(candidate.id, candidate.id, { lastUsed: new Date().toISOString() })
        .catch(() => {});
      return;
    }

    req.log.warn(
      {
        event: rejectionReason ?? 'invalid_api_key',
        requestId: req.id,
        apiKeyPrefix: prefix,
        productIdHeader: req.headers['x-product-id'],
      },
      '[api-key] authentication failed'
    );
    writeApiKeyAudit(req, 'api_key.access_rejected', {
      event: rejectionReason ?? 'invalid_api_key',
      requestId: req.id,
      apiKeyPrefix: prefix,
      productIdHeader: req.headers['x-product-id'],
    });
  });
}