security(mcp): prevent server requests from matching client pending responses

internal/mcp/client.go

@@ -7,11 +7,13 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"os"
 	"os/exec"
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/Gitlawb/zero/internal/execution"
@@ -49,20 +51,32 @@ type Client struct {
 	stdin   io.WriteCloser
 	reader  *messageReader
 	writer  *messageWriter
-	mu      sync.Mutex
 	closeMu sync.Mutex
+	idMu    sync.Mutex
 	nextID  int
 	cleanup func()
 
+	writeQueue      chan writeOp
+	writerOnce      sync.Once
+	droppedCourtesy atomic.Uint64
+
 	// dispatchMu guards the response-dispatch state shared with the single
-	// reader goroutine. It is never held across a blocking read.
+	// reader goroutine. It is never held across a blocking read or write.
 	dispatchMu sync.Mutex
 	readerOnce sync.Once
 	pending    map[int]chan dispatchResult
 	readErr    error
 	readDone   bool
 }
 
+type writeOp struct {
+	ctx     context.Context
+	message rpcMessage
+	done    chan error
+}
+
+const writeQueueCapacity = 32
+
 // dispatchResult carries one matched JSON-RPC response (or a terminal reader
 // error) to a waiting caller.
 type dispatchResult struct {
@@ -307,20 +321,20 @@ func (client *Client) request(ctx context.Context, method string, params any, ta
 		return err
 	}
 
-	// Allocate an id, register a response channel, and write the request while
-	// holding client.mu. The mutex serializes writes and id allocation but is
-	// released before the (potentially unbounded) wait for the response, so a
-	// hung server never holds the lock and blocks other callers/Close.
-	client.mu.Lock()
+	// Allocate an id and register a response channel. ID allocation and dispatch
+	// registrations are fast non-blocking operations. Message transmission is
+	// handled via writeMessage so a caller with a canceled context can stop
+	// waiting even when the peer is not draining its input.
+	client.idMu.Lock()
 	id := client.nextID
 	client.nextID++
+	client.idMu.Unlock()
 
 	responses := make(chan dispatchResult, 1)
 	client.dispatchMu.Lock()
 	if client.readDone {
 		readErr := client.readErr
 		client.dispatchMu.Unlock()
-		client.mu.Unlock()
 		if readErr != nil {
 			return readErr
 		}
@@ -329,16 +343,14 @@ func (client *Client) request(ctx context.Context, method string, params any, ta
 	client.pending[id] = responses
 	client.dispatchMu.Unlock()
 
-	if err := client.writer.write(rpcMessage{
+	if err := client.writeMessage(ctx, rpcMessage{
 		ID:     id,
 		Method: method,
 		Params: rawParams,
 	}); err != nil {
 		client.removePending(id)
-		client.mu.Unlock()
 		return err
 	}
-	client.mu.Unlock()
 
 	select {
 	case <-ctx.Done():
@@ -361,6 +373,53 @@ func (client *Client) request(ctx context.Context, method string, params any, ta
 	}
 }
 
+// ensureWriter lazily starts the single writer goroutine.
+func (client *Client) ensureWriter() {
+	client.writerOnce.Do(func() {
+		if client.writeQueue == nil {
+			client.writeQueue = make(chan writeOp, writeQueueCapacity)
+		}
+		go client.writeLoop()
+	})
+}
+
+func (client *Client) writeLoop() {
+	for op := range client.writeQueue {
+		if op.ctx != nil {
+			select {
+			case <-op.ctx.Done():
+				if op.done != nil {
+					op.done <- op.ctx.Err()
+				}
+				continue
+			default:
+			}
+		}
+		err := client.writer.write(op.message)
+		if op.done != nil {
+			op.done <- err
+		}
+	}
+}
+
+func (client *Client) writeMessage(ctx context.Context, message rpcMessage) error {
+	client.ensureWriter()
+	done := make(chan error, 1)
+	op := writeOp{ctx: ctx, message: message, done: done}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case client.writeQueue <- op:
+	}
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case err := <-done:
+		return err
+	}
+}
+
 // ensureReader lazily starts the single reader goroutine. It runs once per
 // client; subsequent calls are no-ops.
 func (client *Client) ensureReader() {
@@ -385,6 +444,30 @@ func (client *Client) readLoop() {
 			client.failAll(err)
 			return
 		}
+		// A message with a Method is a server-initiated request or notification.
+		// It must never be routed as a response to a pending client request.
+		if message.methodPresent || message.Method != "" {
+			if message.ID != nil && jsonRPCIDEchoable(message.ID) {
+				client.ensureWriter()
+				id := message.ID
+				method := message.Method
+				courtesy := writeOp{
+					message: rpcMessage{
+						ID: id,
+						Error: &rpcError{
+							Code:    -32601,
+							Message: fmt.Sprintf("Method %q not supported", method),
+						},
+					},
+				}
+				select {
+				case client.writeQueue <- courtesy:
+				default:
+					client.droppedCourtesy.Add(1)
+				}
+			}
+			continue
+		}
 		if message.ID == nil {
 			continue
 		}
@@ -471,12 +554,37 @@ func rpcIDMatches(value any, id int) bool {
 	}
 }
 
+// jsonRPCIDEchoable reports whether id is a valid JSON-RPC 2.0 identifier type
+// (string or finite number) that is safe to echo back.
+func jsonRPCIDEchoable(id any) bool {
+	if id == nil {
+		return false
+	}
+	switch v := id.(type) {
+	case string:
+		return true
+	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
+		return true
+	case float64:
+		return !math.IsNaN(v) && !math.IsInf(v, 0)
+	case json.Number:
+		parsed, err := v.Float64()
+		if err != nil || math.IsNaN(parsed) || math.IsInf(parsed, 0) {
+			return false
+		}
+		_, err = json.Marshal(v)
+		return err == nil
+	default:
+		return false
+	}
+}
+
 func (client *Client) notify(method string, params any) error {
 	rawParams, err := json.Marshal(params)
 	if err != nil {
 		return err
 	}
-	return client.writer.write(rpcMessage{
+	return client.writeMessage(context.Background(), rpcMessage{
 		Method: method,
 		Params: rawParams,
 	})