stack/apps/orchestrator/src/spawner/docker-sandbox.service.ts

import { Injectable, Logger } from "@nestjs/common";
import { ConfigService } from "@nestjs/config";
import Docker from "dockerode";
import {
  DockerSandboxOptions,
  ContainerCreateResult,
  DockerSecurityOptions,
  LinuxCapability,
} from "./types/docker-sandbox.types";

/**
 * Maximum allowed length for a Docker image reference.
 * Docker image names rarely exceed 128 characters; 256 provides generous headroom.
 */
export const MAX_IMAGE_TAG_LENGTH = 256;

/**
 * Regex pattern for validating Docker image tag references.
 * Allows: registry/namespace/image:tag or image@sha256:digest
 * Valid characters: alphanumeric, dots, hyphens, underscores, forward slashes, colons, and @.
 * Blocks shell metacharacters (;, &, |, $, backtick, spaces, newlines, etc.) to prevent injection.
 *
 * Uses a simple character-class approach (no alternation or nested quantifiers)
 * to avoid catastrophic backtracking.
 */
export const DOCKER_IMAGE_TAG_PATTERN = /^[a-zA-Z0-9][a-zA-Z0-9./_:@-]*$/;

/**
 * Default whitelist of allowed environment variable names/patterns for Docker containers.
 * Only these variables will be passed to spawned agent containers.
 * This prevents accidental leakage of secrets like API keys, database credentials, etc.
 */
export const DEFAULT_ENV_WHITELIST: readonly string[] = [
  // Agent identification
  "AGENT_ID",
  "TASK_ID",
  // Node.js runtime
  "NODE_ENV",
  "NODE_OPTIONS",
  // Logging
  "LOG_LEVEL",
  "DEBUG",
  // Locale
  "LANG",
  "LC_ALL",
  "TZ",
  // Application-specific safe vars
  "MOSAIC_WORKSPACE_ID",
  "MOSAIC_PROJECT_ID",
  "MOSAIC_AGENT_TYPE",
] as const;

/**
 * Default security hardening options for Docker containers.
 * These settings follow security best practices:
 * - Drop all Linux capabilities (principle of least privilege)
 * - Read-only root filesystem (agents write to mounted /workspace volume)
 * - PID limit to prevent fork bombs
 * - No new privileges to prevent privilege escalation
 */
export const DEFAULT_SECURITY_OPTIONS: Required<DockerSecurityOptions> = {
  capDrop: ["ALL"],
  capAdd: [],
  readonlyRootfs: true,
  pidsLimit: 100,
  noNewPrivileges: true,
} as const;

/**
 * Service for managing Docker container isolation for agents
 * Provides secure sandboxing with resource limits and cleanup
 */
@Injectable()
export class DockerSandboxService {
  private readonly logger = new Logger(DockerSandboxService.name);
  private readonly docker: Docker;
  private readonly sandboxEnabled: boolean;
  private readonly defaultImage: string;
  private readonly defaultMemoryMB: number;
  private readonly defaultCpuLimit: number;
  private readonly defaultNetworkMode: string;
  private readonly envWhitelist: readonly string[];
  private readonly defaultSecurityOptions: Required<DockerSecurityOptions>;

