stack/apps/api/src/prisma/llm-encryption.middleware.ts

/**
 * LLM Encryption Middleware
 *
 * Prisma middleware that transparently encrypts/decrypts LLM provider API keys
 * in the LlmProviderInstance.config JSON field using OpenBao Transit encryption.
 *
 * Encryption happens on:
 * - create: New provider instance records
 * - update/updateMany: Config updates
 * - upsert: Both create and update data
 *
 * Decryption happens on:
 * - findUnique/findMany/findFirst: Read operations
 *
 * Format detection:
 * - `vault:v1:...` = OpenBao Transit encrypted
 * - Otherwise = Legacy plaintext (backward compatible)
 */

import { Logger } from "@nestjs/common";
import type { PrismaClient } from "@prisma/client";
import type { VaultService } from "../vault/vault.service";
import { TransitKey } from "../vault/vault.constants";

/**
 * Prisma middleware parameters interface
 */
interface MiddlewareParams {
  model?: string;
  action: string;
  args: {
    data?: Record<string, unknown>;
    where?: Record<string, unknown>;
    select?: Record<string, unknown>;
    create?: Record<string, unknown>;
    update?: Record<string, unknown>;
  };
  dataPath: string[];
  runInTransaction: boolean;
}

/**
 * LlmProviderInstance data with config field
 */
interface LlmProviderInstanceData extends Record<string, unknown> {
  config?: LlmProviderConfig;
}

/**
 * LLM provider configuration (JSON field)
 */
interface LlmProviderConfig {
  apiKey?: string | null;
  endpoint?: string;
  [key: string]: unknown;
}

/**
 * Register LLM encryption middleware on Prisma client
 *
 * @param prisma - Prisma client instance
 * @param vaultService - Vault service for encryption/decryption
 */
export function registerLlmEncryptionMiddleware(
  prisma: PrismaClient,
  vaultService: VaultService
): void {
  const logger = new Logger("LlmEncryptionMiddleware");

  // TODO: Replace with Prisma Client Extensions (https://www.prisma.io/docs/concepts/components/prisma-client/client-extensions)
  // when stable. Client extensions provide a type-safe alternative to middleware without requiring
  // type assertions or eslint-disable directives. Migration path:
  // 1. Wait for Prisma 6.x stable release with full extension support
  // 2. Create extension using prisma.$extends({ query: { llmProviderInstance: { ... } } })
  // 3. Remove this middleware and eslint-disable comments
  // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
  (prisma as any).$use(
    async (params: MiddlewareParams, next: (params: MiddlewareParams) => Promise<unknown>) => {
      // Only process LlmProviderInstance model operations
      if (params.model !== "LlmProviderInstance") {
        return next(params);
      }

      // Encrypt on write operations
      if (
        params.action === "create" ||
        params.action === "update" ||
        params.action === "updateMany"
      ) {
        if (params.args.data) {
          await encryptConfig(params.args.data as LlmProviderInstanceData, vaultService);
        }
      } else if (params.action === "upsert") {
        // Handle upsert - encrypt both create and update data
        if (params.args.create) {
          await encryptConfig(params.args.create as LlmProviderInstanceData, vaultService);
        }
        if (params.args.update) {
          await encryptConfig(params.args.update as LlmProviderInstanceData, vaultService);
        }
      }

      // Execute query
      const result = await next(params);

      // Decrypt on read operations
      if (params.action === "findUnique" || params.action === "findFirst") {
        if (result && typeof result === "object") {
          await decryptConfig(result as LlmProviderInstanceData, vaultService, logger);
        }
      } else if (params.action === "findMany") {
        if (Array.isArray(result)) {
          for (const instance of result) {
            if (instance && typeof instance === "object") {
              await decryptConfig(instance as LlmProviderInstanceData, vaultService, logger);
            }
          }
        }
      }

      return result;
    }
  );
}

/**
 * Encrypt apiKey in config JSON field
 * Modifies data in-place
 *
 * @param data - LlmProviderInstance data object
 * @param vaultService - Vault service
 */
async function encryptConfig(
  data: LlmProviderInstanceData,
  vaultService: VaultService
): Promise<void> {
  // Skip if no config field
  if (!data.config || typeof data.config !== "object") {
    return;
  }

  const config = data.config;

  // Skip if no apiKey field
  if (!config.apiKey || typeof config.apiKey !== "string") {
    return;
  }

  // Skip if already encrypted (idempotent)
  if (isEncrypted(config.apiKey)) {
    return;
  }

  // Encrypt plaintext apiKey
  const ciphertext = await vaultService.encrypt(config.apiKey, TransitKey.LLM_CONFIG);
  config.apiKey = ciphertext;
}

/**
 * Decrypt apiKey in config JSON field
 * Modifies instance in-place
 *
 * @param instance - LlmProviderInstance record
 * @param vaultService - Vault service
 * @param _logger - NestJS logger (unused, kept for consistency with account middleware)
 * @throws Error with user-facing message when decryption fails
 */
async function decryptConfig(
  instance: LlmProviderInstanceData,
  vaultService: VaultService,
  _logger: Logger
): Promise<void> {
  // Skip if no config field
  if (!instance.config || typeof instance.config !== "object") {
    return;
  }

  const config = instance.config;

  // Skip if no apiKey field
  if (!config.apiKey || typeof config.apiKey !== "string") {
    return;
  }

  // Only decrypt if encrypted (backward compatible with plaintext)
  if (!isEncrypted(config.apiKey)) {
    return;
  }

  // Decrypt ciphertext
  try {
    config.apiKey = await vaultService.decrypt(config.apiKey, TransitKey.LLM_CONFIG);
  } catch (error) {
    const errorMsg = error instanceof Error ? error.message : "Unknown error";
    throw new Error(
      `Failed to decrypt LLM provider configuration. Please re-enter the API key. Details: ${errorMsg}`
    );
  }
}

/**
 * Check if a value is encrypted
 *
 * @param value - String value to check
 * @returns true if value appears to be encrypted
 */
function isEncrypted(value: string): boolean {
  if (!value || typeof value !== "string") {
    return false;
  }

  // Vault format: vault:v1:...
  if (value.startsWith("vault:v1:")) {
    return true;
  }

  // AES format: iv:authTag:encrypted (3 colon-separated hex parts)
  // For future compatibility if AES fallback is used
  if (isAESEncrypted(value)) {
    return true;
  }

  return false;
}

/**
 * Check if a value is AES-256-GCM encrypted
 *
 * @param value - String value to check
 * @returns true if value is in AES format
 */
function isAESEncrypted(value: string): boolean {
  if (!value || typeof value !== "string") {
    return false;
  }

  // AES format: iv:authTag:encrypted (3 parts, all hex)
  const parts = value.split(":");
  if (parts.length !== 3) {
    return false;
  }

  // Verify all parts are hex strings
  return parts.every((part) => /^[0-9a-f]+$/i.test(part));
}