@mcp-audit-gateway/core 0.1.0

package/src/integration.test.ts

@@ -0,0 +1,108 @@
+import { describe, it, expect, afterEach } from "vitest";
+import { Server } from "@modelcontextprotocol/sdk/server/index.js";
+import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { ListToolsRequestSchema, CallToolRequestSchema } from "@modelcontextprotocol/sdk/types.js";
+import { Gateway } from "./proxy/gateway.js";
+import { HmacSigner } from "./attestation/signer.js";
+import { verifyAuditLog } from "./attestation/verify.js";
+import { unlink } from "node:fs/promises";
+import type { GatewayConfig } from "./types.js";
+
+const AUDIT_PATH = "/tmp/integration-test-audit.jsonl";
+
+const config: GatewayConfig = {
+  name: "test-integration",
+  version: "0.1.0",
+  listen: { transport: "streamable-http", port: 3199, host: "127.0.0.1" },
+  upstreams: [],
+  policy: {
+    defaultEffect: "allow",
+    rules: [
+      { effect: "deny", principals: ["agent:blocked"], tools: ["*/secret_*"] },
+    ],
+  },
+  attestation: { enabled: true, algorithm: "hmac-sha256", secret: "c".repeat(64), includeParams: false, includeResult: false },
+  telemetry: { enabled: false, serviceName: "test", sampleRate: 0 },
+  auditLog: { enabled: true, path: AUDIT_PATH, rotateAfterMb: 10 },
+};
+
+describe("Integration: Gateway with in-memory MCP server", () => {
+  afterEach(async () => {
+    try { await unlink(AUDIT_PATH); } catch {}
+  });
+
+  it("routes a tool call through the gateway and produces a verifiable audit record", async () => {
+    const server = new Server(
+      { name: "echo-server", version: "1.0.0" },
+      { capabilities: { tools: {} } },
+    );
+
+    server.setRequestHandler(
+      ListToolsRequestSchema,
+      async () => ({
+        tools: [
+          { name: "echo", description: "Echoes input", inputSchema: { type: "object" as const, properties: { msg: { type: "string" } } } },
+          { name: "secret_tool", description: "Secret", inputSchema: { type: "object" as const } },
+        ],
+      }),
+    );
+
+    server.setRequestHandler(
+      CallToolRequestSchema,
+      async (request) => ({
+        content: [{ type: "text", text: `echo: ${request.params.arguments?.msg}` }],
+      }),
+    );
+
+    const gateway = new Gateway(config);
+    await gateway.init();
+
+    const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+    const client = new Client({ name: "test-client", version: "1.0.0" }, { capabilities: {} });
+
+    await server.connect(serverTransport);
+    await client.connect(clientTransport);
+
+    const toolsResult = await client.listTools();
+    gateway.registerUpstreamTools("echo-server", "echo", toolsResult.tools.map((t) => ({
+      name: t.name,
+      description: t.description,
+      inputSchema: t.inputSchema,
+    })));
+
+    const listed = await gateway.handleToolsList();
+    expect(listed.tools).toHaveLength(2);
+    const names = listed.tools.map((t: any) => t.name);
+    expect(names).toContain("echo/echo");
+    expect(names).toContain("echo/secret_tool");
+
+    const filteredList = await gateway.handleToolsList("agent:blocked");
+    const filteredNames = filteredList.tools.map((t: any) => t.name);
+    expect(filteredNames).toContain("echo/echo");
+    expect(filteredNames).not.toContain("echo/secret_tool");
+
+    // Make a denied call to generate an audit record
+    try {
+      await gateway.handleToolsCall("echo/secret_tool", {}, "agent:blocked");
+    } catch {
+      // expected denial
+    }
+
+    // Make a call to unknown tool to generate another audit record
+    try {
+      await gateway.handleToolsCall("echo/nonexistent", {});
+    } catch {
+      // expected
+    }
+
+    const signer = new HmacSigner("c".repeat(64));
+    const verifyResult = await verifyAuditLog(AUDIT_PATH, signer);
+    expect(verifyResult.total).toBe(2);
+    expect(verifyResult.valid).toBe(2);
+    expect(verifyResult.invalid).toBe(0);
+
+    await client.close();
+    await server.close();
+  });
+});

package/src/policy/engine.test.ts

