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

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';

interface ApiTokenLookupDoc {
  id: string;
  productId: string;
  userId: string;
  userName: string;
  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;
  scopes: string[];
  prefix: string;
}

interface AccessOptions {
  allowJwt?: boolean;
  jwtRoles?: string[];
  apiKeyScopes?: string[];
  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 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 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 ensureApiKeyScopes(req: FastifyRequest, requiredScopes: string[] = []): ApiKeyAuthPayload {
  const apiKey = req.apiKeyAuth;
  if (!apiKey) {
    throw new UnauthorizedError('API key required');
  }

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

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

  return apiKey;
}

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

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

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

export function requireJwtOrApiKey(
  req: FastifyRequest,
  { allowJwt = false, jwtRoles, apiKeyScopes, rateLimitKey }: AccessOptions = {}
): 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);
  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,
    });

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

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

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