stack/apps/api/src/federation/federation.service.spec.ts

/**
 * Federation Service Tests
 */

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { FederationService } from "./federation.service";
import { CryptoService } from "./crypto.service";
import { PrismaService } from "../prisma/prisma.service";
import { ConfigService } from "@nestjs/config";
import { Instance } from "@prisma/client";

describe("FederationService", () => {
  let service: FederationService;
  let prismaService: PrismaService;
  let configService: ConfigService;
  let cryptoService: CryptoService;

  // Mock encrypted private key (simulates encrypted storage)
  const mockEncryptedPrivateKey = "iv:authTag:encryptedData";
  const mockDecryptedPrivateKey = "-----BEGIN PRIVATE KEY-----\nMOCK\n-----END PRIVATE KEY-----";

  const mockInstance: Instance = {
    id: "123e4567-e89b-12d3-a456-426614174000",
    instanceId: "test-instance-id",
    name: "Test Instance",
    url: "https://test.example.com",
    publicKey: "-----BEGIN PUBLIC KEY-----\nMOCK\n-----END PUBLIC KEY-----",
    privateKey: mockEncryptedPrivateKey, // Stored encrypted
    capabilities: {
      supportsQuery: true,
      supportsCommand: true,
      supportsEvent: true,
      protocolVersion: "1.0",
    },
    metadata: {},
    createdAt: new Date("2026-01-01T00:00:00Z"),
    updatedAt: new Date("2026-01-01T00:00:00Z"),
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [
        FederationService,
        {
          provide: PrismaService,
          useValue: {
            instance: {
              findFirst: vi.fn(),
              create: vi.fn(),
              update: vi.fn(),
            },
          },
        },
        {
          provide: ConfigService,
          useValue: {
            get: vi.fn((key: string) => {
              const config: Record<string, string> = {
                INSTANCE_NAME: "Test Instance",
                INSTANCE_URL: "https://test.example.com",
                ENCRYPTION_KEY: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
              };
              return config[key];
            }),
          },
        },
        {
          provide: CryptoService,
          useValue: {
            encrypt: vi.fn((data: string) => mockEncryptedPrivateKey),
            decrypt: vi.fn((encrypted: string) => mockDecryptedPrivateKey),
          },
        },
      ],
    }).compile();

    service = module.get<FederationService>(FederationService);
    prismaService = module.get<PrismaService>(PrismaService);
    configService = module.get<ConfigService>(ConfigService);
    cryptoService = module.get<CryptoService>(CryptoService);
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  describe("getInstanceIdentity", () => {
    it("should return existing instance identity if found", async () => {
      // Arrange
      vi.spyOn(prismaService.instance, "findFirst").mockResolvedValue(mockInstance);

      // Act
      const result = await service.getInstanceIdentity();

      // Assert
      expect(result.privateKey).toEqual(mockDecryptedPrivateKey); // Decrypted
      expect(cryptoService.decrypt).toHaveBeenCalledWith(mockEncryptedPrivateKey);
      expect(prismaService.instance.findFirst).toHaveBeenCalledTimes(1);
    });

    it("should create new instance identity if not found", async () => {
      // Arrange
      vi.spyOn(prismaService.instance, "findFirst").mockResolvedValue(null);
      vi.spyOn(prismaService.instance, "create").mockResolvedValue(mockInstance);
      vi.spyOn(service, "generateKeypair").mockReturnValue({
        publicKey: mockInstance.publicKey,
        privateKey: mockDecryptedPrivateKey,
      });

      // Act
      const result = await service.getInstanceIdentity();

      // Assert
      expect(result.privateKey).toEqual(mockDecryptedPrivateKey);
      expect(cryptoService.encrypt).toHaveBeenCalled(); // Private key encrypted before storage
      expect(prismaService.instance.findFirst).toHaveBeenCalledTimes(1);
      expect(service.generateKeypair).toHaveBeenCalledTimes(1);
      expect(prismaService.instance.create).toHaveBeenCalledTimes(1);
    });

    it("should use config values for instance name and URL", async () => {
      // Arrange
      vi.spyOn(prismaService.instance, "findFirst").mockResolvedValue(null);
      vi.spyOn(prismaService.instance, "create").mockResolvedValue(mockInstance);
      vi.spyOn(service, "generateKeypair").mockReturnValue({
        publicKey: mockInstance.publicKey,
        privateKey: mockDecryptedPrivateKey,
      });

      // Act
      await service.getInstanceIdentity();

      // Assert
      expect(configService.get).toHaveBeenCalledWith("INSTANCE_NAME");
      expect(configService.get).toHaveBeenCalledWith("INSTANCE_URL");
    });

    it("should throw error for invalid URL", async () => {
      // Arrange
      vi.spyOn(prismaService.instance, "findFirst").mockResolvedValue(null);
      vi.spyOn(configService, "get").mockImplementation((key: string) => {
        if (key === "INSTANCE_URL") return "invalid-url";
        return "Test Instance";
      });
      vi.spyOn(service, "generateKeypair").mockReturnValue({
        publicKey: mockInstance.publicKey,
        privateKey: mockDecryptedPrivateKey,
      });

      // Act & Assert
      await expect(service.getInstanceIdentity()).rejects.toThrow("Invalid INSTANCE_URL");
    });
  });

  describe("getPublicIdentity", () => {
    it("should return instance identity without private key", async () => {
      // Arrange
      vi.spyOn(service, "getInstanceIdentity").mockResolvedValue(mockInstance);

      // Act
      const result = await service.getPublicIdentity();

      // Assert
      expect(result).toEqual({
        id: mockInstance.id,
        instanceId: mockInstance.instanceId,
        name: mockInstance.name,
        url: mockInstance.url,
        publicKey: mockInstance.publicKey,
        capabilities: mockInstance.capabilities,
        metadata: mockInstance.metadata,
        createdAt: mockInstance.createdAt,
        updatedAt: mockInstance.updatedAt,
      });
      expect(result).not.toHaveProperty("privateKey");
    });
  });

  describe("generateKeypair", () => {
    it("should generate valid RSA key pair", { timeout: 30000 }, () => {
      // Act
      const result = service.generateKeypair();

      // Assert
      expect(result).toHaveProperty("publicKey");
      expect(result).toHaveProperty("privateKey");
      expect(result.publicKey).toContain("BEGIN PUBLIC KEY");
      expect(result.privateKey).toContain("BEGIN PRIVATE KEY");
    });

    it("should generate different key pairs on each call", { timeout: 60000 }, () => {
      // Act
      const result1 = service.generateKeypair();
      const result2 = service.generateKeypair();

      // Assert
      expect(result1.publicKey).not.toEqual(result2.publicKey);
      expect(result1.privateKey).not.toEqual(result2.privateKey);
    });

    it("should generate RSA-4096 key pairs for future-proof security", { timeout: 30000 }, () => {
      // Act
      const result = service.generateKeypair();

      // Assert - Verify key size by checking approximate length
      // RSA-4096 keys are significantly larger than RSA-2048
      // Private key in PKCS8 format: RSA-2048 ≈ 1700 bytes, RSA-4096 ≈ 3200 bytes
      // Public key in SPKI format: RSA-2048 ≈ 400 bytes, RSA-4096 ≈ 800 bytes
      expect(result.privateKey.length).toBeGreaterThan(3000);
      expect(result.publicKey.length).toBeGreaterThan(700);
    });
  });

  describe("regenerateKeypair", () => {
    it("should generate new keypair and update instance", async () => {
      // Arrange
      const updatedInstance = { ...mockInstance };
      vi.spyOn(service, "getInstanceIdentity").mockResolvedValue({
        ...mockInstance,
        privateKey: mockDecryptedPrivateKey,
      });
      vi.spyOn(service, "generateKeypair").mockReturnValue({
        publicKey: "NEW_PUBLIC_KEY",
        privateKey: "NEW_PRIVATE_KEY",
      });
      vi.spyOn(prismaService.instance, "update").mockResolvedValue(updatedInstance);

      // Act
      const result = await service.regenerateKeypair();

      // Assert
      expect(service.generateKeypair).toHaveBeenCalledTimes(1);
      expect(cryptoService.encrypt).toHaveBeenCalledWith("NEW_PRIVATE_KEY"); // Encrypted before storage
      expect(prismaService.instance.update).toHaveBeenCalled();

      // SECURITY FIX: Verify private key is NOT in response
      expect(result).not.toHaveProperty("privateKey");
      expect(result).toHaveProperty("publicKey");
      expect(result).toHaveProperty("instanceId");
    });
  });

  describe("updateInstanceConfiguration", () => {
    it("should update instance name", async () => {
      // Arrange
      const updatedInstance = { ...mockInstance, name: "Updated Instance" };
      vi.spyOn(service, "getInstanceIdentity").mockResolvedValue({
        ...mockInstance,
        privateKey: mockDecryptedPrivateKey,
      });
      vi.spyOn(prismaService.instance, "update").mockResolvedValue(updatedInstance);

      // Act
      const result = await service.updateInstanceConfiguration({ name: "Updated Instance" });

      // Assert
      expect(prismaService.instance.update).toHaveBeenCalledWith({
        where: { id: mockInstance.id },
        data: { name: "Updated Instance" },
      });
      expect(result.name).toBe("Updated Instance");
      expect(result).not.toHaveProperty("privateKey");
    });

    it("should update instance capabilities", async () => {
      // Arrange
      const newCapabilities = {
        supportsQuery: true,
        supportsCommand: false,
        supportsEvent: true,
        supportsAgentSpawn: false,
        protocolVersion: "1.0",
      };
      const updatedInstance = { ...mockInstance, capabilities: newCapabilities };
      vi.spyOn(service, "getInstanceIdentity").mockResolvedValue({
        ...mockInstance,
        privateKey: mockDecryptedPrivateKey,
      });
      vi.spyOn(prismaService.instance, "update").mockResolvedValue(updatedInstance);

      // Act
      const result = await service.updateInstanceConfiguration({ capabilities: newCapabilities });

      // Assert
      expect(prismaService.instance.update).toHaveBeenCalledWith({
        where: { id: mockInstance.id },
        data: { capabilities: newCapabilities },
      });
      expect(result.capabilities).toEqual(newCapabilities);
    });

    it("should update instance metadata", async () => {
      // Arrange
      const newMetadata = { description: "Test description", region: "us-west-2" };
      const updatedInstance = { ...mockInstance, metadata: newMetadata };
      vi.spyOn(service, "getInstanceIdentity").mockResolvedValue({
        ...mockInstance,
        privateKey: mockDecryptedPrivateKey,
      });
      vi.spyOn(prismaService.instance, "update").mockResolvedValue(updatedInstance);

      // Act
      const result = await service.updateInstanceConfiguration({ metadata: newMetadata });

      // Assert
      expect(prismaService.instance.update).toHaveBeenCalledWith({
        where: { id: mockInstance.id },
        data: { metadata: newMetadata },
      });
      expect(result.metadata).toEqual(newMetadata);
    });

    it("should update multiple fields at once", async () => {
      // Arrange
      const updates = {
        name: "Updated Instance",
        capabilities: {
          supportsQuery: false,
          supportsCommand: false,
          supportsEvent: false,
          supportsAgentSpawn: false,
          protocolVersion: "1.0",
        },
        metadata: { description: "Updated" },
      };
      const updatedInstance = { ...mockInstance, ...updates };
      vi.spyOn(service, "getInstanceIdentity").mockResolvedValue({
        ...mockInstance,
        privateKey: mockDecryptedPrivateKey,
      });
      vi.spyOn(prismaService.instance, "update").mockResolvedValue(updatedInstance);

      // Act
      const result = await service.updateInstanceConfiguration(updates);

      // Assert
      expect(prismaService.instance.update).toHaveBeenCalledWith({
        where: { id: mockInstance.id },
        data: updates,
      });
      expect(result.name).toBe("Updated Instance");
      expect(result.capabilities).toEqual(updates.capabilities);
      expect(result.metadata).toEqual(updates.metadata);
    });

    it("should not expose private key in response", async () => {
      // Arrange
      const updatedInstance = { ...mockInstance, name: "Updated" };
      vi.spyOn(service, "getInstanceIdentity").mockResolvedValue({
        ...mockInstance,
        privateKey: mockDecryptedPrivateKey,
      });
      vi.spyOn(prismaService.instance, "update").mockResolvedValue(updatedInstance);

      // Act
      const result = await service.updateInstanceConfiguration({ name: "Updated" });

      // Assert - SECURITY: Verify private key is NOT in response
      expect(result).not.toHaveProperty("privateKey");
      expect(result).toHaveProperty("publicKey");
      expect(result).toHaveProperty("instanceId");
    });
  });
});