Add audit middleware to vMCP server (#2981)

* Add audit middleware to vMCP server

Integrate the existing ToolHive audit middleware into the Virtual MCP
Server to enable audit logging of MCP operations for security,
compliance, and debugging purposes.

Changes:
  - Add AuditConfig field to vMCP server Config struct
  - Add audit middleware initialization in Start() method
  - Position audit middleware in chain: Auth → Audit → Discovery → Telemetry
  - Use "streamable-http" transport type for vMCP audit events
  - Add unit tests for audit configuration handling

The audit middleware is optional (enabled when AuditConfig is provided)
and uses the existing pkg/audit infrastructure with MCP-specific event
types (mcp_tool_call, mcp_initialize, etc.).

Partially-closes: #2980

* changes from review

---------

Co-authored-by: taskbot <[email protected]>

Author: yrobla

pkg/vmcp/server/integration_test.go

@@ -8,13 +8,16 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"os"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 
+	"github.com/stacklok/toolhive/pkg/audit"
 	"github.com/stacklok/toolhive/pkg/auth"
 	"github.com/stacklok/toolhive/pkg/vmcp"
 	"github.com/stacklok/toolhive/pkg/vmcp/aggregator"
@@ -540,3 +543,306 @@ func TestIntegration_ConflictResolutionStrategies(t *testing.T) {
 		assert.Equal(t, "backend1", result.Tools[0].BackendID)
 	})
 }