@@ -0,0 +1,95 @@
+import { describe, it, expect } from "vitest";
+import { PolicyEngine } from "./engine.js";
+import type { ToolEntry } from "../types.js";
+
+const tool = (name: string, ns: string): ToolEntry => ({
+  name: `${ns}/${name}`,
+  originalName: name,
+  namespace: ns,
+  upstream: `${ns}-server`,
+});
+
+describe("PolicyEngine", () => {
+  describe("default-allow", () => {
+    const engine = new PolicyEngine("allow", [
+      { effect: "deny", principals: ["agent:readonly-*"], tools: ["*/write_*", "*/delete_*"] },
+      { effect: "deny", tools: ["admin/shutdown"] },
+    ]);
+
+    it("allows unmatched tools", () => {
+      const decision = engine.evaluate("agent:dev", tool("read_data", "github"));
+      expect(decision.allowed).toBe(true);
+    });
+
+    it("denies matching principal+tool pattern", () => {
+      const decision = engine.evaluate("agent:readonly-bot", tool("write_file", "fs"));
+      expect(decision.allowed).toBe(false);
+    });
+
+    it("allows write for non-readonly principal", () => {
+      const decision = engine.evaluate("agent:admin", tool("write_file", "fs"));
+      expect(decision.allowed).toBe(true);
+    });
+
+    it("denies tool pattern regardless of principal", () => {
+      const decision = engine.evaluate("agent:admin", tool("shutdown", "admin"));
+      expect(decision.allowed).toBe(false);
+    });
+  });
+
+  describe("default-deny", () => {
+    const engine = new PolicyEngine("deny", [
+      { effect: "allow", principals: ["agent:trusted-*"], namespaces: ["github"] },
+    ]);
+
+    it("denies by default", () => {
+      const decision = engine.evaluate("agent:unknown", tool("create_pr", "github"));
+      expect(decision.allowed).toBe(false);
+    });
+
+    it("allows matching principal+namespace", () => {
+      const decision = engine.evaluate("agent:trusted-ci", tool("create_pr", "github"));
+      expect(decision.allowed).toBe(true);
+    });
+
+    it("denies trusted principal on wrong namespace", () => {
+      const decision = engine.evaluate("agent:trusted-ci", tool("drop_table", "database"));
+      expect(decision.allowed).toBe(false);
+    });
+  });
+
+  describe("rate limiting", () => {
+    it("allows calls within limit", () => {
+      const engine = new PolicyEngine("allow", [
+        { effect: "allow", principals: ["agent:bot"], rateLimit: { maxPerMinute: 3 } },
+      ]);
+
+      const t = tool("read", "api");
+      engine.evaluate("agent:bot", t);
+      engine.recordInvocation("agent:bot", t);
+      engine.recordInvocation("agent:bot", t);
+      engine.recordInvocation("agent:bot", t);
+
+      const decision = engine.evaluate("agent:bot", t);
+      expect(decision.allowed).toBe(false);
+      expect(decision.reason).toContain("rate limit");
+    });
+  });
+
+  describe("filterTools", () => {
+    const engine = new PolicyEngine("allow", [
+      { effect: "deny", principals: ["agent:limited"], tools: ["*/dangerous_*"] },
+    ]);
+
+    it("removes denied tools from list", () => {
+      const tools = [
+        tool("safe_read", "fs"),
+        tool("dangerous_delete", "fs"),
+        tool("safe_write", "fs"),
+      ];
+      const filtered = engine.filterTools("agent:limited", tools);
+      expect(filtered).toHaveLength(2);
+      expect(filtered.map((t) => t.originalName)).not.toContain("dangerous_delete");
+    });
+  });
+});

package/src/policy/engine.ts

