@mcp-audit-gateway/core 0.1.0

package/src/proxy/gateway.ts

@@ -0,0 +1,328 @@
+import type {
+  GatewayConfig,
+  UpstreamConfig,
+  UpstreamStatus,
+  ToolEntry,
+  AuditRecord,
+} from "../types.js";
+import { PolicyEngine } from "../policy/engine.js";
+import { AuditLog } from "../attestation/audit-log.js";
+import { createSigner } from "../attestation/signer.js";
+import { GatewayTracer } from "../telemetry/tracer.js";
+import { GatewayMetrics } from "../telemetry/metrics.js";
+import { UpstreamManager } from "./upstream-manager.js";
+
+export class Gateway {
+  private toolCatalog: Map<string, ToolEntry> = new Map();
+  private manualStatuses?: Map<string, UpstreamStatus>;
+  private policyEngine: PolicyEngine;
+  private auditLog: AuditLog;
+  private tracer: GatewayTracer;
+  private metrics: GatewayMetrics;
+  private upstreamManager: UpstreamManager;
+
+  constructor(private config: GatewayConfig) {
+    this.policyEngine = new PolicyEngine(
+      config.policy.defaultEffect,
+      config.policy.rules,
+    );
+
+    const signer = createSigner(config.attestation);
+    this.auditLog = new AuditLog(
+      config.auditLog.path,
+      signer,
+      config.auditLog.rotateAfterMb * 1024 * 1024,
+    );
+
+    this.tracer = new GatewayTracer(config.telemetry);
+    this.metrics = new GatewayMetrics(config.telemetry);
+    this.upstreamManager = new UpstreamManager();
+  }
+
+  async init(): Promise<void> {
+    await this.auditLog.init();
+    for (const upstream of this.config.upstreams) {
+      await this.connectUpstream(upstream);
+    }
+  }
+
+  private async connectUpstream(upstream: UpstreamConfig): Promise<void> {
+    try {
+      const conn = await this.upstreamManager.connect(upstream);
+      this.registerUpstreamTools(
+        upstream.name,
+        upstream.namespace,
+        conn.tools,
+      );
+    } catch (err) {
+      console.error(`Failed to connect to upstream ${upstream.name}: ${err}`);
+    }
+  }
+
+  async handleToolsList(
+    principal?: string,
+    _meta?: Record<string, unknown>,
+  ): Promise<{ tools: unknown[]; _meta?: Record<string, unknown> }> {
+    const allTools = Array.from(this.toolCatalog.values());
+    const filtered = this.policyEngine.filterTools(principal, allTools);
+
+    const tools = filtered.map((entry) => ({
+      name: entry.name,
+      description: entry.description,
+      inputSchema: entry.inputSchema,
+      annotations: entry.annotations,
+      _meta: {
+        "x-gateway-routing/v1": {
+          upstream: entry.upstream,
+          namespace: entry.namespace,
+        },
+      },
+    }));
+
+    return {
+      tools,
+      _meta: {
+        "x-gateway-routing/v1": {
+          filteringApplied: true,
+          identityResolved: principal ?? "anonymous",
+          totalToolsPreFilter: this.toolCatalog.size,
+          totalToolsPostFilter: filtered.length,
+        },
+      },
+    };
+  }
+
+  async handleToolsCall(
+    toolName: string,
+    args: Record<string, unknown>,
+    principal?: string,
+    traceContext?: { traceparent?: string; tracestate?: string },
+  ): Promise<{ result: unknown; auditRecord: AuditRecord }> {
+    const startTime = Date.now();
+    const tool = this.toolCatalog.get(toolName);
+
+    if (!tool) {
+      const record = await this.auditLog.record("tools/call", {
+        toolName,
+        principal,
+        durationMs: Date.now() - startTime,
+        success: false,
+        errorCode: -32602,
+      });
+      throw new ToolCallError(-32602, `Unknown tool: ${toolName}`, record);
+    }
+
+    const decision = this.policyEngine.evaluate(principal, tool);
+    if (!decision.allowed) {
+      this.metrics.recordPolicyDenial({
+        principal: principal ?? "anonymous",
+        tool: toolName,
+        reason: decision.reason ?? "unauthorized",
+      });
+      const record = await this.auditLog.record("tools/call", {
+        toolName,
+        namespace: tool.namespace,
+        upstream: tool.upstream,
+        principal,
+        durationMs: Date.now() - startTime,
+        success: false,
+        errorCode: -32603,
+      });
+      throw new ToolCallError(
+        -32603,
+        decision.reason ?? "unauthorized",
+        record,
+      );
+    }
+
+    const span = this.tracer.startRouteSpan(
+      "tools/call",
+      toolName,
+      tool.upstream,
+      traceContext,
+    );
+
+    try {
+      this.policyEngine.recordInvocation(principal, tool);
+
+      const result = await this.upstreamManager.callTool(
+        tool.upstream,
+        tool.originalName,
+        args,
+      );
+      const durationMs = Date.now() - startTime;
+
+      this.tracer.recordDuration(span, durationMs);
+      this.tracer.endSpan(span, true);
+
+      this.metrics.recordToolCall({
+        namespace: tool.namespace,
+        tool: toolName,
+        principal: principal ?? "anonymous",
+        success: "true",
+      });
+      this.metrics.recordDuration(durationMs, {
+        namespace: tool.namespace,
+        upstream: tool.upstream,
+      });
+
+      const record = await this.safeAuditRecord("tools/call", {
+        toolName,
+        namespace: tool.namespace,
+        upstream: tool.upstream,
+        principal,
+        durationMs,
+        success: true,
+      });
+
+      return { result, auditRecord: record };
+    } catch (err) {
+      const durationMs = Date.now() - startTime;
+      this.tracer.recordDuration(span, durationMs);
+      this.tracer.endSpan(span, false);
+
+      this.metrics.recordToolCall({
+        namespace: tool.namespace,
+        tool: toolName,
+        principal: principal ?? "anonymous",
+        success: "false",
+      });
+      this.metrics.recordDuration(durationMs, {
+        namespace: tool.namespace,
+        upstream: tool.upstream,
+      });
+
+      if (err instanceof ToolCallError) throw err;
+
+      const record = await this.safeAuditRecord("tools/call", {
+        toolName,
+        namespace: tool.namespace,
+        upstream: tool.upstream,
+        principal,
+        durationMs,
+        success: false,
+        errorCode: -32603,
+      });
+      throw new ToolCallError(-32603, "Upstream server error", record);
+    }
+  }
+
+  getStatus(): {
+    gateway: { status: string; version: string };
+    upstreams: UpstreamStatus[];
+    aggregate: {
+      totalUpstreams: number;
+      healthyUpstreams: number;
+      degradedUpstreams: number;
+      unavailableUpstreams: number;
+      totalTools: number;
+    };
+  } {
+    const managed = this.upstreamManager.getAllStatuses();
+    const manual = this.manualStatuses ? Array.from(this.manualStatuses.values()) : [];
+    const upstreams = managed.length > 0 ? managed : manual;
+    return {
+      gateway: {
+        status: "healthy",
+        version: this.config.version,
+      },
+      upstreams,
+      aggregate: {
+        totalUpstreams: upstreams.length,
+        healthyUpstreams: upstreams.filter((u) => u.status === "healthy").length,
+        degradedUpstreams: upstreams.filter((u) => u.status === "degraded").length,
+        unavailableUpstreams: upstreams.filter((u) => u.status === "unavailable").length,
+        totalTools: this.toolCatalog.size,
+      },
+    };
+  }
+
+  getServerDiscover(): Record<string, unknown> {
+    return {
+      name: this.config.name,
+      version: this.config.version,
+      protocol: "2025-03-26",
+      capabilities: {
+        tools: { listChanged: true },
+        gateway: {
+          namespacing: true,
+          filtering: true,
+          routing: true,
+          healthReporting: true,
+          attestation: this.config.attestation.enabled,
+          upstreamCount: this.config.upstreams.length,
+        },
+      },
+      extensions: ["x-gateway-routing/v1", "x-gateway-attestation/v1"],
+    };
+  }
+
+  registerUpstreamTools(
+    upstreamName: string,
+    namespace: string,
+    tools: Array<{
+      name: string;
+      description?: string;
+      inputSchema?: unknown;
+      annotations?: Record<string, unknown>;
+    }>,
+  ): void {
+    for (const tool of tools) {
+      const namespacedName = `${namespace}/${tool.name}`;
+      const entry: ToolEntry = {
+        name: namespacedName,
+        originalName: tool.name,
+        namespace,
+        upstream: upstreamName,
+        description: tool.description,
+        inputSchema: tool.inputSchema,
+        annotations: tool.annotations,
+      };
+      this.toolCatalog.set(namespacedName, entry);
+    }
+
+    if (!this.manualStatuses) this.manualStatuses = new Map();
+    this.manualStatuses.set(upstreamName, {
+      name: upstreamName,
+      namespace,
+      status: "healthy",
+      toolCount: tools.length,
+      lastSuccessfulContact: new Date().toISOString(),
+    });
+
+    this.metrics.setUpstreamStatus(upstreamName, "healthy", 1);
+  }
+
+  private async safeAuditRecord(
+    method: string,
+    opts: Parameters<AuditLog["record"]>[1],
+  ): Promise<AuditRecord> {
+    try {
+      return await this.auditLog.record(method, opts);
+    } catch (err) {
+      console.error(`Audit log write failed: ${err}`);
+      return {
+        id: "unrecorded",
+        timestamp: new Date().toISOString(),
+        method,
+        ...opts,
+        attestation: undefined,
+      };
+    }
+  }
+
+  async shutdown(): Promise<void> {
+    await this.upstreamManager.disconnectAll();
+  }
+}
+
+export class ToolCallError extends Error {
+  constructor(
+    public code: number,
+    message: string,
+    public auditRecord: AuditRecord,
+  ) {
+    super(message);
+    this.name = "ToolCallError";
+  }
+}

