Skip to content

Commit 6f43bc9

Browse files
dongjiang1989dongjiang1989
committed
mcp: harden paginate() against infinite loops and param mutation
Add safety guards to the paginate() generic helper used by Tools(), Resources(), ResourceTemplates(), and Prompts() iterators: - Add ListMaxPages option to ClientOptions (0 means unlimited, matching the TypeScript SDK's listMaxPages semantics). A positive value caps the number of pages. - Detect cursor cycles via a seen-set to prevent infinite pagination when a server returns a repeating cursor. - Copy the caller's params struct via reflect so pagination does not mutate the original value. Fixes a latent safety gap: without these guards, a buggy or malicious server could cause unbounded CPU/memory consumption via cursor cycles or infinite NextCursor chains. Signed-off-by: dongjiang <dongjiang1989@126.com>
1 parent e761950 commit 6f43bc9

2 files changed

Lines changed: 195 additions & 9 deletions

File tree

‎mcp/client.go‎

Lines changed: 44 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,12 @@ type ClientOptions struct {
204204
// reset" guidance, letting a transient miss pass without tearing down an
205205
// otherwise live session. Has no effect unless KeepAlive is non-zero.
206206
KeepAliveFailureThreshold int
207+
// ListMaxPages is the maximum number of pages to fetch during automatic
208+
// pagination of list operations (Tools, Resources, ResourceTemplates,
209+
// Prompts). A value of 0 (the default) means unlimited. A positive value
210+
// caps the number of pages. This prevents runaway pagination loops caused
211+
// by server-side cursor cycles.
212+
ListMaxPages int
207213
}
208214

209215
// toolContextKeyType is the context key type for passing tool definitions
@@ -1550,7 +1556,7 @@ func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) ite
15501556
if params == nil {
15511557
params = &ListToolsParams{}
15521558
}
1553-
return paginate(ctx, params, cs.ListTools, func(res *ListToolsResult) []*Tool {
1559+
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListTools, func(res *ListToolsResult) []*Tool {
15541560
return res.Tools
15551561
})
15561562
}
@@ -1563,7 +1569,7 @@ func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesPar
15631569
if params == nil {
15641570
params = &ListResourcesParams{}
15651571
}
1566-
return paginate(ctx, params, cs.ListResources, func(res *ListResourcesResult) []*Resource {
1572+
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListResources, func(res *ListResourcesResult) []*Resource {
15671573
return res.Resources
15681574
})
15691575
}
@@ -1576,7 +1582,7 @@ func (cs *ClientSession) ResourceTemplates(ctx context.Context, params *ListReso
15761582
if params == nil {
15771583
params = &ListResourceTemplatesParams{}
15781584
}
1579-
return paginate(ctx, params, cs.ListResourceTemplates, func(res *ListResourceTemplatesResult) []*ResourceTemplate {
1585+
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListResourceTemplates, func(res *ListResourceTemplatesResult) []*ResourceTemplate {
15801586
return res.ResourceTemplates
15811587
})
15821588
}
@@ -1589,16 +1595,33 @@ func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams)
15891595
if params == nil {
15901596
params = &ListPromptsParams{}
15911597
}
1592-
return paginate(ctx, params, cs.ListPrompts, func(res *ListPromptsResult) []*Prompt {
1598+
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListPrompts, func(res *ListPromptsResult) []*Prompt {
15931599
return res.Prompts
15941600
})
15951601
}
15961602

