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

/**
 * Federation Controller Tests
 */

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { FederationController } from "./federation.controller";
import { FederationService } from "./federation.service";
import { FederationAuditService } from "./audit.service";
import { ConnectionService } from "./connection.service";
import { FederationAgentService } from "./federation-agent.service";
import { AuthGuard } from "../auth/guards/auth.guard";
import { AdminGuard } from "../auth/guards/admin.guard";
import { CsrfGuard } from "../common/guards/csrf.guard";
import { WorkspaceGuard } from "../common/guards/workspace.guard";
import { FederationConnectionStatus } from "@prisma/client";
import type { PublicInstanceIdentity } from "./types/instance.types";
import type { ConnectionDetails } from "./types/connection.types";

describe("FederationController", () => {
  let controller: FederationController;
  let service: FederationService;
  let auditService: FederationAuditService;
  let connectionService: ConnectionService;

  const mockPublicIdentity: PublicInstanceIdentity = {
    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-----",
    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"),
  };

  const mockUser = {
    id: "user-123",
    email: "admin@example.com",
    name: "Admin User",
    workspaceId: "workspace-123",
  };

  const mockConnection: ConnectionDetails = {
    id: "conn-123",
    workspaceId: "workspace-123",
    remoteInstanceId: "remote-instance-456",
    remoteUrl: "https://remote.example.com",
    remotePublicKey: "-----BEGIN PUBLIC KEY-----\nREMOTE\n-----END PUBLIC KEY-----",
    remoteCapabilities: { supportsQuery: true },
    status: FederationConnectionStatus.PENDING,
    metadata: {},
    createdAt: new Date(),
    updatedAt: new Date(),
    connectedAt: null,
    disconnectedAt: null,
  };

  // Store original env value
  const originalDefaultWorkspaceId = process.env.DEFAULT_WORKSPACE_ID;

  beforeEach(async () => {
    // Set environment variable for tests that use getDefaultWorkspaceId()
    process.env.DEFAULT_WORKSPACE_ID = "12345678-1234-4123-8123-123456789abc";

    const module: TestingModule = await Test.createTestingModule({
      controllers: [FederationController],
      providers: [
        {
          provide: FederationService,
          useValue: {
            getPublicIdentity: vi.fn(),
            regenerateKeypair: vi.fn(),
          },
        },
        {
          provide: FederationAuditService,
          useValue: {
            logKeypairRegeneration: vi.fn(),
          },
        },
        {
          provide: ConnectionService,
          useValue: {
            initiateConnection: vi.fn(),
            acceptConnection: vi.fn(),
            rejectConnection: vi.fn(),
            disconnect: vi.fn(),
            getConnections: vi.fn(),
            getConnection: vi.fn(),
            handleIncomingConnectionRequest: vi.fn(),
          },
        },
        {
          provide: FederationAgentService,
          useValue: {
            spawnAgentOnRemote: vi.fn(),
            getAgentStatus: vi.fn(),
            killAgentOnRemote: vi.fn(),
          },
        },
      ],
    })
      .overrideGuard(AuthGuard)
      .useValue({ canActivate: () => true })
      .overrideGuard(AdminGuard)
      .useValue({ canActivate: () => true })
      .overrideGuard(CsrfGuard)
      .useValue({ canActivate: () => true })
      .overrideGuard(WorkspaceGuard)
      .useValue({ canActivate: () => true })
      .compile();

    controller = module.get<FederationController>(FederationController);
    service = module.get<FederationService>(FederationService);
    auditService = module.get<FederationAuditService>(FederationAuditService);
    connectionService = module.get<ConnectionService>(ConnectionService);
  });

  afterEach(() => {
    // Restore original env value
    if (originalDefaultWorkspaceId !== undefined) {
      process.env.DEFAULT_WORKSPACE_ID = originalDefaultWorkspaceId;
    } else {
      delete process.env.DEFAULT_WORKSPACE_ID;
    }
  });

  describe("GET /instance", () => {
    it("should return public instance identity", async () => {
      // Arrange
      vi.spyOn(service, "getPublicIdentity").mockResolvedValue(mockPublicIdentity);

      // Act
      const result = await controller.getInstance();

      // Assert
      expect(result).toEqual(mockPublicIdentity);
      expect(service.getPublicIdentity).toHaveBeenCalledTimes(1);
    });

    it("should not expose private key", async () => {
      // Arrange
      vi.spyOn(service, "getPublicIdentity").mockResolvedValue(mockPublicIdentity);

      // Act
      const result = await controller.getInstance();

      // Assert
      expect(result).not.toHaveProperty("privateKey");
    });

    it("should return consistent identity across multiple calls", async () => {
      // Arrange
      vi.spyOn(service, "getPublicIdentity").mockResolvedValue(mockPublicIdentity);

      // Act
      const result1 = await controller.getInstance();
      const result2 = await controller.getInstance();

      // Assert
      expect(result1).toEqual(result2);
      expect(result1.instanceId).toEqual(result2.instanceId);
    });
  });

  describe("POST /instance/regenerate-keys", () => {
    it("should regenerate keypair and return public identity only", async () => {
      // Arrange
      const updatedIdentity = {
        ...mockPublicIdentity,
        publicKey: "NEW_PUBLIC_KEY",
      };
      vi.spyOn(service, "regenerateKeypair").mockResolvedValue(updatedIdentity);

      const mockRequest = {
        user: mockUser,
      } as any;

      // Act
      const result = await controller.regenerateKeys(mockRequest);

      // Assert
      expect(result).toEqual(updatedIdentity);
      expect(service.regenerateKeypair).toHaveBeenCalledTimes(1);

      // SECURITY FIX: Verify audit logging
      expect(auditService.logKeypairRegeneration).toHaveBeenCalledWith(
        mockUser.id,
        updatedIdentity.instanceId
      );
    });

    it("should NOT expose private key in response", async () => {
      // Arrange
      const updatedIdentity = {
        ...mockPublicIdentity,
        publicKey: "NEW_PUBLIC_KEY",
      };
      vi.spyOn(service, "regenerateKeypair").mockResolvedValue(updatedIdentity);

      const mockRequest = {
        user: mockUser,
      } as any;

      // Act
      const result = await controller.regenerateKeys(mockRequest);

      // Assert - CRITICAL SECURITY TEST
      expect(result).not.toHaveProperty("privateKey");
      expect(result).toHaveProperty("publicKey");
      expect(result).toHaveProperty("instanceId");
    });
  });

  describe("POST /connections/initiate", () => {
    it("should initiate connection to remote instance", async () => {
      const dto = { remoteUrl: "https://remote.example.com" };
      vi.spyOn(connectionService, "initiateConnection").mockResolvedValue(mockConnection);

      const mockRequest = { user: mockUser } as never;
      const result = await controller.initiateConnection(mockRequest, dto);

      expect(result).toEqual(mockConnection);
      expect(connectionService.initiateConnection).toHaveBeenCalledWith(
        mockUser.workspaceId,
        dto.remoteUrl
      );
    });
  });

  describe("POST /connections/:id/accept", () => {
    it("should accept pending connection", async () => {
      const activeConnection = { ...mockConnection, status: FederationConnectionStatus.ACTIVE };
      vi.spyOn(connectionService, "acceptConnection").mockResolvedValue(activeConnection);

      const mockRequest = { user: mockUser } as never;
      const result = await controller.acceptConnection(mockRequest, "conn-123", {});

      expect(result.status).toBe(FederationConnectionStatus.ACTIVE);
      expect(connectionService.acceptConnection).toHaveBeenCalledWith(
        mockUser.workspaceId,
        "conn-123",
        undefined
      );
    });
  });

  describe("POST /connections/:id/reject", () => {
    it("should reject pending connection", async () => {
      const rejectedConnection = {
        ...mockConnection,
        status: FederationConnectionStatus.DISCONNECTED,
      };
      vi.spyOn(connectionService, "rejectConnection").mockResolvedValue(rejectedConnection);

      const mockRequest = { user: mockUser } as never;
      const result = await controller.rejectConnection(mockRequest, "conn-123", {
        reason: "Not approved",
      });

      expect(result.status).toBe(FederationConnectionStatus.DISCONNECTED);
      expect(connectionService.rejectConnection).toHaveBeenCalledWith(
        mockUser.workspaceId,
        "conn-123",
        "Not approved"
      );
    });
  });

  describe("POST /connections/:id/disconnect", () => {
    it("should disconnect active connection", async () => {
      const disconnectedConnection = {
        ...mockConnection,
        status: FederationConnectionStatus.DISCONNECTED,
      };
      vi.spyOn(connectionService, "disconnect").mockResolvedValue(disconnectedConnection);

      const mockRequest = { user: mockUser } as never;
      const result = await controller.disconnectConnection(mockRequest, "conn-123", {
        reason: "Manual disconnect",
      });

      expect(result.status).toBe(FederationConnectionStatus.DISCONNECTED);
      expect(connectionService.disconnect).toHaveBeenCalledWith(
        mockUser.workspaceId,
        "conn-123",
        "Manual disconnect"
      );
    });
  });

  describe("GET /connections", () => {
    it("should list all connections for workspace", async () => {
      vi.spyOn(connectionService, "getConnections").mockResolvedValue([mockConnection]);

      const mockRequest = { user: mockUser } as never;
      const result = await controller.getConnections(mockRequest);

      expect(result).toEqual([mockConnection]);
      expect(connectionService.getConnections).toHaveBeenCalledWith(
        mockUser.workspaceId,
        undefined
      );
    });

    it("should filter connections by status", async () => {
      vi.spyOn(connectionService, "getConnections").mockResolvedValue([mockConnection]);

      const mockRequest = { user: mockUser } as never;
      await controller.getConnections(mockRequest, FederationConnectionStatus.ACTIVE);

      expect(connectionService.getConnections).toHaveBeenCalledWith(
        mockUser.workspaceId,
        FederationConnectionStatus.ACTIVE
      );
    });
  });

  describe("GET /connections/:id", () => {
    it("should return connection details", async () => {
      vi.spyOn(connectionService, "getConnection").mockResolvedValue(mockConnection);

      const mockRequest = { user: mockUser } as never;
      const result = await controller.getConnection(mockRequest, "conn-123");

      expect(result).toEqual(mockConnection);
      expect(connectionService.getConnection).toHaveBeenCalledWith(
        mockUser.workspaceId,
        "conn-123"
      );
    });
  });

  describe("POST /incoming/connect", () => {
    it("should handle incoming connection request", async () => {
      const dto = {
        instanceId: "remote-instance-456",
        instanceUrl: "https://remote.example.com",
        publicKey: "PUBLIC_KEY",
        capabilities: { supportsQuery: true },
        timestamp: Date.now(),
        signature: "valid-signature",
      };
      vi.spyOn(connectionService, "handleIncomingConnectionRequest").mockResolvedValue(
        mockConnection
      );

      const result = await controller.handleIncomingConnection(dto);

      expect(result).toEqual({
        status: "pending",
        connectionId: mockConnection.id,
      });
      expect(connectionService.handleIncomingConnectionRequest).toHaveBeenCalled();
    });

    it("should validate capabilities structure with valid data", async () => {
      const dto = {
        instanceId: "remote-instance-456",
        instanceUrl: "https://remote.example.com",
        publicKey: "PUBLIC_KEY",
        capabilities: {
          supportsQuery: true,
          supportsCommand: false,
          supportsEvent: true,
          supportsAgentSpawn: false,
          protocolVersion: "1.0",
        },
        timestamp: Date.now(),
        signature: "valid-signature",
      };
      vi.spyOn(connectionService, "handleIncomingConnectionRequest").mockResolvedValue(
        mockConnection
      );

      const result = await controller.handleIncomingConnection(dto);

      expect(result.status).toBe("pending");
      expect(connectionService.handleIncomingConnectionRequest).toHaveBeenCalled();
    });
  });
});