package/src/proxy/mcp-server-adapter.test.ts

@@ -0,0 +1,197 @@
+import { describe, it, expect, beforeEach, afterEach } from "vitest";
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
+import {
+  ListToolsResultSchema,
+  CompatibilityCallToolResultSchema,
+} from "@modelcontextprotocol/sdk/types.js";
+import { Gateway } from "./gateway.js";
+import { McpServerAdapter } from "./mcp-server-adapter.js";
+import type { GatewayConfig } from "../types.js";
+import { unlink } from "node:fs/promises";
+
+const AUDIT_PATH = "/tmp/mcp-adapter-test-audit.jsonl";
+
+const testConfig: GatewayConfig = {
+  name: "adapter-test-gateway",
+  version: "0.2.0",
+  listen: { transport: "streamable-http", port: 3200, host: "127.0.0.1" },
+  upstreams: [
+    {
+      name: "mock-server",
+      namespace: "mock",
+      transport: { type: "stdio", command: "echo", args: [] },
+    },
+  ],
+  policy: {
+    defaultEffect: "allow",
+    rules: [
+      {
+        effect: "deny",
+        principals: ["agent:restricted"],
+        tools: ["mock/secret_tool"],
+      },
+    ],
+  },
+  attestation: {
+    enabled: true,
+    algorithm: "hmac-sha256",
+    secret: "b".repeat(64),
+    includeParams: false,
+    includeResult: false,
+  },
+  telemetry: { enabled: false, serviceName: "test", sampleRate: 0 },
+  auditLog: { enabled: true, path: AUDIT_PATH, rotateAfterMb: 10 },
+};
+
+describe("McpServerAdapter", () => {
+  let gateway: Gateway;
+  let adapter: McpServerAdapter;
+  let client: Client;
+
+  beforeEach(async () => {
+    gateway = new Gateway(testConfig);
+    await gateway.init();
+
+    // Register tools manually (simulating upstream discovery)
+    gateway.registerUpstreamTools("mock-server", "mock", [
+      {
+        name: "greet",
+        description: "Greets a user",
+        inputSchema: { type: "object", properties: { name: { type: "string" } } },
+      },
+      {
+        name: "secret_tool",
+        description: "A restricted tool",
+        inputSchema: { type: "object" },
+      },
+    ]);
+
+    adapter = new McpServerAdapter(gateway, {
+      name: "adapter-test-gateway",
+      version: "0.2.0",
+      resolvePrincipal: (meta) => meta?.requestIdentity as string | undefined,
+    });
+
+    const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+    client = new Client(
+      { name: "test-client", version: "1.0.0" },
+      { capabilities: {} },
+    );
+
+    await adapter.connect(serverTransport);
+    await client.connect(clientTransport);
+  });
+
+  afterEach(async () => {
+    await client.close();
+    await adapter.close();
+    try {
+      await unlink(AUDIT_PATH);
+    } catch {
+      // file may not exist
+    }
+  });
+
+  describe("tools/list", () => {
+    it("lists tools with namespace prefixes via the SDK client", async () => {
+      const result = await client.listTools();
+
+      expect(result.tools.length).toBe(2);
+      const names = result.tools.map((t) => t.name);
+      expect(names).toContain("mock/greet");
+      expect(names).toContain("mock/secret_tool");
+    });
+
+    it("includes tool descriptions through the adapter", async () => {
+      const result = await client.listTools();
+
+      const greet = result.tools.find((t) => t.name === "mock/greet");
+      expect(greet).toBeDefined();
+      expect(greet!.description).toBe("Greets a user");
+    });
+
+    it("applies policy filtering when requestIdentity is provided", async () => {
+      // Use a custom request with _meta to pass identity
+      const result = await client.request(
+        {
+          method: "tools/list",
+          params: {
+            _meta: { requestIdentity: "agent:restricted" },
+          },
+        },
+        ListToolsResultSchema,
+      );
+
+      const tools = result.tools;
+      const names = tools.map((t) => t.name);
+      expect(names).toContain("mock/greet");
+      expect(names).not.toContain("mock/secret_tool");
+    });
+  });
+
+  describe("tools/call", () => {
+    it("returns attestation metadata in tool call responses", async () => {
+      // Call a tool that will be denied by policy, which produces attestation
+      const result = await client.request(
+        {
+          method: "tools/call",
+          params: {
+            name: "mock/secret_tool",
+            arguments: {},
+            _meta: { requestIdentity: "agent:restricted" },
+          },
+        },
+        CompatibilityCallToolResultSchema,
+      );
+
+      // Should get an error response with attestation metadata
+      const meta = (result as any)._meta;
+      expect(meta).toBeDefined();
+      expect(meta["x-gateway-attestation/v1"]).toBeDefined();
+      expect(meta["x-gateway-attestation/v1"].auditId).toBeDefined();
+      expect(meta["x-gateway-attestation/v1"].attestation).toBeDefined();
+      expect(meta["x-gateway-attestation/v1"].timestamp).toBeDefined();
+
+      // Should be flagged as an error
+      expect((result as any).isError).toBe(true);
+    });
+
+    it("returns error content for unknown tools with attestation", async () => {
+      const result = await client.request(
+        {
+          method: "tools/call",
+          params: {
+            name: "mock/nonexistent",
+            arguments: {},
+          },
+        },
+        CompatibilityCallToolResultSchema,
+      );
+
+      expect((result as any).isError).toBe(true);
+      const content = (result as any).content;
+      expect(content).toHaveLength(1);
+      expect(content[0].type).toBe("text");
+      expect(content[0].text).toContain("Unknown tool");
+
+      // Attestation metadata should still be present
+      const meta = (result as any)._meta;
+      expect(meta["x-gateway-attestation/v1"].auditId).toBeDefined();
+    });
+  });
+
+  describe("server capabilities", () => {
+    it("advertises tools capability", () => {
+      const caps = client.getServerCapabilities();
+      expect(caps?.tools).toBeDefined();
+    });
+
+    it("advertises gateway extension in experimental capabilities", () => {
+      const caps = client.getServerCapabilities();
+      expect(caps?.experimental).toBeDefined();
+      expect((caps?.experimental as any)["x-gateway-routing/v1"]).toBeDefined();
+      expect((caps?.experimental as any)["x-gateway-attestation/v1"]).toBeDefined();
+    });
+  });
+});

