stack/apps/api/src/federation/guards/capability.guard.spec.ts

/**
 * Capability Guard Tests
 *
 * Verifies capability-based authorization enforcement.
 */

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { ExecutionContext, ForbiddenException, UnauthorizedException } from "@nestjs/common";
import { Reflector } from "@nestjs/core";
import { Test, TestingModule } from "@nestjs/testing";
import { FederationConnectionStatus } from "@prisma/client";
import { CapabilityGuard, RequireCapability } from "./capability.guard";
import { ConnectionService } from "../connection.service";
import { FederationAuditService } from "../audit.service";
import type { ConnectionDetails } from "../types/connection.types";

describe("CapabilityGuard", () => {
  let guard: CapabilityGuard;
  let connectionService: vi.Mocked<ConnectionService>;
  let auditService: vi.Mocked<FederationAuditService>;
  let reflector: Reflector;

  const mockConnection: ConnectionDetails = {
    id: "conn-123",
    workspaceId: "ws-456",
    remoteInstanceId: "remote-instance",
    remoteUrl: "https://remote.example.com",
    remotePublicKey: "public-key",
    remoteCapabilities: {
      supportsQuery: true,
      supportsCommand: false,
      supportsEvent: true,
      supportsAgentSpawn: undefined, // Explicitly test undefined
    },
    status: FederationConnectionStatus.ACTIVE,
    metadata: {},
    createdAt: new Date(),
    updatedAt: new Date(),
    connectedAt: new Date(),
    disconnectedAt: null,
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [
        CapabilityGuard,
        {
          provide: ConnectionService,
          useValue: {
            getConnectionById: vi.fn(),
          },
        },
        {
          provide: FederationAuditService,
          useValue: {
            logCapabilityDenied: vi.fn(),
          },
        },
        Reflector,
      ],
    }).compile();

    guard = module.get<CapabilityGuard>(CapabilityGuard);
    connectionService = module.get(ConnectionService) as vi.Mocked<ConnectionService>;
    auditService = module.get(FederationAuditService) as vi.Mocked<FederationAuditService>;
    reflector = module.get<Reflector>(Reflector);
  });

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

  /**
   * Helper to create mock execution context
   */
  function createMockContext(
    requiredCapability: string | undefined,
    requestData: {
      body?: Record<string, unknown>;
      headers?: Record<string, string>;
      params?: Record<string, string>;
      url?: string;
    }
  ): ExecutionContext {
    const mockHandler = vi.fn();

    // Mock reflector to return required capability
    vi.spyOn(reflector, "get").mockReturnValue(requiredCapability);

    return {
      getHandler: () => mockHandler,
      switchToHttp: () => ({
        getRequest: () => ({
          body: requestData.body || {},
          headers: requestData.headers || {},
          params: requestData.params || {},
          url: requestData.url || "/api/test",
        }),
      }),
    } as unknown as ExecutionContext;
  }

  describe("Capability Enforcement", () => {
    it("should allow access when no capability required", async () => {
      const context = createMockContext(undefined, {});

      const result = await guard.canActivate(context);

      expect(result).toBe(true);
      expect(connectionService.getConnectionById).not.toHaveBeenCalled();
    });

    it("should deny access when connection ID is missing", async () => {
      const context = createMockContext("supportsQuery", {});

      await expect(guard.canActivate(context)).rejects.toThrow(UnauthorizedException);
      await expect(guard.canActivate(context)).rejects.toThrow(
        "Federation connection not identified"
      );
    });

    it("should deny access when connection does not exist", async () => {
      const context = createMockContext("supportsQuery", {
        body: { connectionId: "nonexistent" },
      });

      connectionService.getConnectionById.mockResolvedValue(null);

      await expect(guard.canActivate(context)).rejects.toThrow(UnauthorizedException);
      await expect(guard.canActivate(context)).rejects.toThrow("Invalid federation connection");
    });

    it("should deny access when connection is not CONNECTED", async () => {
      const context = createMockContext("supportsQuery", {
        body: { connectionId: "conn-123" },
      });

      connectionService.getConnectionById.mockResolvedValue({
        ...mockConnection,
        status: FederationConnectionStatus.PENDING,
      });

      await expect(guard.canActivate(context)).rejects.toThrow(ForbiddenException);
      await expect(guard.canActivate(context)).rejects.toThrow("Connection is not active");
    });

    it("should deny access when capability is not granted", async () => {
      const context = createMockContext("supportsCommand", {
        body: { connectionId: "conn-123" },
        url: "/api/execute-command",
      });

      connectionService.getConnectionById.mockResolvedValue(mockConnection);

      await expect(guard.canActivate(context)).rejects.toThrow(ForbiddenException);
      await expect(guard.canActivate(context)).rejects.toThrow(
        "Operation requires capability: supportsCommand"
      );

      expect(auditService.logCapabilityDenied).toHaveBeenCalledWith(
        "remote-instance",
        "supportsCommand",
        "/api/execute-command"
      );
    });

    it("should allow access when capability is granted", async () => {
      const context = createMockContext("supportsQuery", {
        body: { connectionId: "conn-123" },
      });

      connectionService.getConnectionById.mockResolvedValue(mockConnection);

      const result = await guard.canActivate(context);

      expect(result).toBe(true);
      expect(auditService.logCapabilityDenied).not.toHaveBeenCalled();
    });
  });

  describe("Connection ID Extraction", () => {
    it("should extract connection ID from request body", async () => {
      const context = createMockContext("supportsQuery", {
        body: { connectionId: "conn-123" },
      });

      connectionService.getConnectionById.mockResolvedValue(mockConnection);

      await guard.canActivate(context);

      expect(connectionService.getConnectionById).toHaveBeenCalledWith("conn-123");
    });

    it("should extract connection ID from headers", async () => {
      const context = createMockContext("supportsQuery", {
        headers: { "x-federation-connection-id": "conn-456" },
      });

      connectionService.getConnectionById.mockResolvedValue(mockConnection);

      await guard.canActivate(context);

      expect(connectionService.getConnectionById).toHaveBeenCalledWith("conn-456");
    });

    it("should extract connection ID from route params", async () => {
      const context = createMockContext("supportsQuery", {
        params: { connectionId: "conn-789" },
      });

      connectionService.getConnectionById.mockResolvedValue(mockConnection);

      await guard.canActivate(context);

      expect(connectionService.getConnectionById).toHaveBeenCalledWith("conn-789");
    });
  });

  describe("Fail-Closed Security", () => {
    it("should deny access when capability is undefined (fail-closed)", async () => {
      const context = createMockContext("supportsAgentSpawn", {
        body: { connectionId: "conn-123" },
      });

      connectionService.getConnectionById.mockResolvedValue(mockConnection);

      await expect(guard.canActivate(context)).rejects.toThrow(ForbiddenException);
      await expect(guard.canActivate(context)).rejects.toThrow(
        "Operation requires capability: supportsAgentSpawn"
      );
    });

    it("should only allow explicitly true values (not truthy)", async () => {
      const context = createMockContext("supportsEvent", {
        body: { connectionId: "conn-123" },
      });

      // Test with explicitly true (should pass)
      connectionService.getConnectionById.mockResolvedValue({
        ...mockConnection,
        remoteCapabilities: { supportsEvent: true },
      });

      const resultTrue = await guard.canActivate(context);
      expect(resultTrue).toBe(true);

      // Test with truthy but not true (should fail)
      connectionService.getConnectionById.mockResolvedValue({
        ...mockConnection,
        remoteCapabilities: { supportsEvent: 1 as unknown as boolean },
      });

      await expect(guard.canActivate(context)).rejects.toThrow(ForbiddenException);
    });
  });

  describe("RequireCapability Decorator", () => {
    it("should create decorator with correct metadata", () => {
      const decorator = RequireCapability("supportsQuery");

      expect(decorator).toBeDefined();
      expect(typeof decorator).toBe("function");
    });
  });
});