  constructor(
    private readonly configService: ConfigService,
    docker?: Docker
  ) {
    const socketPath = this.configService.get<string>(
      "orchestrator.docker.socketPath",
      "/var/run/docker.sock"
    );

    this.docker = docker ?? new Docker({ socketPath });

    this.sandboxEnabled = this.configService.get<boolean>("orchestrator.sandbox.enabled", false);

    this.defaultImage = this.configService.get<string>(
      "orchestrator.sandbox.defaultImage",
      "node:20-alpine"
    );

    this.defaultMemoryMB = this.configService.get<number>(
      "orchestrator.sandbox.defaultMemoryMB",
      512
    );

    this.defaultCpuLimit = this.configService.get<number>(
      "orchestrator.sandbox.defaultCpuLimit",
      1.0
    );

    this.defaultNetworkMode = this.configService.get<string>(
      "orchestrator.sandbox.networkMode",
      "bridge"
    );

    // Load custom whitelist from config, or use defaults
    const customWhitelist = this.configService.get<string[]>("orchestrator.sandbox.envWhitelist");
    this.envWhitelist = customWhitelist ?? DEFAULT_ENV_WHITELIST;

    // Load security options from config, merging with secure defaults
    const configCapDrop = this.configService.get<LinuxCapability[]>(
      "orchestrator.sandbox.security.capDrop"
    );
    const configCapAdd = this.configService.get<LinuxCapability[]>(
      "orchestrator.sandbox.security.capAdd"
    );
    const configReadonlyRootfs = this.configService.get<boolean>(
      "orchestrator.sandbox.security.readonlyRootfs"
    );
    const configPidsLimit = this.configService.get<number>(
      "orchestrator.sandbox.security.pidsLimit"
    );
    const configNoNewPrivileges = this.configService.get<boolean>(
      "orchestrator.sandbox.security.noNewPrivileges"
    );

    this.defaultSecurityOptions = {
      capDrop: configCapDrop ?? DEFAULT_SECURITY_OPTIONS.capDrop,
      capAdd: configCapAdd ?? DEFAULT_SECURITY_OPTIONS.capAdd,
      readonlyRootfs: configReadonlyRootfs ?? DEFAULT_SECURITY_OPTIONS.readonlyRootfs,
      pidsLimit: configPidsLimit ?? DEFAULT_SECURITY_OPTIONS.pidsLimit,
      noNewPrivileges: configNoNewPrivileges ?? DEFAULT_SECURITY_OPTIONS.noNewPrivileges,
    };

    // Validate default image tag at startup to fail fast on misconfiguration
    this.validateImageTag(this.defaultImage);

    this.logger.log(
      `DockerSandboxService initialized (enabled: ${this.sandboxEnabled.toString()}, socket: ${socketPath})`
    );
    this.logger.log(
      `Security hardening: capDrop=${this.defaultSecurityOptions.capDrop.join(",") || "none"}, ` +
        `readonlyRootfs=${this.defaultSecurityOptions.readonlyRootfs.toString()}, ` +
        `pidsLimit=${this.defaultSecurityOptions.pidsLimit.toString()}, ` +
        `noNewPrivileges=${this.defaultSecurityOptions.noNewPrivileges.toString()}`
    );

    if (!this.sandboxEnabled) {
      this.logger.warn(
        "SECURITY WARNING: Docker sandbox is DISABLED. Agents will run directly on the host without container isolation."
      );
    }
  }

  /**
   * Validate a Docker image tag reference.
   * Ensures the image tag only contains safe characters and is within length limits.
   * Blocks shell metacharacters and suspicious patterns to prevent injection attacks.
   * @param imageTag The Docker image tag to validate
   * @throws Error if the image tag is invalid
   */
  validateImageTag(imageTag: string): void {
    if (!imageTag || imageTag.trim().length === 0) {
      throw new Error("Docker image tag must not be empty");
    }

    if (imageTag.length > MAX_IMAGE_TAG_LENGTH) {
      throw new Error(
        `Docker image tag exceeds maximum length of ${MAX_IMAGE_TAG_LENGTH.toString()} characters`
      );
    }

    if (!DOCKER_IMAGE_TAG_PATTERN.test(imageTag)) {
      throw new Error(
        `Docker image tag contains invalid characters: "${imageTag}". ` +
          "Only alphanumeric characters, dots, hyphens, underscores, forward slashes, colons, and sha256 digests are allowed."
      );
    }
  }