package/src/proxy/mcp-server-adapter.ts

@@ -0,0 +1,140 @@
+import { Server } from "@modelcontextprotocol/sdk/server/index.js";
+import {
+  ListToolsRequestSchema,
+  CallToolRequestSchema,
+} from "@modelcontextprotocol/sdk/types.js";
+import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
+import { Gateway, ToolCallError } from "./gateway.js";
+
+export interface McpServerAdapterOptions {
+  name: string;
+  version: string;
+  resolvePrincipal?: (meta?: Record<string, unknown>) => string | undefined;
+}
+
+export class McpServerAdapter {
+  private server: Server;
+  private gateway: Gateway;
+  private resolvePrincipal?: (meta?: Record<string, unknown>) => string | undefined;
+
+  constructor(gateway: Gateway, options: McpServerAdapterOptions) {
+    this.gateway = gateway;
+    this.resolvePrincipal = options.resolvePrincipal;
+
+    const discover = gateway.getServerDiscover();
+    const gatewayCapabilities = (discover.capabilities as Record<string, unknown>)?.gateway;
+
+    this.server = new Server(
+      { name: options.name, version: options.version },
+      {
+        capabilities: {
+          tools: { listChanged: true },
+          experimental: {
+            "x-gateway-routing/v1": gatewayCapabilities ?? {},
+            "x-gateway-attestation/v1": { enabled: true },
+          },
+        },
+      },
+    );
+
+    this.registerHandlers();
+  }
+
+  async connect(transport: Transport): Promise<void> {
+    await this.server.connect(transport);
+  }
+
+  async close(): Promise<void> {
+    await this.server.close();
+  }
+
+  getServer(): Server {
+    return this.server;
+  }
+
+  private registerHandlers(): void {
+    this.server.setRequestHandler(
+      ListToolsRequestSchema,
+      async (request) => {
+        const principal = this.extractPrincipal(request.params?._meta);
+        const listResult = await this.gateway.handleToolsList(principal);
+
+        return {
+          tools: listResult.tools.map((tool: any) => ({
+            name: tool.name,
+            description: tool.description,
+            inputSchema: tool.inputSchema ?? { type: "object" as const },
+            annotations: tool.annotations,
+          })),
+          _meta: listResult._meta,
+        };
+      },
+    );
+
+    this.server.setRequestHandler(
+      CallToolRequestSchema,
+      async (request) => {
+        const principal = this.extractPrincipal(request.params?._meta);
+        const traceContext = this.extractTraceContext(request.params?._meta);
+        const toolName = request.params.name;
+        const args = (request.params.arguments ?? {}) as Record<string, unknown>;
+
+        try {
+          const { result, auditRecord } = await this.gateway.handleToolsCall(
+            toolName,
+            args,
+            principal,
+            traceContext,
+          );
+
+          const upstreamResult = result as {
+            content?: Array<{ type: string; text?: string; [key: string]: unknown }>;
+            isError?: boolean;
+          };
+
+          return {
+            content: upstreamResult?.content ?? [{ type: "text" as const, text: JSON.stringify(result) }],
+            isError: upstreamResult?.isError,
+            _meta: {
+              "x-gateway-attestation/v1": {
+                auditId: auditRecord.id,
+                attestation: auditRecord.attestation,
+                timestamp: auditRecord.timestamp,
+              },
+            },
+          };
+        } catch (err) {
+          if (err instanceof ToolCallError) {
+            return {
+              content: [{ type: "text" as const, text: err.message }],
+              isError: true,
+              _meta: {
+                "x-gateway-attestation/v1": {
+                  auditId: err.auditRecord.id,
+                  attestation: err.auditRecord.attestation,
+                  timestamp: err.auditRecord.timestamp,
+                },
+              },
+            };
+          }
+          throw err;
+        }
+      },
+    );
+  }
+
+  private extractPrincipal(meta?: Record<string, unknown>): string | undefined {
+    if (this.resolvePrincipal) {
+      return this.resolvePrincipal(meta);
+    }
+    return undefined;
+  }
+
+  private extractTraceContext(
+    meta?: Record<string, unknown>,
+  ): { traceparent?: string; tracestate?: string } | undefined {
+    if (!meta) return undefined;
+    const tc = meta.traceContext as { traceparent?: string; tracestate?: string } | undefined;
+    return tc;
+  }
+}