fix transport for auditor middleware

Author: amirejaz

pkg/audit/auditor.go

@@ -15,6 +15,7 @@ import (
 	"github.com/stacklok/toolhive/pkg/auth"
 	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/mcp"
+	"github.com/stacklok/toolhive/pkg/transport/types"
 )
 
 // LevelAudit is a custom audit log level - between Info and Warn
@@ -35,12 +36,13 @@ func NewAuditLogger(w io.Writer) *slog.Logger {
 
 // Auditor handles audit logging for HTTP requests.
 type Auditor struct {
-	config      *Config
-	auditLogger *slog.Logger
+	config        *Config
+	auditLogger   *slog.Logger
+	transportType string // e.g., "sse", "streamable-http"
 }
 
-// NewAuditor creates a new Auditor with the given configuration.
-func NewAuditor(config *Config) (*Auditor, error) {
+// NewAuditorWithTransport creates a new Auditor with the given configuration and transport information.
+func NewAuditorWithTransport(config *Config, transportType string) (*Auditor, error) {
 	var logWriter io.Writer = os.Stdout // default to stdout
 
 	if config != nil {
@@ -54,11 +56,17 @@ func NewAuditor(config *Config) (*Auditor, error) {
 	}
 
 	return &Auditor{
-		config:      config,
-		auditLogger: NewAuditLogger(logWriter),
+		config:        config,
+		auditLogger:   NewAuditLogger(logWriter),
+		transportType: transportType,
 	}, nil
 }
 
+// isSSETransport checks if the current transport is SSE
+func (a *Auditor) isSSETransport() bool {
+	return a.transportType == types.TransportTypeSSE.String()
+}
+
 // responseWriter wraps http.ResponseWriter to capture response data and status.
 type responseWriter struct {
 	http.ResponseWriter
@@ -88,7 +96,7 @@ func (a *Auditor) Middleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		// Handle SSE endpoints specially - log the connection event immediately
 		// since SSE connections are long-lived and don't follow normal request/response pattern
-		if r.URL.Path == "/sse" {
+		if a.isSSETransport() {
 			// Log SSE connection event immediately
 			a.logSSEConnectionEvent(r)
 
@@ -164,7 +172,7 @@ func (a *Auditor) logAuditEvent(r *http.Request, rw *responseWriter, requestData
 	}
 
 	// Add metadata
-	a.addMetadata(event, r, duration, rw)
+	a.addMetadata(event, duration, rw)
 
 	// Add request/response data if configured
 	a.addEventData(event, r, rw, requestData)
@@ -184,7 +192,7 @@ func (a *Auditor) determineEventType(r *http.Request) string {
 	path := r.URL.Path
 
 	// Handle SSE connection establishment
-	if strings.Contains(path, "/sse") {
+	if a.isSSETransport() {
 		return EventTypeMCPInitialize
 	}
 
@@ -372,7 +380,7 @@ func (*Auditor) extractTarget(r *http.Request, eventType string) map[string]stri
 }
 
 // addMetadata adds metadata to the audit event.
-func (*Auditor) addMetadata(event *AuditEvent, r *http.Request, duration time.Duration, rw *responseWriter) {
+func (a *Auditor) addMetadata(event *AuditEvent, duration time.Duration, rw *responseWriter) {
 	if event.Metadata.Extra == nil {
 		event.Metadata.Extra = make(map[string]any)
 	}
@@ -381,11 +389,7 @@ func (*Auditor) addMetadata(event *AuditEvent, r *http.Request, duration time.Du
 	event.Metadata.Extra[MetadataExtraKeyDuration] = duration.Milliseconds()
 
 	// Add transport information
-	if strings.Contains(r.URL.Path, "/sse") {
-		event.Metadata.Extra[MetadataExtraKeyTransport] = "sse"
-	} else {
-		event.Metadata.Extra[MetadataExtraKeyTransport] = "http"
-	}
+	event.Metadata.Extra[MetadataExtraKeyTransport] = a.transportType
 
 	// Add response size if available
 	if rw.body != nil {
@@ -454,7 +458,7 @@ func (a *Auditor) logSSEConnectionEvent(r *http.Request) {
 
 	// Add metadata
 	event.Metadata.Extra = map[string]any{
-		"transport":  "sse",
+		"transport":  a.transportType,
 		"user_agent": r.Header.Get("User-Agent"),
 	}
 

pkg/audit/auditor_test.go

@@ -26,7 +26,7 @@ func init() {
 func TestNewAuditor(t *testing.T) {
 	t.Parallel()
 	config := &Config{}
-	auditor, err := NewAuditor(config)
+	auditor, err := NewAuditorWithTransport(config, "sse")
 
 	assert.NoError(t, err)
 	assert.NotNil(t, auditor)
@@ -36,7 +36,7 @@ func TestNewAuditor(t *testing.T) {
 func TestAuditorMiddlewareDisabled(t *testing.T) {
 	t.Parallel()
 	config := &Config{}
-	auditor, err := NewAuditor(config)
+	auditor, err := NewAuditorWithTransport(config, "sse")
 	require.NoError(t, err)
 
 	handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -61,7 +61,7 @@ func TestAuditorMiddlewareWithRequestData(t *testing.T) {
 		IncludeRequestData: true,
 		MaxDataSize:        1024,
 	}
-	auditor, err := NewAuditor(config)
+	auditor, err := NewAuditorWithTransport(config, "sse")
 	require.NoError(t, err)
 
 	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -91,7 +91,7 @@ func TestAuditorMiddlewareWithResponseData(t *testing.T) {
 		IncludeResponseData: true,
 		MaxDataSize:         1024,
 	}
-	auditor, err := NewAuditor(config)
+	auditor, err := NewAuditorWithTransport(config, "sse")
 	require.NoError(t, err)
 
 	responseData := `{"result": "success"}`
@@ -114,38 +114,43 @@ func TestAuditorMiddlewareWithResponseData(t *testing.T) {
 
 func TestDetermineEventType(t *testing.T) {
 	t.Parallel()
-	auditor, err := NewAuditor(&Config{})
-	require.NoError(t, err)
 
 	tests := []struct {
-		name     string
-		path     string
-		method   string
-		expected string
+		name      string
+		path      string
+		method    string
+		transport string
+		expected  string
 	}{
 		{
-			name:     "SSE endpoint",
-			path:     "/sse",
-			method:   "GET",
-			expected: EventTypeMCPInitialize,
+			name:      "SSE endpoint",
+			path:      "/sse",
+			method:    "GET",
+			transport: "sse",
+			expected:  EventTypeMCPInitialize,
 		},
 		{
-			name:     "MCP messages endpoint",
-			path:     "/messages",
-			method:   "POST",
-			expected: "mcp_request", // Since extractMCPMethod returns empty
+			name:      "MCP messages endpoint",
+			path:      "/messages",
+			method:    "POST",
+			transport: "streamable-http",
+			expected:  "mcp_request", // Since extractMCPMethod returns empty
 		},
 		{
-			name:     "Regular HTTP request",
-			path:     "/api/health",
-			method:   "GET",
-			expected: "http_request",
+			name:      "Regular HTTP request",
+			path:      "/api/health",
+			method:    "GET",
+			transport: "streamable-http",
+			expected:  "http_request",
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
+			auditor, err := NewAuditorWithTransport(&Config{}, tt.transport)
+			require.NoError(t, err)
+
 			req := httptest.NewRequest(tt.method, tt.path, nil)
 			result := auditor.determineEventType(req)
 			assert.Equal(t, tt.expected, result)
@@ -174,7 +179,7 @@ func TestMapMCPMethodToEventType(t *testing.T) {
 		{"unknown_method", "mcp_request"},
 	}
 
-	auditor, err := NewAuditor(&Config{})
+	auditor, err := NewAuditorWithTransport(&Config{}, "sse")
 	require.NoError(t, err)
 	for _, tt := range tests {
 		t.Run(tt.mcpMethod, func(t *testing.T) {
@@ -187,7 +192,7 @@ func TestMapMCPMethodToEventType(t *testing.T) {
 
 func TestDetermineOutcome(t *testing.T) {
 	t.Parallel()
-	auditor, err := NewAuditor(&Config{})
+	auditor, err := NewAuditorWithTransport(&Config{}, "sse")
 	require.NoError(t, err)
 
 	tests := []struct {
@@ -218,7 +223,7 @@ func TestDetermineOutcome(t *testing.T) {
 
 func TestGetClientIP(t *testing.T) {
 	t.Parallel()
-	auditor, err := NewAuditor(&Config{})
+	auditor, err := NewAuditorWithTransport(&Config{}, "sse")
 	require.NoError(t, err)
 
 	tests := []struct {
@@ -268,7 +273,7 @@ func TestGetClientIP(t *testing.T) {
 
 func TestExtractSubjects(t *testing.T) {
 	t.Parallel()
-	auditor, err := NewAuditor(&Config{})
+	auditor, err := NewAuditorWithTransport(&Config{}, "sse")
 	require.NoError(t, err)
 
 	t.Run("with JWT claims", func(t *testing.T) {
@@ -342,7 +347,7 @@ func TestDetermineComponent(t *testing.T) {
 	t.Run("with configured component", func(t *testing.T) {
 		t.Parallel()
 		config := &Config{Component: "custom-component"}
-		auditor, err := NewAuditor(config)
+		auditor, err := NewAuditorWithTransport(config, "sse")
 		require.NoError(t, err)
 
 		req := httptest.NewRequest("GET", "/test", nil)
@@ -354,7 +359,7 @@ func TestDetermineComponent(t *testing.T) {
 	t.Run("without configured component", func(t *testing.T) {
 		t.Parallel()
 		config := &Config{}
-		auditor, err := NewAuditor(config)
+		auditor, err := NewAuditorWithTransport(config, "sse")
 		require.NoError(t, err)
 
 		req := httptest.NewRequest("GET", "/test", nil)
@@ -366,7 +371,7 @@ func TestDetermineComponent(t *testing.T) {
 
 func TestExtractTarget(t *testing.T) {
 	t.Parallel()
-	auditor, err := NewAuditor(&Config{})
+	auditor, err := NewAuditorWithTransport(&Config{}, "sse")
 	require.NoError(t, err)
 
 	tests := []struct {
@@ -423,18 +428,17 @@ func TestExtractTarget(t *testing.T) {
 
 func TestAddMetadata(t *testing.T) {
 	t.Parallel()
-	auditor, err := NewAuditor(&Config{})
+	auditor, err := NewAuditorWithTransport(&Config{}, "sse")
 	require.NoError(t, err)
 
 	event := NewAuditEvent("test", EventSource{}, OutcomeSuccess, map[string]string{}, "test")
-	req := httptest.NewRequest("GET", "/sse/test", nil)
 	duration := 150 * time.Millisecond
 	rw := &responseWriter{
 		ResponseWriter: httptest.NewRecorder(),
 		body:           bytes.NewBufferString("test response"),
 	}
 
-	auditor.addMetadata(event, req, duration, rw)
+	auditor.addMetadata(event, duration, rw)
 
 	require.NotNil(t, event.Metadata.Extra)
 	assert.Equal(t, int64(150), event.Metadata.Extra[MetadataExtraKeyDuration])
@@ -450,7 +454,7 @@ func TestAddEventData(t *testing.T) {
 			IncludeRequestData:  true,
 			IncludeResponseData: true,
 		}
-		auditor, err := NewAuditor(config)
+		auditor, err := NewAuditorWithTransport(config, "sse")
 		require.NoError(t, err)
 
 		event := NewAuditEvent("test", EventSource{}, OutcomeSuccess, map[string]string{}, "test")
@@ -483,7 +487,7 @@ func TestAddEventData(t *testing.T) {
 			IncludeRequestData:  true,
 			IncludeResponseData: true,
 		}
-		auditor, err := NewAuditor(config)
+		auditor, err := NewAuditorWithTransport(config, "sse")
 		require.NoError(t, err)
 
 		event := NewAuditEvent("test", EventSource{}, OutcomeSuccess, map[string]string{}, "test")
@@ -511,7 +515,7 @@ func TestAddEventData(t *testing.T) {
 			IncludeRequestData:  false,
 			IncludeResponseData: false,
 		}
-		auditor, err := NewAuditor(config)
+		auditor, err := NewAuditorWithTransport(config, "sse")
 		require.NoError(t, err)
 
 		event := NewAuditEvent("test", EventSource{}, OutcomeSuccess, map[string]string{}, "test")
@@ -531,7 +535,7 @@ func TestResponseWriterCapture(t *testing.T) {
 		IncludeResponseData: true,
 		MaxDataSize:         10, // Small limit for testing
 	}
-	auditor, err := NewAuditor(config)
+	auditor, err := NewAuditorWithTransport(config, "sse")
 	require.NoError(t, err)
 
 	rw := &responseWriter{
@@ -568,7 +572,7 @@ func TestResponseWriterStatusCode(t *testing.T) {
 
 func TestExtractSourceWithHeaders(t *testing.T) {
 	t.Parallel()
-	auditor, err := NewAuditor(&Config{})
+	auditor, err := NewAuditorWithTransport(&Config{}, "sse")
 	require.NoError(t, err)
 
 	req := httptest.NewRequest("GET", "/test", nil)

pkg/audit/config.go

@@ -104,25 +104,26 @@ func (c *Config) ShouldAuditEvent(eventType string) bool {
 	return true
 }
 
-// CreateMiddleware creates an HTTP middleware from the audit configuration.
-func (c *Config) CreateMiddleware() (types.MiddlewareFunction, error) {
-	auditor, err := NewAuditor(c)
+// CreateMiddlewareWithTransport creates an HTTP middleware from the audit configuration with transport information.
+func (c *Config) CreateMiddlewareWithTransport(transportType string) (types.MiddlewareFunction, error) {
+	auditor, err := NewAuditorWithTransport(c, transportType)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create auditor: %w", err)
 	}
 	return auditor.Middleware, nil
 }
 
 // GetMiddlewareFromFile loads the audit configuration from a file and creates an HTTP middleware.
-func GetMiddlewareFromFile(path string) (func(http.Handler) http.Handler, error) {
+// Note: This function requires a transport type to be provided separately.
+func GetMiddlewareFromFile(path string, transportType string) (func(http.Handler) http.Handler, error) {
 	// Load the configuration
 	config, err := LoadFromFile(path)
 	if err != nil {
 		return nil, fmt.Errorf("failed to load audit config: %w", err)
 	}
 
-	// Create the middleware
-	return config.CreateMiddleware()
+	// Create the middleware with transport information
+	return config.CreateMiddlewareWithTransport(transportType)
 }
 
 // Validate validates the audit configuration.

pkg/audit/config_test.go

@@ -112,7 +112,7 @@ func TestCreateMiddleware(t *testing.T) {
 	t.Parallel()
 	config := &Config{}
 
-	middleware, err := config.CreateMiddleware()
+	middleware, err := config.CreateMiddlewareWithTransport("sse")
 	assert.NoError(t, err)
 	assert.NotNil(t, middleware)
 }
@@ -236,7 +236,7 @@ func TestConfigMinimalJSON(t *testing.T) {
 func TestGetMiddlewareFromFileError(t *testing.T) {
 	t.Parallel()
 	// Test with non-existent file
-	_, err := GetMiddlewareFromFile("/non/existent/file.json")
+	_, err := GetMiddlewareFromFile("/non/existent/file.json", "sse")
 	assert.Error(t, err)
 	assert.Contains(t, err.Error(), "failed to load audit config")
 }