+
+// TestIntegration_AuditLogging tests that the vMCP server logs MCP operations
+// when audit middleware is enabled.
+// Note: This test does not use t.Parallel() because subtests share the same
+// server instance and audit log file, and must run sequentially.
+//
+//nolint:paralleltest // Subtests must run sequentially as they share server state
+func TestIntegration_AuditLogging(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	t.Cleanup(ctrl.Finish)
+
+	ctx := context.Background()
+
+	// Create temp file for audit logs
+	auditLogFile, err := os.CreateTemp("", "vmcp-audit-test-*.log")
+	require.NoError(t, err)
+	auditLogPath := auditLogFile.Name()
+	auditLogFile.Close()
+	t.Cleanup(func() {
+		os.Remove(auditLogPath)
+	})
+
+	// Create audit config that writes to temp file
+	auditConfig := &audit.Config{
+		Component:           "vmcp-server-test",
+		IncludeRequestData:  true,
+		IncludeResponseData: false,
+		MaxDataSize:         2048,
+		LogFile:             auditLogPath,
+	}
+
+	// Create mock backend client
+	mockBackendClient := mocks.NewMockBackendClient(ctrl)
+
+	// Define backend capabilities
+	backendCapabilities := &vmcp.CapabilityList{
+		Tools: []vmcp.Tool{
+			{
+				Name:        "get_weather",
+				Description: "Get weather information",
+				InputSchema: map[string]any{
+					"type": "object",
+					"properties": map[string]any{
+						"location": map[string]any{"type": "string"},
+					},
+				},
+				BackendID: "weather-service",
+			},
+		},
+		Resources: []vmcp.Resource{
+			{
+				URI:         "weather://current",
+				Name:        "Current Weather",
+				Description: "Current weather data",
+				MimeType:    "application/json",
+				BackendID:   "weather-service",
+			},
+		},
+		Prompts: []vmcp.Prompt{
+			{
+				Name:        "weather_summary",
+				Description: "Generate weather summary",
+				Arguments:   []vmcp.PromptArgument{},
+				BackendID:   "weather-service",
+			},
+		},
+	}
+
+	// Mock backend responses
+	mockBackendClient.EXPECT().
+		ListCapabilities(gomock.Any(), gomock.Any()).
+		Return(backendCapabilities, nil).
+		AnyTimes()
+
+	mockBackendClient.EXPECT().
+		CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+		Return(map[string]any{
+			"result": "Sunny, 72°F",
+		}, nil).
+		AnyTimes()
+
+	mockBackendClient.EXPECT().
+		ReadResource(gomock.Any(), gomock.Any(), gomock.Any()).
+		Return([]byte(`{"temp": 72, "condition": "sunny"}`), nil).
+		AnyTimes()
+
+	// Create backends
+	backends := []vmcp.Backend{
+		{
+			ID:   "weather-service",
+			Name: "Weather Service",
+		},
+	}
+
+	// Create router
+	rt := router.NewDefaultRouter()
+
+	// Create discovery manager
+	mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl)
+	mockDiscoveryMgr.EXPECT().
+		Discover(gomock.Any(), gomock.Any()).
+		DoAndReturn(func(_ context.Context, _ []vmcp.Backend) (*aggregator.AggregatedCapabilities, error) {
+			resolver := aggregator.NewPrefixConflictResolver("{workload}_")
+			agg := aggregator.NewDefaultAggregator(mockBackendClient, resolver, nil)
+			return agg.AggregateCapabilities(ctx, backends)
+		}).
+		AnyTimes()
+	mockDiscoveryMgr.EXPECT().Stop().AnyTimes()
+
+	// Helper function to read audit log file
+	readAuditLog := func() string {
+		data, err := os.ReadFile(auditLogPath)
+		if err != nil {
+			return ""
+		}
+		return string(data)
+	}
+
+	// Create server with audit config
+	srv, err := server.New(ctx, &server.Config{
+		Host:        "127.0.0.1",
+		Port:        0, // Random port
+		AuditConfig: auditConfig,
+	}, rt, mockBackendClient, mockDiscoveryMgr, backends, nil)
+	require.NoError(t, err)
+
+	// Start server
+	serverCtx, cancelServer := context.WithCancel(ctx)
+	t.Cleanup(cancelServer)
+
+	serverErrCh := make(chan error, 1)
+	go func() {
+		if err := srv.Start(serverCtx); err != nil && err != context.Canceled {
+			serverErrCh <- err
+		}
+	}()
+
+	// Wait for server ready
+	select {
+	case <-srv.Ready():
+	case err := <-serverErrCh:
+		t.Fatalf("Server failed to start: %v", err)
+	case <-time.After(5 * time.Second):
+		t.Fatal("Server timeout waiting for ready")
+	}
+
+	baseURL := "http://" + srv.Address()
+
+	// Test 1: Initialize request should be logged
+	t.Run("initialize request is logged", func(t *testing.T) {
+		initReq := map[string]any{
+			"method": "initialize",
+			"params": map[string]any{
+				"protocolVersion": "2024-11-05",
+				"capabilities":    map[string]any{},
+				"clientInfo": map[string]any{
+					"name":    "audit-test-client",
+					"version": "1.0.0",
+				},
+			},
+		}
+
+		reqBody, err := json.Marshal(initReq)
+		require.NoError(t, err)
+
+		resp, err := http.Post(baseURL+"/mcp", "application/json", bytes.NewReader(reqBody))
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+
+		// Wait for audit event to be written
+		time.Sleep(500 * time.Millisecond)
+
+		// Verify audit log contains initialize event
+		auditLog := readAuditLog()
+		assert.Contains(t, auditLog, "vmcp-server-test", "Should contain component name")
+		assert.Contains(t, auditLog, "\"method\":\"initialize\"", "Should log initialize method in request data")
+		assert.Contains(t, auditLog, "audit-test-client", "Should capture client name")
+	})
+
+	// Test 2: Tool list request should be logged
+	t.Run("tools/list request is logged", func(t *testing.T) {
+
+		toolsReq := map[string]any{
+			"method": "tools/list",
+		}
+
+		reqBody, err := json.Marshal(toolsReq)
+		require.NoError(t, err)
+
+		req, err := http.NewRequest("POST", baseURL+"/mcp", bytes.NewReader(reqBody))
+		require.NoError(t, err)
+		req.Header.Set("Content-Type", "application/json")
+
+		resp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		// Wait for audit event
+		time.Sleep(500 * time.Millisecond)
+
+		auditLog := readAuditLog()
+		assert.Contains(t, auditLog, "\"method\":\"tools/list\"", "Should log tools/list method in request data")
+		assert.Contains(t, auditLog, "vmcp-server-test", "Should contain component name")
+	})
+
+	// Test 3: Tool call should be logged
+	t.Run("tool call is logged", func(t *testing.T) {
+
+		toolCallReq := map[string]any{
+			"method": "tools/call",
+			"params": map[string]any{
+				"name": "get_weather",
+				"arguments": map[string]any{
+					"location": "San Francisco",
+				},
+			},
+		}
+
+		reqBody, err := json.Marshal(toolCallReq)
+		require.NoError(t, err)
+
+		req, err := http.NewRequest("POST", baseURL+"/mcp", bytes.NewReader(reqBody))
+		require.NoError(t, err)
+		req.Header.Set("Content-Type", "application/json")
+
+		resp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		// Wait for audit event
+		time.Sleep(500 * time.Millisecond)
+
+		auditLog := readAuditLog()
+		assert.Contains(t, auditLog, "\"method\":\"tools/call\"", "Should log tools/call method in request data")
+		assert.Contains(t, auditLog, "get_weather", "Should capture tool name in request data")
+		assert.Contains(t, auditLog, "San Francisco", "Should capture tool arguments in request data")
+		assert.Contains(t, auditLog, "vmcp-server-test", "Should contain component name")
+	})
+
+	// Test 4: Resource read should be logged
+	t.Run("resource read is logged", func(t *testing.T) {
+
+		resourceReq := map[string]any{
+			"method": "resources/read",
+			"params": map[string]any{
+				"uri": "weather://current",
+			},
+		}
+
+		reqBody, err := json.Marshal(resourceReq)
+		require.NoError(t, err)
+
+		req, err := http.NewRequest("POST", baseURL+"/mcp", bytes.NewReader(reqBody))
+		require.NoError(t, err)
+		req.Header.Set("Content-Type", "application/json")
+
+		resp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		// Wait for audit event
+		time.Sleep(500 * time.Millisecond)
+
+		auditLog := readAuditLog()
+		assert.Contains(t, auditLog, "\"method\":\"resources/read\"", "Should log resources/read method in request data")
+		assert.Contains(t, auditLog, "weather://current", "Should capture resource URI in request data")
+		assert.Contains(t, auditLog, "vmcp-server-test", "Should contain component name")
+	})
+
+	// Test 5: Verify audit events have required fields
+	t.Run("audit events contain required fields", func(t *testing.T) {
+		// Get all audit logs
+		auditLog := readAuditLog()
+
+		// Split into individual log lines
+		lines := strings.Split(strings.TrimSpace(auditLog), "\n")
+		require.Greater(t, len(lines), 0, "Should have at least one audit event")
+
+		// Parse first audit event
+		var auditEvent map[string]any
+		err := json.Unmarshal([]byte(lines[0]), &auditEvent)
+		require.NoError(t, err, "Audit log should be valid JSON")
+
+		// Verify required fields
+		assert.Contains(t, auditEvent, "audit_id", "Should have audit_id")
+		assert.Contains(t, auditEvent, "type", "Should have type")
+		assert.Contains(t, auditEvent, "logged_at", "Should have logged_at")
+		assert.Contains(t, auditEvent, "outcome", "Should have outcome")
+		assert.Contains(t, auditEvent, "component", "Should have component")
+		assert.Contains(t, auditEvent, "source", "Should have source")
+
+		// Verify component value
+		assert.Equal(t, "vmcp-server-test", auditEvent["component"])
+
+		// Verify source has network information
+		source, ok := auditEvent["source"].(map[string]any)
+		require.True(t, ok, "Source should be an object")
+		assert.Equal(t, "network", source["type"])
+		assert.Contains(t, source, "value", "Source should have IP address")
+	})
+}