  /**
   * Create a Docker container for agent isolation
   * @param agentId Unique agent identifier
   * @param taskId Task identifier
   * @param workspacePath Path to workspace directory to mount
   * @param options Optional container configuration
   * @returns Container creation result
   */
  async createContainer(
    agentId: string,
    taskId: string,
    workspacePath: string,
    options?: DockerSandboxOptions
  ): Promise<ContainerCreateResult> {
    try {
      const image = options?.image ?? this.defaultImage;

      // Validate image tag format before any Docker operations
      this.validateImageTag(image);

      const memoryMB = options?.memoryMB ?? this.defaultMemoryMB;
      const cpuLimit = options?.cpuLimit ?? this.defaultCpuLimit;
      const networkMode = options?.networkMode ?? this.defaultNetworkMode;

      // Merge security options with defaults
      const security = this.mergeSecurityOptions(options?.security);

      // Convert memory from MB to bytes
      const memoryBytes = memoryMB * 1024 * 1024;

      // Convert CPU limit to NanoCPUs (1.0 = 1,000,000,000 nanocpus)
      const nanoCpus = Math.floor(cpuLimit * 1000000000);

      // Build environment variables with whitelist filtering
      const env = [`AGENT_ID=${agentId}`, `TASK_ID=${taskId}`];

      if (options?.env) {
        const { allowed, filtered } = this.filterEnvVars(options.env);

        // Add allowed vars
        Object.entries(allowed).forEach(([key, value]) => {
          env.push(`${key}=${value}`);
        });

        // Log warning for filtered vars
        if (filtered.length > 0) {
          this.logger.warn(
            `SECURITY: Filtered ${filtered.length.toString()} non-whitelisted env var(s) for agent ${agentId}: ${filtered.join(", ")}`
          );
        }
      }

      // Container name with timestamp to ensure uniqueness
      const containerName = `mosaic-agent-${agentId}-${Date.now().toString()}`;

      this.logger.log(
        `Creating container for agent ${agentId} (image: ${image}, memory: ${memoryMB.toString()}MB, cpu: ${cpuLimit.toString()})`
      );

      // Build HostConfig with security hardening
      const hostConfig: Docker.HostConfig = {
        Memory: memoryBytes,
        NanoCpus: nanoCpus,
        NetworkMode: networkMode,
        Binds: [`${workspacePath}:/workspace`],
        AutoRemove: false, // Manual cleanup for audit trail
        ReadonlyRootfs: security.readonlyRootfs,
        PidsLimit: security.pidsLimit > 0 ? security.pidsLimit : undefined,
        SecurityOpt: security.noNewPrivileges ? ["no-new-privileges:true"] : undefined,
      };

      // Add capability dropping if configured
      if (security.capDrop.length > 0) {
        hostConfig.CapDrop = security.capDrop;
      }

      // Add capabilities back if configured (useful when dropping ALL first)
      if (security.capAdd.length > 0) {
        hostConfig.CapAdd = security.capAdd;
      }

      const container = await this.docker.createContainer({
        Image: image,
        name: containerName,
        User: "node:node", // Non-root user for security
        HostConfig: hostConfig,
        WorkingDir: "/workspace",
        Env: env,
      });

      const createdAt = new Date();

      this.logger.log(`Container created successfully: ${container.id} for agent ${agentId}`);

      return {
        containerId: container.id,
        agentId,
        taskId,
        createdAt,
      };
    } catch (error) {
      const enhancedError = error instanceof Error ? error : new Error(String(error));
      enhancedError.message = `Failed to create container for agent ${agentId}: ${enhancedError.message}`;
      this.logger.error(enhancedError.message, enhancedError);
      throw enhancedError;
    }
  }

  /**
   * Start a Docker container
   * @param containerId Container ID to start
   */
  async startContainer(containerId: string): Promise<void> {
    try {
      this.logger.log(`Starting container: ${containerId}`);
      const container = this.docker.getContainer(containerId);
      await container.start();
      this.logger.log(`Container started successfully: ${containerId}`);
    } catch (error) {
      const enhancedError = error instanceof Error ? error : new Error(String(error));
      enhancedError.message = `Failed to start container ${containerId}: ${enhancedError.message}`;
      this.logger.error(enhancedError.message, enhancedError);
      throw enhancedError;
    }
  }

  /**
   * Stop a Docker container
   * @param containerId Container ID to stop
   * @param timeout Timeout in seconds (default: 10)
   */
  async stopContainer(containerId: string, timeout = 10): Promise<void> {
    try {
      this.logger.log(`Stopping container: ${containerId} (timeout: ${timeout.toString()}s)`);
      const container = this.docker.getContainer(containerId);
      await container.stop({ t: timeout });
      this.logger.log(`Container stopped successfully: ${containerId}`);
    } catch (error) {
      const enhancedError = error instanceof Error ? error : new Error(String(error));
      enhancedError.message = `Failed to stop container ${containerId}: ${enhancedError.message}`;
      this.logger.error(enhancedError.message, enhancedError);
      throw enhancedError;
    }
  }

