From 15a54a6c6df3e38692b35413c60f0804ef5f67df Mon Sep 17 00:00:00 2001 From: P4ST4S Date: Mon, 24 Aug 2026 00:39:59 +0200 Subject: [PATCH] feat(audit): finalize incomplete stdio requests --- CHANGELOG.md | 3 + internal/proxy/stdio.go | 83 +++++++++++++++++---- internal/proxy/stdio_test.go | 135 +++++++++++++++++++++++++++++++++++ 3 files changed, 208 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 25f9b8f..34c813b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,9 @@ All notable changes to mcp-audit are documented in this file. - Terminal audit `outcome` values and UUIDv7 `audit_operation_id` correlation, with an exactly-once operation finalizer shared by HTTP and stdio transports. +- Complete stdio terminal auditing for pending-call expiry, malformed upstream + messages, upstream termination/write failures, cancellation, and client + disconnects. - Security invariants and release gates for v1.2.0. - Installation cookbook with platform-specific notes in `INSTALL.md`. - VS Code stdio configuration example under `examples/vscode/`. diff --git a/internal/proxy/stdio.go b/internal/proxy/stdio.go index 2669064..d0ad029 100644 --- a/internal/proxy/stdio.go +++ b/internal/proxy/stdio.go @@ -65,7 +65,9 @@ func (p *StdioProxy) Run(ctx context.Context) error { } cleanupCtx, stopCleanup := context.WithCancel(ctx) defer stopCleanup() - go p.state.cleanupLoop(cleanupCtx, pendingCallTTL, pendingCallCleanupInterval) + go p.state.cleanupLoop(cleanupCtx, pendingCallTTL, pendingCallCleanupInterval, func(calls []pendingCall) { + p.finalizeCalls(calls, audit.OutcomeTimeout) + }) cmd := exec.Command("/bin/sh", "-c", p.config.Upstream) stdin, err := cmd.StdinPipe() @@ -106,11 +108,15 @@ func (p *StdioProxy) Run(ctx context.Context) error { select { case <-ctx.Done(): + stopCleanup() _ = p.shutdownUpstream(cmd, waitErr) p.waitForPipes(pipesDone) + p.finalizeAll(audit.OutcomeCancelled) return nil case err := <-waitErr: + stopCleanup() p.waitForPipes(pipesDone) + p.finalizeAll(audit.OutcomeUpstreamError) if err != nil { return fmt.Errorf("proxy: stdio: wait: %w", err) } @@ -163,6 +169,7 @@ func (p *StdioProxy) pipeClientToServer(ctx context.Context, src io.Reader, upst } if _, err := upstream.Write(append(line, '\n')); err != nil { p.log.Error("failed to write to upstream", "error", err) + p.finalizeAll(audit.OutcomeUpstreamError) return } select { @@ -173,6 +180,7 @@ func (p *StdioProxy) pipeClientToServer(ctx context.Context, src io.Reader, upst } if err := scanner.Err(); err != nil { p.log.Error("failed to read client stdin", "error", err) + p.finalizeAll(audit.OutcomeClientDisconnect) } } @@ -185,6 +193,7 @@ func (p *StdioProxy) pipeServerToClient(ctx context.Context, src io.Reader, clie clientMu.Unlock() if err != nil { p.log.Error("failed to write to client stdout", "error", err) + p.finalizeAll(audit.OutcomeClientDisconnect) return } p.observeServerMessage(line) @@ -196,6 +205,7 @@ func (p *StdioProxy) pipeServerToClient(ctx context.Context, src io.Reader, clie } if err := scanner.Err(); err != nil { p.log.Error("failed to read upstream stdout", "error", err) + p.finalizeAll(audit.OutcomeUpstreamError) } } @@ -218,7 +228,7 @@ func (p *StdioProxy) observeClientMessage(raw []byte) messageAction { for _, msg := range messages { if msg.Method != "" { toolName := toolNameFromParams(msg.Method, msg.Params) - call := p.newPendingCall(msg.Method, jsonRPCID(msg.ID), toolName, msg.Params, time.Now()) + call := p.newPendingCall(msg.Method, jsonRPCID(msg.ID), toolName, msg.Params, audit.DirectionServerToClient, time.Now()) if msg.Method == "tools/call" { decision := p.evaluatePolicy(toolName) p.recordPolicyDecision(decision) @@ -264,12 +274,13 @@ func (p *StdioProxy) observeServerMessage(raw []byte) { messages, err := decodeMessages(raw) if err != nil { p.log.Warn("failed to inspect server message", "error", err) + p.finalizeCalls(p.state.takeClients(), audit.OutcomeMalformedUpstreamResponse) return } for _, msg := range messages { if msg.Method != "" { toolName := toolNameFromParams(msg.Method, msg.Params) - call := p.newPendingCall(msg.Method, jsonRPCID(msg.ID), toolName, msg.Params, time.Now()) + call := p.newPendingCall(msg.Method, jsonRPCID(msg.ID), toolName, msg.Params, audit.DirectionClientToServer, time.Now()) if len(msg.ID) > 0 { p.state.rememberServer(string(msg.ID), call) continue @@ -289,7 +300,7 @@ func (p *StdioProxy) observeServerMessage(raw []byte) { } } -func (p *StdioProxy) newPendingCall(method, requestID, toolName string, params json.RawMessage, startedAt time.Time) pendingCall { +func (p *StdioProxy) newPendingCall(method, requestID, toolName string, params json.RawMessage, completionDirection string, startedAt time.Time) pendingCall { operation, err := audit.NewOperation(p.config.Audit, audit.Entry{ Method: method, RequestID: requestID, @@ -301,7 +312,19 @@ func (p *StdioProxy) newPendingCall(method, requestID, toolName string, params j if err != nil { p.log.Error("failed to start audit operation", "method", method, "error", err) } - return pendingCall{operation: operation, startedAt: startedAt} + return pendingCall{operation: operation, startedAt: startedAt, completionDirection: completionDirection} +} + +func (p *StdioProxy) finalizeAll(outcome audit.Outcome) { + p.finalizeCalls(p.state.takeAll(), outcome) +} + +func (p *StdioProxy) finalizeCalls(calls []pendingCall, outcome audit.Outcome) { + for _, call := range calls { + if err := p.record(call, outcome, call.completionDirection, nil, nil); err != nil { + p.log.Error("failed to finalize incomplete audit operation", "outcome", outcome, "error", err) + } + } } func (p *StdioProxy) record(call pendingCall, outcome audit.Outcome, direction string, result json.RawMessage, rpcErr *audit.RPCError) error { @@ -335,8 +358,9 @@ func (p *StdioProxy) recordPolicyDecision(decision policy.Decision) { } type pendingCall struct { - operation *audit.Operation - startedAt time.Time + operation *audit.Operation + startedAt time.Time + completionDirection string } type rpcState struct { @@ -352,7 +376,7 @@ func newRPCState() *rpcState { } } -func (s *rpcState) cleanupLoop(ctx context.Context, ttl time.Duration, interval time.Duration) { +func (s *rpcState) cleanupLoop(ctx context.Context, ttl time.Duration, interval time.Duration, onExpired func([]pendingCall)) { if ttl <= 0 { return } @@ -366,29 +390,62 @@ func (s *rpcState) cleanupLoop(ctx context.Context, ttl time.Duration, interval case <-ctx.Done(): return case now := <-ticker.C: - s.purgeExpired(now, ttl) + expired := s.takeExpired(now, ttl) + if len(expired) > 0 && onExpired != nil { + onExpired(expired) + } } } } func (s *rpcState) purgeExpired(now time.Time, ttl time.Duration) int { + return len(s.takeExpired(now, ttl)) +} + +func (s *rpcState) takeExpired(now time.Time, ttl time.Duration) []pendingCall { cutoff := now.Add(-ttl) s.mu.Lock() defer s.mu.Unlock() - purged := 0 + var expired []pendingCall for id, call := range s.clientPending { if call.startedAt.Before(cutoff) { delete(s.clientPending, id) - purged++ + expired = append(expired, call) } } for id, call := range s.serverPending { if call.startedAt.Before(cutoff) { delete(s.serverPending, id) - purged++ + expired = append(expired, call) } } - return purged + return expired +} + +func (s *rpcState) takeClients() []pendingCall { + s.mu.Lock() + defer s.mu.Unlock() + calls := make([]pendingCall, 0, len(s.clientPending)) + for id, call := range s.clientPending { + calls = append(calls, call) + delete(s.clientPending, id) + } + return calls +} + +func (s *rpcState) takeAll() []pendingCall { + s.mu.Lock() + defer s.mu.Unlock() + calls := make([]pendingCall, 0, len(s.clientPending)+len(s.serverPending)) + for id, call := range s.clientPending { + calls = append(calls, call) + delete(s.clientPending, id) + } + for id, call := range s.serverPending { + calls = append(calls, call) + delete(s.serverPending, id) + } + return calls } func (s *rpcState) rememberClient(id string, call pendingCall) { diff --git a/internal/proxy/stdio_test.go b/internal/proxy/stdio_test.go index fee3449..525ce36 100644 --- a/internal/proxy/stdio_test.go +++ b/internal/proxy/stdio_test.go @@ -1,7 +1,12 @@ package proxy import ( + "context" "encoding/json" + "errors" + "io" + "strings" + "sync" "testing" "time" @@ -92,6 +97,136 @@ func TestStdioPolicyDeniesToolCallBeforeUpstream(t *testing.T) { } } +func TestStdioMalformedUpstreamResponseFinalizesPendingCall(t *testing.T) { + store := &memoryAuditStore{} + proxy := NewStdioProxy(StdioConfig{ + Audit: audit.NewLogger(audit.LoggerConfig{Store: store, Transport: "stdio"}), + Limiter: middleware.NewRateLimiter(false, 0), + }) + + proxy.observeClientMessage([]byte(`{"jsonrpc":"2.0","id":1,"method":"resources/list"}`)) + proxy.observeServerMessage([]byte(`not-json`)) + + if len(store.entries) != 1 { + t.Fatalf("stored entries = %d, want 1", len(store.entries)) + } + entry := store.entries[0] + if entry.Outcome != audit.OutcomeMalformedUpstreamResponse { + t.Fatalf("outcome = %q, want malformed_upstream_response", entry.Outcome) + } + if entry.Direction != audit.DirectionServerToClient { + t.Fatalf("direction = %q, want server-to-client", entry.Direction) + } + if entry.AuditOperationID == "" { + t.Fatal("audit_operation_id is empty") + } + + proxy.observeServerMessage([]byte(`{"jsonrpc":"2.0","id":1,"result":{}}`)) + if len(store.entries) != 1 { + t.Fatalf("late response created a duplicate terminal entry: %d", len(store.entries)) + } +} + +func TestStdioExpiredPendingCallFinalizesAsTimeout(t *testing.T) { + store := &memoryAuditStore{} + proxy := NewStdioProxy(StdioConfig{ + Audit: audit.NewLogger(audit.LoggerConfig{Store: store, Transport: "stdio"}), + Limiter: middleware.NewRateLimiter(false, 0), + }) + now := time.Date(2026, 8, 24, 12, 0, 0, 0, time.UTC) + + proxy.observeClientMessage([]byte(`{"jsonrpc":"2.0","id":1,"method":"resources/read"}`)) + call, ok := proxy.state.takeClient("1") + if !ok { + t.Fatal("pending call was not registered") + } + call.startedAt = now.Add(-pendingCallTTL - time.Second) + proxy.state.rememberClient("1", call) + + proxy.finalizeCalls(proxy.state.takeExpired(now, pendingCallTTL), audit.OutcomeTimeout) + if len(store.entries) != 1 { + t.Fatalf("stored entries = %d, want 1", len(store.entries)) + } + if store.entries[0].Outcome != audit.OutcomeTimeout { + t.Fatalf("outcome = %q, want timeout", store.entries[0].Outcome) + } + if _, ok := proxy.state.takeClient("1"); ok { + t.Fatal("expired call remains pending") + } +} + +func TestStdioUpstreamTerminationFinalizesBothDirections(t *testing.T) { + store := &memoryAuditStore{} + proxy := NewStdioProxy(StdioConfig{ + Audit: audit.NewLogger(audit.LoggerConfig{Store: store, Transport: "stdio"}), + Limiter: middleware.NewRateLimiter(false, 0), + }) + + proxy.observeClientMessage([]byte(`{"jsonrpc":"2.0","id":1,"method":"resources/list"}`)) + proxy.observeServerMessage([]byte(`{"jsonrpc":"2.0","id":2,"method":"sampling/createMessage"}`)) + proxy.finalizeAll(audit.OutcomeUpstreamError) + + if len(store.entries) != 2 { + t.Fatalf("stored entries = %d, want 2", len(store.entries)) + } + directions := make(map[string]bool) + for _, entry := range store.entries { + if entry.Outcome != audit.OutcomeUpstreamError { + t.Fatalf("outcome = %q, want upstream_error", entry.Outcome) + } + directions[entry.Direction] = true + } + if !directions[audit.DirectionServerToClient] || !directions[audit.DirectionClientToServer] { + t.Fatalf("terminal directions = %#v, want both directions", directions) + } +} + +func TestStdioUpstreamWriteFailureFinalizesAcceptedCall(t *testing.T) { + store := &memoryAuditStore{} + proxy := NewStdioProxy(StdioConfig{ + Audit: audit.NewLogger(audit.LoggerConfig{Store: store, Transport: "stdio"}), + Limiter: middleware.NewRateLimiter(false, 0), + }) + + proxy.pipeClientToServer( + context.Background(), + strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"resources/list"}`+"\n"), + errorWriter{}, + io.Discard, + &sync.Mutex{}, + ) + + if len(store.entries) != 1 || store.entries[0].Outcome != audit.OutcomeUpstreamError { + t.Fatalf("entries = %#v, want one upstream_error", store.entries) + } +} + +func TestStdioClientWriteFailureFinalizesPendingCall(t *testing.T) { + store := &memoryAuditStore{} + proxy := NewStdioProxy(StdioConfig{ + Audit: audit.NewLogger(audit.LoggerConfig{Store: store, Transport: "stdio"}), + Limiter: middleware.NewRateLimiter(false, 0), + }) + proxy.observeClientMessage([]byte(`{"jsonrpc":"2.0","id":1,"method":"resources/list"}`)) + + proxy.pipeServerToClient( + context.Background(), + strings.NewReader(`{"jsonrpc":"2.0","id":1,"result":{}}`+"\n"), + errorWriter{}, + &sync.Mutex{}, + ) + + if len(store.entries) != 1 || store.entries[0].Outcome != audit.OutcomeClientDisconnect { + t.Fatalf("entries = %#v, want one client_disconnect", store.entries) + } +} + +type errorWriter struct{} + +func (errorWriter) Write([]byte) (int, error) { + return 0, errors.New("write failed") +} + type memoryAuditStore struct { entries []audit.Entry }