1package mcp
2
3import (
4 "bytes"
5 "context"
6 "encoding/json"
7 "errors"
8 "io"
9 "net/url"
10 "strings"
11 "testing"
12 "time"
13
14 "github.com/bnema/zut/packages/core"
15)
16
17// Regression tests for the second review round.
18
19// ---- GO-015: write cancellation ----
20
21func TestStdioQueuedFrameIsWithdrawnButStartedWriteFailsTransport(t *testing.T) {
22 c, events := connectFake(t, nil)
23 st := c.t.(*stdioTransport)
24
25 // The server stops reading stdin.
26 stallCtx, stallCancel := context.WithCancel(context.Background())
27 defer stallCancel()
28 go func() { _, _ = c.CallTool(stallCtx, "stall", []byte(`{}`)) }()
29 waitEvent(t, events, "stalled")
30
31 // Frame A is far larger than the pipe buffer: its write starts and blocks.
32 started := st.writesStarted.Load()
33 aCtx, aCancel := context.WithCancel(context.Background())
34 big := bytes.Repeat([]byte("a"), 4<<20)
35 aDone := make(chan struct {
36 sent bool
37 err error
38 }, 1)
39 go func() {
40 sent, err := st.writeFrame(aCtx, big)
41 aDone <- struct {
42 sent bool
43 err error
44 }{sent, err}
45 }()
46 eventually(t, "frame A to start writing", func() bool { return st.writesStarted.Load() == started+1 })
47
48 // Frame B queues behind A. Its caller gives up before B is dispatched:
49 // B must be withdrawn, never written, and the transport must stay usable.
50 bCtx, bCancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
51 defer bCancel()
52 sent, err := st.writeFrame(bCtx, []byte(`{"jsonrpc":"2.0","method":"notifications/never"}`))
53 if sent || !errors.Is(err, context.DeadlineExceeded) {
54 t.Fatalf("queued frame: sent=%v err=%v, want withdrawn with deadline error", sent, err)
55 }
56 select {
57 case <-st.done():
58 t.Fatal("withdrawing a queued frame must not fail the transport")
59 default:
60 }
61
62 // Now A's caller gives up mid-write: the frame cannot be withdrawn, so the
63 // transport fails and stdin is closed.
64 aCancel()
65 select {
66 case r := <-aDone:
67 if !r.sent || !errors.Is(r.err, context.Canceled) {
68 t.Errorf("started frame: sent=%v err=%v", r.sent, r.err)
69 }
70 case <-time.After(5 * time.Second):
71 t.Fatal("caller still blocked after cancellation")
72 }
73 select {
74 case <-st.done():
75 case <-time.After(5 * time.Second):
76 t.Fatal("transport still live with a half-written frame")
77 }
78 if err := st.failure(); err == nil || !strings.Contains(err.Error(), "interrupted") {
79 t.Errorf("failure = %v", err)
80 }
81 // B was never handed to the writer.
82 if got := st.writesStarted.Load(); got != started+1 {
83 t.Errorf("writes started = %d, want %d (frame B must never be written)", got, started+1)
84 }
85}
86
87func TestManagerReconnectsAfterInterruptedWrite(t *testing.T) {
88 // Cancellation is driven by the test, not by toolTimeoutSec: a timer could
89 // fire before the 4 MiB frame is even queued on a slow (-race) runner, and
90 // a withdrawn frame correctly leaves the transport up.
91 m, events := startFake(t, nil, nil)
92 st := m.servers["fake"].live().t.(*stdioTransport)
93 // Park the server: the call is delivered, then it stops reading stdin.
94 // The stall call stays pending (no cancel notification is queued) and
95 // ends when the transport fails.
96 stall, echo := findTool(t, m, "mcp__fake__stall"), findTool(t, m, "mcp__fake__echo")
97 stallDone := make(chan struct{})
98 go func() {
99 defer close(stallDone)
100 _, _ = stall.Execute(context.Background(), json.RawMessage(`{}`), nil)
101 }()
102 waitEvent(t, events, "stalled")
103 // A request too large for the pipe is cancelled once its write started.
104 started := st.writesStarted.Load()
105 bigCtx, bigCancel := context.WithCancel(context.Background())
106 bigDone := make(chan core.ToolResult, 1)
107 go func() {
108 big := `{"text":"` + strings.Repeat("a", 4<<20) + `"}`
109 res, _ := echo.Execute(bigCtx, json.RawMessage(big), nil)
110 bigDone <- res
111 }()
112 eventually(t, "the big frame write to start", func() bool { return st.writesStarted.Load() > started })
113 bigCancel()
114 if res := <-bigDone; !res.IsError {
115 t.Fatalf("big call result = %q", resultText(res))
116 }
117 eventually(t, "status to show the dropped connection", func() bool { return !m.Status()[0].Connected })
118 <-stallDone
119 // The next call starts a fresh server and works.
120 if got := resultText(execTool(t, findTool(t, m, "mcp__fake__echo"), `{"text":"again"}`)); got != "echo: again" {
121 t.Fatalf("after reconnect = %q", got)
122 }
123 if n := countEvents(events, "start"); n != 2 {
124 t.Errorf("server starts = %d, want 2", n)
125 }
126}
127
128// ---- GO-004, GO-018: stderr redaction ----
129
130func newTestStderrTransport(cfg ServerConfig) *stdioTransport {
131 return &stdioTransport{red: newRedactor(cfg), stderrDone: make(chan struct{})}
132}
133
134func TestStderrSecretStraddlingTruncationIsRedacted(t *testing.T) {
135 const secret = "SECRETVALUE-0123456789"
136 tr := newTestStderrTransport(ServerConfig{Env: map[string]string{"K": secret}})
137 // The secret starts before the 4096-byte display cut and ends after it.
138 line := strings.Repeat("x", maxStderrLine-len(secret)/2) + secret + strings.Repeat("y", 300) + "\nnext line\n"
139 tr.readStderr(strings.NewReader(line))
140 tail := tr.stderrTail()
141 for _, frag := range []string{"SECRET", "VALUE", "0123", secret[:8]} {
142 if strings.Contains(tail, frag) {
143 t.Errorf("secret fragment %q leaked: %.80q…", frag, tail[maxStderrLine-60:])
144 }
145 }
146 if !strings.Contains(tail, redactedMarker) || !strings.Contains(tail, "[line truncated]") || !strings.Contains(tail, "next line") {
147 t.Errorf("unexpected tail shape (len %d)", len(tail))
148 }
149}
150
151func TestStderrSecretOnTheCutIsNotPartiallyShown(t *testing.T) {
152 // A line whose redacted form still ends in a secret prefix would leak it;
153 // trimPartialSecret removes such a tail.
154 red := newRedactor(ServerConfig{Env: map[string]string{"K": "TOPSECRET99"}})
155 for _, in := range []string{"abc TOP", "abc TOPSECRET9", "abc TOPSECRET", "abc"} {
156 got := red.trimPartialSecret(in)
157 if strings.Contains(got, "TOP") {
158 t.Errorf("trimPartialSecret(%q) = %q", in, got)
159 }
160 }
161 if got := red.trimPartialSecret("keep TOPS and"); got != "keep TOPS and" {
162 t.Errorf("interior text changed: %q", got)
163 }
164 var nilRed *redactor
165 if got := nilRed.trimPartialSecret("x"); got != "x" {
166 t.Error("nil redactor must pass text through")
167 }
168}
169
170func TestStdioStderrSecretEndToEnd(t *testing.T) {
171 const secret = "E2ESECRET-abcdefghij"
172 c, _ := connectFake(t, map[string]string{"FAKE_SECRET": secret})
173 if got := callText(t, c, "stderr_secret", `{}`); got != "written" {
174 t.Fatal(got)
175 }
176 st := c.t.(*stdioTransport)
177 eventually(t, "stderr to be read", func() bool { return strings.Contains(st.stderrTail(), "after-secret") })
178 tail := st.stderrTail()
179 if strings.Contains(tail, "E2ESECRET") || strings.Contains(tail, "abcdef") {
180 t.Errorf("secret (or prefix) leaked")
181 }
182}
183
184func TestRedactionSurvivesControlSequencesInsideSecret(t *testing.T) {
185 // Case 1: the server's output splits the secret with escape sequences.
186 red := newRedactor(ServerConfig{Env: map[string]string{"K": "token-abcdef123"}})
187 out := red.scrub("leak: tok\x1b[0men-abc\x1b]0;x\x07def123 end")
188 if strings.Contains(out, "abcdef123") || strings.Contains(out, "token") || !strings.Contains(out, redactedMarker) {
189 t.Errorf("out = %q", out)
190 }
191 // Case 2: the configured secret itself contains control bytes.
192 red = newRedactor(ServerConfig{Env: map[string]string{"K": "se\x1b[1mcret-value-9"}})
193 out = red.scrub("echo se\x1b[1mcret-value-9 / secret-value-9")
194 if strings.Contains(out, "cret-value-9") {
195 t.Errorf("out = %q", out)
196 }
197}
198
199// ---- GO-016: request-building errors ----
200
201func TestHTTPRequestBuildErrorIsScrubbed(t *testing.T) {
202 const secret = "QUERYSECRET-42"
203 cfg := ServerConfig{Type: "http", URL: "http://bad host/path?token=" + secret}
204 tr, err := newHTTPTransport("x", cfg, nil)
205 if err != nil {
206 t.Fatal(err)
207 }
208 defer tr.close()
209 _, err = tr.request(context.Background(), 1, "initialize", nil)
210 if err == nil {
211 t.Fatal("want an error for an unparsable URL")
212 }
213 if strings.Contains(err.Error(), secret) {
214 t.Errorf("error leaks the query token: %v", err)
215 }
216 // The same path also feeds close() and the GET stream.
217 if _, err := tr.newRequest(context.Background(), "GET", nil); err == nil || strings.Contains(err.Error(), secret) {
218 t.Errorf("newRequest error = %v", err)
219 }
220}
221
222// ---- GO-017: encoded URL secrets ----
223
224func TestURLSecretsIncludePercentDecodedForms(t *testing.T) {
225 raw := "https://us%40er:p%40ssw0rd@h.example/p?api_key=a%2Bb%20c&other=plain-value&semi=x;y=z-token"
226 got := map[string]bool{}
227 for _, v := range urlSecrets(raw) {
228 got[v] = true
229 }
230 for _, want := range []string{"us%40er", "us@er", "p%40ssw0rd", "p@ssw0rd", "a%2Bb%20c", "a+b c", "plain-value"} {
231 if !got[want] {
232 t.Errorf("urlSecrets missing %q (have %v)", want, got)
233 }
234 }
235 red := newRedactor(ServerConfig{URL: raw})
236 out := red.scrub("server said p@ssw0rd and a+b c and a%2Bb%20c and plain-value")
237 for _, secret := range []string{"p@ssw0rd", "a+b c", "a%2Bb%20c", "plain-value"} {
238 if strings.Contains(out, secret) {
239 t.Errorf("%q leaked in %q", secret, out)
240 }
241 }
242 // A URL that net/url refuses to parse still yields its secrets.
243 unparsable := strings.Replace("http://bad_host/x?tok=SECRETQ99", "_", " ", 1) // space in the host
244 if _, err := url.Parse(unparsable); err == nil {
245 t.Fatal("test premise: URL should not parse")
246 }
247 bad := urlSecrets(unparsable)
248 found := false
249 for _, v := range bad {
250 found = found || v == "SECRETQ99"
251 }
252 if !found {
253 t.Errorf("unparsable URL secrets = %v", bad)
254 }
255}
256
257// ---- GO-014: per-field budget ----
258
259func TestBudgetedFieldIsBoundedBeforeFormatting(t *testing.T) {
260 b := &budgetedText{}
261 huge := strings.Repeat("é", 10<<20) // 20 MiB of two-byte runes
262 got := b.field(huge)
263 if len(got) > maxResultTextBytes+1 {
264 t.Errorf("field not clipped: %d bytes", len(got))
265 }
266 if !strings.HasPrefix(huge, got) && got != "" {
267 t.Error("clipped field is not a prefix")
268 }
269 if clipUTF8("日本語", 4) != "日" {
270 t.Errorf("clipUTF8 split a rune: %q", clipUTF8("日本語", 4))
271 }
272 b.write(strings.Repeat("a", maxResultTextBytes-10))
273 if got := b.field(strings.Repeat("b", 1000)); len(got) > 11 {
274 t.Errorf("field ignores remaining budget: %d", len(got))
275 }
276}
277
278func TestListResourcesHugeDescription(t *testing.T) {
279 m, _ := startFake(t, map[string]string{"FAKE_HUGE_DESC": "5242880"}, nil) // 5 MiB
280 text := resultText(execTool(t, findTool(t, m, ListResourcesToolName), `{}`))
281 if len(text) > maxResultTextBytes+len("\n... [truncated at 51200 bytes]") {
282 t.Errorf("output %d bytes exceeds the budget", len(text))
283 }
284 if !strings.Contains(text, "mem://huge") || !strings.Contains(text, "[truncated") {
285 t.Errorf("unexpected output head: %.100q", text)
286 }
287}
288
289// ---- readStderr bookkeeping ----
290
291func TestReadStderrHandlesEOFWithoutNewline(t *testing.T) {
292 tr := newTestStderrTransport(ServerConfig{})
293 tr.readStderr(io.MultiReader(strings.NewReader("one\ntwo"), strings.NewReader("")))
294 if got := tr.stderrTail(); got != "one\ntwo" {
295 t.Errorf("tail = %q", got)
296 }
297}