  /**
   * Remove a Docker container
   * @param containerId Container ID to remove
   */
  async removeContainer(containerId: string): Promise<void> {
    try {
      this.logger.log(`Removing container: ${containerId}`);
      const container = this.docker.getContainer(containerId);
      await container.remove({ force: true });
      this.logger.log(`Container removed successfully: ${containerId}`);
    } catch (error) {
      const enhancedError = error instanceof Error ? error : new Error(String(error));
      enhancedError.message = `Failed to remove container ${containerId}: ${enhancedError.message}`;
      this.logger.error(enhancedError.message, enhancedError);
      throw enhancedError;
    }
  }

  /**
   * Get container status
   * @param containerId Container ID to inspect
   * @returns Container status string
   */
  async getContainerStatus(containerId: string): Promise<string> {
    try {
      const container = this.docker.getContainer(containerId);
      const info = await container.inspect();
      return info.State.Status;
    } catch (error) {
      const enhancedError = error instanceof Error ? error : new Error(String(error));
      enhancedError.message = `Failed to get container status for ${containerId}: ${enhancedError.message}`;
      this.logger.error(enhancedError.message, enhancedError);
      throw enhancedError;
    }
  }

  /**
   * Cleanup container (stop and remove)
   * @param containerId Container ID to cleanup
   */
  async cleanup(containerId: string): Promise<void> {
    this.logger.log(`Cleaning up container: ${containerId}`);

    try {
      // Try to stop first
      await this.stopContainer(containerId);
    } catch (error) {
      this.logger.warn(
        `Failed to stop container ${containerId} during cleanup (may already be stopped): ${error instanceof Error ? error.message : String(error)}`
      );
    }

    try {
      // Always try to remove
      await this.removeContainer(containerId);
    } catch (error) {
      const enhancedError = error instanceof Error ? error : new Error(String(error));
      enhancedError.message = `Failed to cleanup container ${containerId}: ${enhancedError.message}`;
      this.logger.error(enhancedError.message, enhancedError);
      throw enhancedError;
    }

    this.logger.log(`Container cleanup completed: ${containerId}`);
  }

  /**
   * Check if sandbox mode is enabled
   * @returns True if sandbox is enabled
   */
  isEnabled(): boolean {
    return this.sandboxEnabled;
  }

  /**
   * Get the current environment variable whitelist
   * @returns The configured whitelist of allowed env var names
   */
  getEnvWhitelist(): readonly string[] {
    return this.envWhitelist;
  }

  /**
   * Filter environment variables against the whitelist
   * @param envVars Object of environment variables to filter
   * @returns Object with allowed vars and array of filtered var names
   */
  filterEnvVars(envVars: Record<string, string>): {
    allowed: Record<string, string>;
    filtered: string[];
  } {
    const allowed: Record<string, string> = {};
    const filtered: string[] = [];

    for (const [key, value] of Object.entries(envVars)) {
      if (this.isEnvVarAllowed(key)) {
        allowed[key] = value;
      } else {
        filtered.push(key);
      }
    }

    return { allowed, filtered };
  }

  /**
   * Check if an environment variable name is allowed by the whitelist
   * @param varName Environment variable name to check
   * @returns True if allowed
   */
  private isEnvVarAllowed(varName: string): boolean {
    return this.envWhitelist.includes(varName);
  }

  /**
   * Get the current default security options
   * @returns The configured security options
   */
  getSecurityOptions(): Required<DockerSecurityOptions> {
    return { ...this.defaultSecurityOptions };
  }

  /**
   * Merge provided security options with defaults
   * @param options Optional security options to merge
   * @returns Complete security options with all fields
   */
  private mergeSecurityOptions(options?: DockerSecurityOptions): Required<DockerSecurityOptions> {
    if (!options) {
      return { ...this.defaultSecurityOptions };
    }

    return {
      capDrop: options.capDrop ?? this.defaultSecurityOptions.capDrop,
      capAdd: options.capAdd ?? this.defaultSecurityOptions.capAdd,
      readonlyRootfs: options.readonlyRootfs ?? this.defaultSecurityOptions.readonlyRootfs,
      pidsLimit: options.pidsLimit ?? this.defaultSecurityOptions.pidsLimit,
      noNewPrivileges: options.noNewPrivileges ?? this.defaultSecurityOptions.noNewPrivileges,
    };
  }
}