pkg/vmcp/server/server.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/mark3labs/mcp-go/server"
 
+	"github.com/stacklok/toolhive/pkg/audit"
 	"github.com/stacklok/toolhive/pkg/auth"
 	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/telemetry"
@@ -87,6 +88,11 @@ type Config struct {
 	// TelemetryProvider is the optional telemetry provider.
 	// If nil, no telemetry is recorded.
 	TelemetryProvider *telemetry.Provider
+
+	// AuditConfig is the optional audit configuration.
+	// If nil, no audit logging is performed.
+	// Component should be set to "vmcp-server" to distinguish vMCP audit logs.
+	AuditConfig *audit.Config
 }
 
 // Server is the Virtual MCP Server that aggregates multiple backends.
@@ -386,20 +392,33 @@ func (s *Server) Start(ctx context.Context) error {
 		logger.Info("RFC 9728 OAuth discovery endpoints enabled at /.well-known/")
 	}
 
-	// MCP endpoint - apply middleware chain: auth → discovery → telemetry
+	// MCP endpoint - apply middleware chain: auth → audit → discovery → telemetry
 	var mcpHandler http.Handler = streamableServer
 
 	if s.config.TelemetryProvider != nil {
 		mcpHandler = s.config.TelemetryProvider.Middleware(s.config.Name, "streamable-http")(mcpHandler)
 		logger.Info("Telemetry middleware enabled for MCP endpoints")
 	}
 
-	// Apply discovery middleware (runs after auth middleware)
+	// Apply discovery middleware (runs after audit/auth middleware)
 	// Discovery middleware performs per-request capability aggregation with user context
 	// Pass sessionManager to enable session-based capability retrieval for subsequent requests
 	mcpHandler = discovery.Middleware(s.discoveryMgr, s.backends, s.sessionManager)(mcpHandler)
 	logger.Info("Discovery middleware enabled for lazy per-user capability discovery")
 
+	// Apply audit middleware if configured (runs after auth, before discovery)
+	if s.config.AuditConfig != nil {
+		auditor, err := audit.NewAuditorWithTransport(
+			s.config.AuditConfig,
+			"streamable-http", // vMCP uses streamable HTTP transport
+		)
+		if err != nil {
+			return fmt.Errorf("failed to create auditor: %w", err)
+		}
+		mcpHandler = auditor.Middleware(mcpHandler)
+		logger.Info("Audit middleware enabled for MCP endpoints")
+	}
+
 	// Apply authentication middleware if configured (runs first in chain)
 	if s.config.AuthMiddleware != nil {
 		mcpHandler = s.config.AuthMiddleware(mcpHandler)