@@ -0,0 +1,162 @@
+import type { PolicyRule, ToolEntry } from "../types.js";
+
+export interface PolicyDecision {
+  allowed: boolean;
+  reason?: string;
+  rateLimit?: { maxPerMinute?: number; maxPerHour?: number };
+}
+
+interface RateLimitState {
+  minuteCounts: Map<string, { count: number; resetAt: number }>;
+  hourCounts: Map<string, { count: number; resetAt: number }>;
+}
+
+export class PolicyEngine {
+  private rateLimitState: RateLimitState = {
+    minuteCounts: new Map(),
+    hourCounts: new Map(),
+  };
+
+  constructor(
+    private defaultEffect: "allow" | "deny",
+    private rules: PolicyRule[],
+  ) {}
+
+  evaluate(principal: string | undefined, tool: ToolEntry): PolicyDecision {
+    let matchedRule: PolicyRule | null = null;
+
+    for (const rule of this.rules) {
+      if (this.ruleMatches(rule, principal, tool)) {
+        matchedRule = rule;
+        break;
+      }
+    }
+
+    if (!matchedRule) {
+      return { allowed: this.defaultEffect === "allow" };
+    }
+
+    if (matchedRule.effect === "deny") {
+      return { allowed: false, reason: "denied by policy rule" };
+    }
+
+    if (matchedRule.rateLimit) {
+      const key = `${principal ?? "anonymous"}:${tool.name}`;
+      const rateLimited = this.checkRateLimit(key, matchedRule.rateLimit);
+      if (rateLimited) {
+        return { allowed: false, reason: "rate limit exceeded", rateLimit: matchedRule.rateLimit };
+      }
+    }
+
+    return { allowed: true, rateLimit: matchedRule.rateLimit };
+  }
+
+  filterTools(principal: string | undefined, tools: ToolEntry[]): ToolEntry[] {
+    return tools.filter((tool) => this.evaluate(principal, tool).allowed);
+  }
+
+  recordInvocation(principal: string | undefined, tool: ToolEntry): void {
+    const key = `${principal ?? "anonymous"}:${tool.name}`;
+    this.incrementCounter(key);
+  }
+
+  private ruleMatches(
+    rule: PolicyRule,
+    principal: string | undefined,
+    tool: ToolEntry,
+  ): boolean {
+    if (rule.principals && rule.principals.length > 0) {
+      if (!principal || !this.matchesPattern(principal, rule.principals)) {
+        return false;
+      }
+    }
+
+    if (rule.namespaces && rule.namespaces.length > 0) {
+      if (!rule.namespaces.includes(tool.namespace)) {
+        return false;
+      }
+    }
+
+    if (rule.tools && rule.tools.length > 0) {
+      if (!this.matchesPattern(tool.name, rule.tools) &&
+          !this.matchesPattern(tool.originalName, rule.tools)) {
+        return false;
+      }
+    }
+
+    return true;
+  }
+
+  private matchesPattern(value: string, patterns: string[]): boolean {
+    return patterns.some((pattern) => this.globMatch(pattern, value));
+  }
+
+  private globMatch(pattern: string, value: string): boolean {
+    if (!pattern.includes("*")) return pattern === value;
+
+    const parts = pattern.split("*");
+    let pos = 0;
+
+    for (let i = 0; i < parts.length; i++) {
+      const part = parts[i];
+      if (part === "") continue;
+
+      if (i === 0) {
+        if (!value.startsWith(part)) return false;
+        pos = part.length;
+      } else if (i === parts.length - 1) {
+        if (value.length - pos < part.length) return false;
+        if (!value.endsWith(part)) return false;
+        const suffixStart = value.length - part.length;
+        if (suffixStart < pos) return false;
+      } else {
+        const idx = value.indexOf(part, pos);
+        if (idx === -1) return false;
+        pos = idx + part.length;
+      }
+    }
+
+    return true;
+  }
+
+  private checkRateLimit(
+    key: string,
+    limit: { maxPerMinute?: number; maxPerHour?: number },
+  ): boolean {
+    const now = Date.now();
+
+    if (limit.maxPerMinute != null) {
+      const entry = this.rateLimitState.minuteCounts.get(key);
+      if (entry && entry.resetAt > now && entry.count >= limit.maxPerMinute) {
+        return true;
+      }
+    }
+
+    if (limit.maxPerHour != null) {
+      const entry = this.rateLimitState.hourCounts.get(key);
+      if (entry && entry.resetAt > now && entry.count >= limit.maxPerHour) {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
+  private incrementCounter(key: string): void {
+    const now = Date.now();
+
+    const minute = this.rateLimitState.minuteCounts.get(key);
+    if (!minute || minute.resetAt <= now) {
+      this.rateLimitState.minuteCounts.set(key, { count: 1, resetAt: now + 60_000 });
+    } else {
+      minute.count++;
+    }
+
+    const hour = this.rateLimitState.hourCounts.get(key);
+    if (!hour || hour.resetAt <= now) {
+      this.rateLimitState.hourCounts.set(key, { count: 1, resetAt: now + 3_600_000 });
+    } else {
+      hour.count++;
+    }
+  }
+}

package/src/proxy/gateway.test.ts

@@ -0,0 +1,108 @@
+import { describe, it, expect, beforeEach } from "vitest";
+import { Gateway } from "./gateway.js";
+import type { GatewayConfig } from "../types.js";
+
+const testConfig: GatewayConfig = {
+  name: "test-gateway",
+  version: "0.1.0",
+  listen: { transport: "streamable-http", port: 3100, host: "127.0.0.1" },
+  upstreams: [
+    {
+      name: "test-server",
+      namespace: "test",
+      transport: { type: "stdio", command: "echo", args: [] },
+    },
+  ],
+  policy: {
+    defaultEffect: "allow",
+    rules: [
+      {
+        effect: "deny",
+        principals: ["agent:blocked"],
+        tools: ["test/dangerous_tool"],
+      },
+    ],
+  },
+  attestation: { enabled: true, algorithm: "hmac-sha256", secret: "a".repeat(64), includeParams: false, includeResult: false },
+  telemetry: { enabled: false, serviceName: "test", sampleRate: 0 },
+  auditLog: { enabled: true, path: "/tmp/test-audit.jsonl", rotateAfterMb: 10 },
+};
+
+describe("Gateway", () => {
+  let gateway: Gateway;
+
+  beforeEach(async () => {
+    gateway = new Gateway(testConfig);
+    await gateway.init();
+
+    gateway.registerUpstreamTools("test-server", "test", [
+      { name: "safe_tool", description: "A safe tool" },
+      { name: "dangerous_tool", description: "A dangerous tool" },
+      { name: "read_data", description: "Read some data" },
+    ]);
+  });
+
+  describe("server/discover", () => {
+    it("advertises gateway capabilities", () => {
+      const discover = gateway.getServerDiscover();
+      expect(discover.capabilities).toHaveProperty("gateway");
+      const gw = (discover.capabilities as any).gateway;
+      expect(gw.namespacing).toBe(true);
+      expect(gw.filtering).toBe(true);
+      expect(gw.routing).toBe(true);
+      expect(gw.attestation).toBe(true);
+      expect(gw.upstreamCount).toBe(1);
+    });
+  });
+
+  describe("tools/list", () => {
+    it("returns namespaced tools", async () => {
+      const result = await gateway.handleToolsList();
+      expect(result.tools).toHaveLength(3);
+      const names = result.tools.map((t: any) => t.name);
+      expect(names).toContain("test/safe_tool");
+      expect(names).toContain("test/dangerous_tool");
+      expect(names).toContain("test/read_data");
+    });
+
+    it("filters tools based on policy", async () => {
+      const result = await gateway.handleToolsList("agent:blocked");
+      const names = result.tools.map((t: any) => t.name);
+      expect(names).toContain("test/safe_tool");
+      expect(names).not.toContain("test/dangerous_tool");
+      expect(names).toContain("test/read_data");
+    });
+
+    it("includes filtering metadata", async () => {
+      const result = await gateway.handleToolsList("agent:blocked");
+      const meta = result._meta as any;
+      expect(meta["x-gateway-routing/v1"].filteringApplied).toBe(true);
+      expect(meta["x-gateway-routing/v1"].totalToolsPreFilter).toBe(3);
+      expect(meta["x-gateway-routing/v1"].totalToolsPostFilter).toBe(2);
+    });
+  });
+
+  describe("tools/call", () => {
+    it("rejects unknown tools", async () => {
+      await expect(
+        gateway.handleToolsCall("nonexistent/tool", {}),
+      ).rejects.toThrow("Unknown tool");
+    });
+
+    it("rejects unauthorized tool calls", async () => {
+      await expect(
+        gateway.handleToolsCall("test/dangerous_tool", {}, "agent:blocked"),
+      ).rejects.toThrow();
+    });
+  });
+
+  describe("status", () => {
+    it("reports gateway and upstream health", () => {
+      const status = gateway.getStatus();
+      expect(status.gateway.status).toBe("healthy");
+      expect(status.upstreams).toHaveLength(1);
+      expect(status.upstreams[0].name).toBe("test-server");
+      expect(status.aggregate.totalTools).toBe(3);
+    });
+  });
+});