15971603
// paginate is a generic helper function to provide a paginated iterator.
1598-
func paginate[P listParams, R listResult[T], T any](ctx context.Context, params P, listFunc func(context.Context, P) (R, error), items func(R) []*T) iter.Seq2[*T, error] {
1604+
//
1605+
// It fetches pages by calling listFunc until the result has no NextCursor,
1606+
// maxPages is exceeded (if non-zero), or a cursor cycle is detected.
1607+
// The caller's params struct is not mutated; a local copy is used instead.
1608+
func paginate[P listParams, R listResult[T], T any](ctx context.Context, params P, maxPages int, listFunc func(context.Context, P) (R, error), items func(R) []*T) iter.Seq2[*T, error] {
15991609
return func(yield func(*T, error) bool) {
1610+
// Copy the underlying struct so we don't mutate the caller's params.
1611+
// P is always a pointer to a struct (e.g. *ListToolsParams).
1612+
// We use reflect to create a shallow copy of the pointed-to struct.
1613+
localParams := params
1614+
if v := reflect.ValueOf(params); v.Kind() == reflect.Pointer {
1615+
cp := reflect.New(v.Type().Elem())
1616+
cp.Elem().Set(v.Elem())
1617+
localParams = cp.Interface().(P)
1618+
}
1619+
var seen map[string]bool
1620+
pages := 0
1621+
16001622
for {
1601-
res, err := listFunc(ctx, params)
1623+
pages++
1624+
res, err := listFunc(ctx, localParams)
16021625
if err != nil {
16031626
yield(nil, err)
16041627
return
@@ -1612,7 +1635,21 @@ func paginate[P listParams, R listResult[T], T any](ctx context.Context, params
16121635
if nextCursorVal == nil || *nextCursorVal == "" {
16131636
return
16141637
}
1615-
*params.cursorPtr() = *nextCursorVal
1638+
// Check max pages limit. 0 means unlimited.
1639+
if maxPages > 0 && pages >= maxPages {
1640+
yield(nil, fmt.Errorf("mcp: pagination exceeded maximum page limit of %d", maxPages))
1641+
return
1642+
}
1643+
// Detect cursor cycles to prevent infinite loops.
1644+
if seen == nil {
1645+
seen = make(map[string]bool)
1646+
}
1647+
if seen[*nextCursorVal] {
1648+
yield(nil, fmt.Errorf("mcp: pagination detected cursor cycle: %q", *nextCursorVal))
1649+
return
1650+
}
1651+
seen[*nextCursorVal] = true
1652+
*localParams.cursorPtr() = *nextCursorVal
16161653
}
16171654
}
16181655
}

‎mcp/client_test.go‎

Lines changed: 151 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -147,7 +147,7 @@ func TestClientPaginateBasic(t *testing.T) {
147147

148148
var gotItems []*Item
149149
var iterationErr error
150-
seq := paginate(ctx, params, listFunc, func(r *ListTestResult) []*Item { return r.Items })
150+
seq := paginate(ctx, params, 0, listFunc, func(r *ListTestResult) []*Item { return r.Items })
151151
for item, err := range seq {
152152
if err != nil {
153153
iterationErr = err
@@ -200,7 +200,7 @@ func TestClientPaginateVariousPageSizes(t *testing.T) {
200200
return res, nil
201201
}
202202
var gotItems []*Item
203-
seq := paginate(ctx, &ListTestParams{}, listFunc, func(r *ListTestResult) []*Item { return r.Items })
203+
seq := paginate(ctx, &ListTestParams{}, 0, listFunc, func(r *ListTestResult) []*Item { return r.Items })
204204
for item, err := range seq {
205205
if err != nil {
206206
t.Fatalf("paginate() unexpected error during iteration: %v", err)
@@ -214,6 +214,155 @@ func TestClientPaginateVariousPageSizes(t *testing.T) {
214214
}
215215
}
216216

217+
func TestPaginateMaxPages(t *testing.T) {
218+
ctx := context.Background()
219+
// Create 10 pages of results (1 item each, all with next cursor).
220+
all := allItems[:10]
221+
results := generatePaginatedResults(all, 1)
222+
223+
listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
224+
if len(results) == 0 {
225+
t.Fatal("listFunc called more times than expected")
226+
}
227+
res := results[0]
228+
results = results[1:]
229+
return res, nil
230+
}
231+
232+
t.Run("UnlimitedByDefault", func(t *testing.T) {
233+
// With maxPages=0 (unlimited), 10 pages should succeed.
234+
results = generatePaginatedResults(all, 1)
235+
var got []*Item
236+
var iterErr error
237+
seq := paginate(ctx, &ListTestParams{}, 0, listFunc, func(r *ListTestResult) []*Item { return r.Items })
238+
for item, err := range seq {
239+
if err != nil {
240+
iterErr = err
241+
break
242+
}
243+
got = append(got, item)
244+
}
245+
if iterErr != nil {
246+
t.Fatalf("unexpected error: %v", iterErr)
247+
}
248+
if len(got) != len(all) {
249+
t.Fatalf("got %d items, want %d", len(got), len(all))
250+
}
251+
})
252+
253+
t.Run("MaxPagesExceeded", func(t *testing.T) {
254+
results = generatePaginatedResults(all, 1)
255+
var got []*Item
256+
var iterErr error
257+
// maxPages=3: should stop after 3 pages with an error.
258+
seq := paginate(ctx, &ListTestParams{}, 3, listFunc, func(r *ListTestResult) []*Item { return r.Items })
259+
for item, err := range seq {
260+
if err != nil {
261+
iterErr = err
262+
break
263+
}
264+
got = append(got, item)
265+
}
266+
if iterErr == nil {
267+
t.Fatal("expected pagination error, got nil")
268+
}
269+
if len(got) != 3 {
270+
t.Fatalf("got %d items, want 3", len(got))
271+
}
272+
})
273+
274+
t.Run("PositiveLimitNotReached", func(t *testing.T) {
275+
results = generatePaginatedResults(all, 1)
276+
var got []*Item
277+
var iterErr error
278+
// maxPages=100: limit not reached for 10 pages, should succeed.
279+
seq := paginate(ctx, &ListTestParams{}, 100, listFunc, func(r *ListTestResult) []*Item { return r.Items })
280+
for item, err := range seq {
281+
if err != nil {
282+
iterErr = err
283+
break
284+
}
285+
got = append(got, item)
286+
}
287+
if iterErr != nil {
288+
t.Fatalf("unexpected error: %v", iterErr)
289+
}
290+
if len(got) != len(all) {
291+
t.Fatalf("got %d items, want %d", len(got), len(all))
292+
}
293+
})
294+
}
295+
296+
func TestPaginateCursorCycle(t *testing.T) {
297+
ctx := context.Background()
298+
callCount := 0
299+
// Server returns a cursor cycle: page1 → cursorA, page2 → cursorB, page3 → cursorA again.
300+
listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
301+
callCount++
302+
switch callCount {
303+
case 1:
304+
return &ListTestResult{Items: []*Item{{Name: "a"}}, NextCursor: "cursorA"}, nil
305+
case 2:
306+
return &ListTestResult{Items: []*Item{{Name: "b"}}, NextCursor: "cursorB"}, nil
307+
case 3:
308+
// Cycle: cursorA was already seen.
309+
return &ListTestResult{Items: []*Item{{Name: "c"}}, NextCursor: "cursorA"}, nil
310+
default:
311+
t.Fatal("listFunc called after cycle should have been detected")
312+
return nil, fmt.Errorf("unreachable")
313+
}
314+
}
315+
316+
var got []*Item
317+
var iterErr error
318+
seq := paginate(ctx, &ListTestParams{}, -1, listFunc, func(r *ListTestResult) []*Item { return r.Items })
319+
for item, err := range seq {
320+
if err != nil {
321+
iterErr = err
322+
break
323+
}
324+
got = append(got, item)
325+
}
326+
if iterErr == nil {
327+
t.Fatal("expected cursor cycle error, got nil")
328+
}
329+
if len(got) != 3 {
330+
t.Fatalf("got %d items, want 3 (items on cycling page are yielded before detection)", len(got))
331+
}
332+
}
333+
334+
func TestPaginateNoMutation(t *testing.T) {
335+
ctx := context.Background()
336+
params := &ListTestParams{Cursor: "initial"}
337+
results := generatePaginatedResults(allItems[:3], 1)
338+
339+
listFunc := func(ctx context.Context, p *ListTestParams) (*ListTestResult, error) {
340+
if len(results) == 0 {
341+
return &ListTestResult{Items: nil, NextCursor: ""}, nil
342+
}
343+
res := results[0]
344+
results = results[1:]
345+
return res, nil
346+
}
347+
348+
var got []*Item
349+
seq := paginate(ctx, params, -1, listFunc, func(r *ListTestResult) []*Item { return r.Items })
350+
for item, err := range seq {
351+
if err != nil {
352+
t.Fatalf("unexpected error: %v", err)
353+
}
354+
got = append(got, item)
355+
}
356+
357+
// The original params should not have been mutated.
358+
if params.Cursor != "initial" {
359+
t.Errorf("params.Cursor was mutated to %q, want %q", params.Cursor, "initial")
360+
}
361+
if len(got) != 3 {
362+
t.Fatalf("got %d items, want 3", len(got))
363+
}
364+
}
365+
217366
func TestClientCapabilities(t *testing.T) {
218367
testCases := []struct {
219368
name string

0 commit comments

Comments
 (0)