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

import { ConfigService } from "@nestjs/config";
import { Logger } from "@nestjs/common";
import { describe, it, expect, beforeEach, vi, afterEach } from "vitest";
import { DockerSandboxService, DEFAULT_ENV_WHITELIST } from "./docker-sandbox.service";
import Docker from "dockerode";

describe("DockerSandboxService", () => {
  let service: DockerSandboxService;
  let mockConfigService: ConfigService;
  let mockDocker: Docker;
  let mockContainer: Docker.Container;

  beforeEach(() => {
    // Create mock Docker container
    mockContainer = {
      id: "container-123",
      start: vi.fn().mockResolvedValue(undefined),
      stop: vi.fn().mockResolvedValue(undefined),
      remove: vi.fn().mockResolvedValue(undefined),
      inspect: vi.fn().mockResolvedValue({
        State: { Status: "running" },
      }),
    } as unknown as Docker.Container;

    // Create mock Docker instance
    mockDocker = {
      createContainer: vi.fn().mockResolvedValue(mockContainer),
      getContainer: vi.fn().mockReturnValue(mockContainer),
    } as unknown as Docker;

    // Create mock ConfigService
    mockConfigService = {
      get: vi.fn((key: string, defaultValue?: unknown) => {
        const config: Record<string, unknown> = {
          "orchestrator.docker.socketPath": "/var/run/docker.sock",
          "orchestrator.sandbox.enabled": true,
          "orchestrator.sandbox.defaultImage": "node:20-alpine",
          "orchestrator.sandbox.defaultMemoryMB": 512,
          "orchestrator.sandbox.defaultCpuLimit": 1.0,
          "orchestrator.sandbox.networkMode": "bridge",
        };
        return config[key] !== undefined ? config[key] : defaultValue;
      }),
    } as unknown as ConfigService;

    // Create service with mock Docker instance
    service = new DockerSandboxService(mockConfigService, mockDocker);
  });

  describe("constructor", () => {
    it("should be defined", () => {
      expect(service).toBeDefined();
    });

    it("should use provided Docker instance", () => {
      expect(service).toBeDefined();
      // Service should use the mockDocker instance we provided
    });
  });

  describe("createContainer", () => {
    it("should create a container with default configuration", async () => {
      const agentId = "agent-123";
      const taskId = "task-456";
      const workspacePath = "/workspace/agent-123";

      const result = await service.createContainer(agentId, taskId, workspacePath);

      expect(result.containerId).toBe("container-123");
      expect(result.agentId).toBe(agentId);
      expect(result.taskId).toBe(taskId);
      expect(result.createdAt).toBeInstanceOf(Date);
      expect(mockDocker.createContainer).toHaveBeenCalledWith({
        Image: "node:20-alpine",
        name: expect.stringContaining(`mosaic-agent-${agentId}`),
        User: "node:node",
        HostConfig: {
          Memory: 512 * 1024 * 1024, // 512MB in bytes
          NanoCpus: 1000000000, // 1.0 CPU
          NetworkMode: "bridge",
          Binds: [`${workspacePath}:/workspace`],
          AutoRemove: false,
          ReadonlyRootfs: false,
        },
        WorkingDir: "/workspace",
        Env: [`AGENT_ID=${agentId}`, `TASK_ID=${taskId}`],
      });
    });

    it("should create a container with custom resource limits", async () => {
      const agentId = "agent-123";
      const taskId = "task-456";
      const workspacePath = "/workspace/agent-123";
      const options = {
        memoryMB: 1024,
        cpuLimit: 2.0,
      };

      await service.createContainer(agentId, taskId, workspacePath, options);

      expect(mockDocker.createContainer).toHaveBeenCalledWith(
        expect.objectContaining({
          HostConfig: expect.objectContaining({
            Memory: 1024 * 1024 * 1024, // 1024MB in bytes
            NanoCpus: 2000000000, // 2.0 CPU
          }),
        })
      );
    });

    it("should create a container with network isolation", async () => {
      const agentId = "agent-123";
      const taskId = "task-456";
      const workspacePath = "/workspace/agent-123";
      const options = {
        networkMode: "none" as const,
      };

      await service.createContainer(agentId, taskId, workspacePath, options);

      expect(mockDocker.createContainer).toHaveBeenCalledWith(
        expect.objectContaining({
          HostConfig: expect.objectContaining({
            NetworkMode: "none",
          }),
        })
      );
    });

    it("should create a container with whitelisted environment variables", async () => {
      const agentId = "agent-123";
      const taskId = "task-456";
      const workspacePath = "/workspace/agent-123";
      const options = {
        env: {
          NODE_ENV: "production",
          LOG_LEVEL: "debug",
        },
      };

      await service.createContainer(agentId, taskId, workspacePath, options);

      expect(mockDocker.createContainer).toHaveBeenCalledWith(
        expect.objectContaining({
          Env: expect.arrayContaining([
            `AGENT_ID=${agentId}`,
            `TASK_ID=${taskId}`,
            "NODE_ENV=production",
            "LOG_LEVEL=debug",
          ]),
        })
      );
    });

    it("should throw error if container creation fails", async () => {
      const agentId = "agent-123";
      const taskId = "task-456";
      const workspacePath = "/workspace/agent-123";

      (mockDocker.createContainer as ReturnType<typeof vi.fn>).mockRejectedValue(
        new Error("Docker daemon not available")
      );

      await expect(service.createContainer(agentId, taskId, workspacePath)).rejects.toThrow(
        "Failed to create container for agent agent-123"
      );
    });
  });

  describe("startContainer", () => {
    it("should start a container by ID", async () => {
      const containerId = "container-123";

      await service.startContainer(containerId);

      expect(mockDocker.getContainer).toHaveBeenCalledWith(containerId);
      expect(mockContainer.start).toHaveBeenCalled();
    });

    it("should throw error if container start fails", async () => {
      const containerId = "container-123";

      (mockContainer.start as ReturnType<typeof vi.fn>).mockRejectedValue(
        new Error("Container not found")
      );

      await expect(service.startContainer(containerId)).rejects.toThrow(
        "Failed to start container container-123"
      );
    });
  });

  describe("stopContainer", () => {
    it("should stop a container by ID", async () => {
      const containerId = "container-123";

      await service.stopContainer(containerId);

      expect(mockDocker.getContainer).toHaveBeenCalledWith(containerId);
      expect(mockContainer.stop).toHaveBeenCalledWith({ t: 10 });
    });

    it("should stop a container with custom timeout", async () => {
      const containerId = "container-123";
      const timeout = 30;

      await service.stopContainer(containerId, timeout);

      expect(mockContainer.stop).toHaveBeenCalledWith({ t: timeout });
    });

    it("should throw error if container stop fails", async () => {
      const containerId = "container-123";

      (mockContainer.stop as ReturnType<typeof vi.fn>).mockRejectedValue(
        new Error("Container already stopped")
      );

      await expect(service.stopContainer(containerId)).rejects.toThrow(
        "Failed to stop container container-123"
      );
    });
  });

  describe("removeContainer", () => {
    it("should remove a container by ID", async () => {
      const containerId = "container-123";

      await service.removeContainer(containerId);

      expect(mockDocker.getContainer).toHaveBeenCalledWith(containerId);
      expect(mockContainer.remove).toHaveBeenCalledWith({ force: true });
    });

    it("should throw error if container removal fails", async () => {
      const containerId = "container-123";

      (mockContainer.remove as ReturnType<typeof vi.fn>).mockRejectedValue(
        new Error("Container not found")
      );

      await expect(service.removeContainer(containerId)).rejects.toThrow(
        "Failed to remove container container-123"
      );
    });
  });

  describe("getContainerStatus", () => {
    it("should return container status", async () => {
      const containerId = "container-123";

      const status = await service.getContainerStatus(containerId);

      expect(status).toBe("running");
      expect(mockDocker.getContainer).toHaveBeenCalledWith(containerId);
      expect(mockContainer.inspect).toHaveBeenCalled();
    });

    it("should throw error if container inspect fails", async () => {
      const containerId = "container-123";

      (mockContainer.inspect as ReturnType<typeof vi.fn>).mockRejectedValue(
        new Error("Container not found")
      );

      await expect(service.getContainerStatus(containerId)).rejects.toThrow(
        "Failed to get container status for container-123"
      );
    });
  });

  describe("cleanup", () => {
    it("should stop and remove container", async () => {
      const containerId = "container-123";

      await service.cleanup(containerId);

      expect(mockContainer.stop).toHaveBeenCalledWith({ t: 10 });
      expect(mockContainer.remove).toHaveBeenCalledWith({ force: true });
    });

    it("should remove container even if stop fails", async () => {
      const containerId = "container-123";

      (mockContainer.stop as ReturnType<typeof vi.fn>).mockRejectedValue(
        new Error("Container already stopped")
      );

      await service.cleanup(containerId);

      expect(mockContainer.remove).toHaveBeenCalledWith({ force: true });
    });

    it("should throw error if both stop and remove fail", async () => {
      const containerId = "container-123";

      (mockContainer.stop as ReturnType<typeof vi.fn>).mockRejectedValue(
        new Error("Container not found")
      );
      (mockContainer.remove as ReturnType<typeof vi.fn>).mockRejectedValue(
        new Error("Container not found")
      );

      await expect(service.cleanup(containerId)).rejects.toThrow(
        "Failed to cleanup container container-123"
      );
    });
  });

  describe("isEnabled", () => {
    it("should return true if sandbox is enabled in config", () => {
      expect(service.isEnabled()).toBe(true);
    });

    it("should return false if sandbox is disabled in config", () => {
      const disabledConfigService = {
        get: vi.fn((key: string, defaultValue?: unknown) => {
          const config: Record<string, unknown> = {
            "orchestrator.docker.socketPath": "/var/run/docker.sock",
            "orchestrator.sandbox.enabled": false,
            "orchestrator.sandbox.defaultImage": "node:20-alpine",
            "orchestrator.sandbox.defaultMemoryMB": 512,
            "orchestrator.sandbox.defaultCpuLimit": 1.0,
            "orchestrator.sandbox.networkMode": "bridge",
          };
          return config[key] !== undefined ? config[key] : defaultValue;
        }),
      } as unknown as ConfigService;

      const disabledService = new DockerSandboxService(disabledConfigService, mockDocker);

      expect(disabledService.isEnabled()).toBe(false);
    });
  });

  describe("security warning", () => {
    let warnSpy: ReturnType<typeof vi.spyOn>;

    beforeEach(() => {
      warnSpy = vi.spyOn(Logger.prototype, "warn").mockImplementation(() => undefined);
    });

    afterEach(() => {
      warnSpy.mockRestore();
    });

    it("should log security warning when sandbox is disabled", () => {
      const disabledConfigService = {
        get: vi.fn((key: string, defaultValue?: unknown) => {
          const config: Record<string, unknown> = {
            "orchestrator.docker.socketPath": "/var/run/docker.sock",
            "orchestrator.sandbox.enabled": false,
            "orchestrator.sandbox.defaultImage": "node:20-alpine",
            "orchestrator.sandbox.defaultMemoryMB": 512,
            "orchestrator.sandbox.defaultCpuLimit": 1.0,
            "orchestrator.sandbox.networkMode": "bridge",
          };
          return config[key] !== undefined ? config[key] : defaultValue;
        }),
      } as unknown as ConfigService;

      new DockerSandboxService(disabledConfigService, mockDocker);

      expect(warnSpy).toHaveBeenCalledWith(
        "SECURITY WARNING: Docker sandbox is DISABLED. Agents will run directly on the host without container isolation."
      );
    });

    it("should not log security warning when sandbox is enabled", () => {
      // Use the default mockConfigService which has sandbox enabled
      new DockerSandboxService(mockConfigService, mockDocker);

      expect(warnSpy).not.toHaveBeenCalledWith(expect.stringContaining("SECURITY WARNING"));
    });
  });

  describe("environment variable whitelist", () => {
    describe("getEnvWhitelist", () => {
      it("should return default whitelist when no custom whitelist is configured", () => {
        const whitelist = service.getEnvWhitelist();

        expect(whitelist).toEqual(DEFAULT_ENV_WHITELIST);
        expect(whitelist).toContain("AGENT_ID");
        expect(whitelist).toContain("TASK_ID");
        expect(whitelist).toContain("NODE_ENV");
        expect(whitelist).toContain("LOG_LEVEL");
      });

      it("should return custom whitelist when configured", () => {
        const customWhitelist = ["CUSTOM_VAR_1", "CUSTOM_VAR_2"];
        const customConfigService = {
          get: vi.fn((key: string, defaultValue?: unknown) => {
            const config: Record<string, unknown> = {
              "orchestrator.docker.socketPath": "/var/run/docker.sock",
              "orchestrator.sandbox.enabled": true,
              "orchestrator.sandbox.defaultImage": "node:20-alpine",
              "orchestrator.sandbox.defaultMemoryMB": 512,
              "orchestrator.sandbox.defaultCpuLimit": 1.0,
              "orchestrator.sandbox.networkMode": "bridge",
              "orchestrator.sandbox.envWhitelist": customWhitelist,
            };
            return config[key] !== undefined ? config[key] : defaultValue;
          }),
        } as unknown as ConfigService;

        const customService = new DockerSandboxService(customConfigService, mockDocker);
        const whitelist = customService.getEnvWhitelist();

        expect(whitelist).toEqual(customWhitelist);
      });
    });

    describe("filterEnvVars", () => {
      it("should allow whitelisted environment variables", () => {
        const envVars = {
          NODE_ENV: "production",
          LOG_LEVEL: "debug",
          TZ: "UTC",
        };

        const result = service.filterEnvVars(envVars);

        expect(result.allowed).toEqual({
          NODE_ENV: "production",
          LOG_LEVEL: "debug",
          TZ: "UTC",
        });
        expect(result.filtered).toEqual([]);
      });

      it("should filter non-whitelisted environment variables", () => {
        const envVars = {
          NODE_ENV: "production",
          DATABASE_URL: "postgres://secret@host/db",
          API_KEY: "sk-secret-key",
          AWS_SECRET_ACCESS_KEY: "super-secret",
        };

        const result = service.filterEnvVars(envVars);

        expect(result.allowed).toEqual({
          NODE_ENV: "production",
        });
        expect(result.filtered).toContain("DATABASE_URL");
        expect(result.filtered).toContain("API_KEY");
        expect(result.filtered).toContain("AWS_SECRET_ACCESS_KEY");
        expect(result.filtered).toHaveLength(3);
      });

      it("should handle empty env vars object", () => {
        const result = service.filterEnvVars({});

        expect(result.allowed).toEqual({});
        expect(result.filtered).toEqual([]);
      });

      it("should handle all vars being filtered", () => {
        const envVars = {
          SECRET_KEY: "secret",
          PASSWORD: "password123",
          PRIVATE_TOKEN: "token",
        };

        const result = service.filterEnvVars(envVars);

        expect(result.allowed).toEqual({});
        expect(result.filtered).toEqual(["SECRET_KEY", "PASSWORD", "PRIVATE_TOKEN"]);
      });
    });

    describe("createContainer with filtering", () => {
      let warnSpy: ReturnType<typeof vi.spyOn>;

      beforeEach(() => {
        warnSpy = vi.spyOn(Logger.prototype, "warn").mockImplementation(() => undefined);
      });

      afterEach(() => {
        warnSpy.mockRestore();
      });

      it("should filter non-whitelisted vars and only pass allowed vars to container", async () => {
        const agentId = "agent-123";
        const taskId = "task-456";
        const workspacePath = "/workspace/agent-123";
        const options = {
          env: {
            NODE_ENV: "production",
            DATABASE_URL: "postgres://secret@host/db",
            LOG_LEVEL: "info",
          },
        };

        await service.createContainer(agentId, taskId, workspacePath, options);

        // Should include whitelisted vars
        expect(mockDocker.createContainer).toHaveBeenCalledWith(
          expect.objectContaining({
            Env: expect.arrayContaining([
              `AGENT_ID=${agentId}`,
              `TASK_ID=${taskId}`,
              "NODE_ENV=production",
              "LOG_LEVEL=info",
            ]),
          })
        );

        // Should NOT include filtered vars
        const callArgs = (mockDocker.createContainer as ReturnType<typeof vi.fn>).mock.calls[0][0];
        expect(callArgs.Env).not.toContain("DATABASE_URL=postgres://secret@host/db");
      });

      it("should log warning when env vars are filtered", async () => {
        const agentId = "agent-123";
        const taskId = "task-456";
        const workspacePath = "/workspace/agent-123";
        const options = {
          env: {
            DATABASE_URL: "postgres://secret@host/db",
            API_KEY: "sk-secret",
          },
        };

        await service.createContainer(agentId, taskId, workspacePath, options);

        expect(warnSpy).toHaveBeenCalledWith(
          expect.stringContaining("SECURITY: Filtered 2 non-whitelisted env var(s)")
        );
        expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("DATABASE_URL"));
        expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("API_KEY"));
      });

      it("should not log warning when all vars are whitelisted", async () => {
        const agentId = "agent-123";
        const taskId = "task-456";
        const workspacePath = "/workspace/agent-123";
        const options = {
          env: {
            NODE_ENV: "production",
            LOG_LEVEL: "debug",
          },
        };

        await service.createContainer(agentId, taskId, workspacePath, options);

        expect(warnSpy).not.toHaveBeenCalledWith(expect.stringContaining("SECURITY: Filtered"));
      });

      it("should not log warning when no env vars are provided", async () => {
        const agentId = "agent-123";
        const taskId = "task-456";
        const workspacePath = "/workspace/agent-123";

        await service.createContainer(agentId, taskId, workspacePath);

        expect(warnSpy).not.toHaveBeenCalledWith(expect.stringContaining("SECURITY: Filtered"));
      });
    });
  });

  describe("DEFAULT_ENV_WHITELIST", () => {
    it("should contain essential agent identification vars", () => {
      expect(DEFAULT_ENV_WHITELIST).toContain("AGENT_ID");
      expect(DEFAULT_ENV_WHITELIST).toContain("TASK_ID");
    });

    it("should contain Node.js runtime vars", () => {
      expect(DEFAULT_ENV_WHITELIST).toContain("NODE_ENV");
      expect(DEFAULT_ENV_WHITELIST).toContain("NODE_OPTIONS");
    });

    it("should contain logging vars", () => {
      expect(DEFAULT_ENV_WHITELIST).toContain("LOG_LEVEL");
      expect(DEFAULT_ENV_WHITELIST).toContain("DEBUG");
    });

    it("should contain locale vars", () => {
      expect(DEFAULT_ENV_WHITELIST).toContain("LANG");
      expect(DEFAULT_ENV_WHITELIST).toContain("LC_ALL");
      expect(DEFAULT_ENV_WHITELIST).toContain("TZ");
    });

    it("should contain Mosaic-specific safe vars", () => {
      expect(DEFAULT_ENV_WHITELIST).toContain("MOSAIC_WORKSPACE_ID");
      expect(DEFAULT_ENV_WHITELIST).toContain("MOSAIC_PROJECT_ID");
      expect(DEFAULT_ENV_WHITELIST).toContain("MOSAIC_AGENT_TYPE");
    });

    it("should NOT contain sensitive var patterns", () => {
      // Verify common sensitive vars are not in the whitelist
      expect(DEFAULT_ENV_WHITELIST).not.toContain("DATABASE_URL");
      expect(DEFAULT_ENV_WHITELIST).not.toContain("API_KEY");
      expect(DEFAULT_ENV_WHITELIST).not.toContain("SECRET");
      expect(DEFAULT_ENV_WHITELIST).not.toContain("PASSWORD");
      expect(DEFAULT_ENV_WHITELIST).not.toContain("AWS_SECRET_ACCESS_KEY");
      expect(DEFAULT_ENV_WHITELIST).not.toContain("ANTHROPIC_API_KEY");
    });
  });
});