Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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/`.
Expand Down
83 changes: 70 additions & 13 deletions internal/proxy/stdio.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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 {
Expand All @@ -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)
}
}

Expand All @@ -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)
Expand All @@ -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)
}
}

Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -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
}
Expand All @@ -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) {
Expand Down
135 changes: 135 additions & 0 deletions internal/proxy/stdio_test.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
package proxy

import (
"context"
"encoding/json"
"errors"
"io"
"strings"
"sync"
"testing"
"time"

Expand Down Expand Up @@ -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
}
Expand Down
Loading