Skip to content

Commit 481cc4a

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 with a default of DefaultListMaxPages (64), matching the TypeScript SDK's DEFAULT_LIST_MAX_PAGES. A value of 0 uses the default; a negative value means unlimited. - 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 481cc4a

2 files changed

Lines changed: 238 additions & 9 deletions

File tree

‎mcp/client.go‎

Lines changed: 54 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -204,8 +204,19 @@ 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 uses the default of [DefaultListMaxPages] (64).
210+
// A negative value means unlimited. A positive value caps the number of
211+
// pages. This prevents runaway pagination loops caused by server-side
212+
// cursor cycles.
213+
ListMaxPages int
207214
}
208215

216+
// DefaultListMaxPages is the default value for [ClientOptions.ListMaxPages],
217+
// matching the TypeScript SDK's DEFAULT_LIST_MAX_PAGES.
218+
const DefaultListMaxPages = 64
219+
209220
// toolContextKeyType is the context key type for passing tool definitions
210221
// from CallTool to the transport layer.
211222
type toolContextKeyType struct{}
@@ -1550,7 +1561,7 @@ func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) ite
15501561
if params == nil {
15511562
params = &ListToolsParams{}
15521563
}
1553-
return paginate(ctx, params, cs.ListTools, func(res *ListToolsResult) []*Tool {
1564+
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListTools, func(res *ListToolsResult) []*Tool {
15541565
return res.Tools
15551566
})
15561567
}
@@ -1563,7 +1574,7 @@ func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesPar
15631574
if params == nil {
15641575
params = &ListResourcesParams{}
15651576
}
1566-
return paginate(ctx, params, cs.ListResources, func(res *ListResourcesResult) []*Resource {
1577+
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListResources, func(res *ListResourcesResult) []*Resource {
15671578
return res.Resources
15681579
})
15691580
}
@@ -1576,7 +1587,7 @@ func (cs *ClientSession) ResourceTemplates(ctx context.Context, params *ListReso
15761587
if params == nil {
15771588
params = &ListResourceTemplatesParams{}
15781589
}
1579-
return paginate(ctx, params, cs.ListResourceTemplates, func(res *ListResourceTemplatesResult) []*ResourceTemplate {
1590+
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListResourceTemplates, func(res *ListResourceTemplatesResult) []*ResourceTemplate {
15801591
return res.ResourceTemplates
15811592
})
15821593
}
@@ -1589,16 +1600,38 @@ func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams)
15891600
if params == nil {
15901601
params = &ListPromptsParams{}
15911602
}
1592-
return paginate(ctx, params, cs.ListPrompts, func(res *ListPromptsResult) []*Prompt {
1603+
return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListPrompts, func(res *ListPromptsResult) []*Prompt {
15931604
return res.Prompts
15941605
})
15951606
}
15961607

15971608
// 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] {
1609+
//
1610+
// It fetches pages by calling listFunc until the result has no NextCursor,
1611+
// maxPages is exceeded (if non-zero), or a cursor cycle is detected.
1612+
// The caller's params struct is not mutated; a local copy is used instead.
1613+
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] {
15991614
return func(yield func(*T, error) bool) {
1615+
// Copy the underlying struct so we don't mutate the caller's params.
1616+
// P is always a pointer to a struct (e.g. *ListToolsParams).
1617+
// We use reflect to create a shallow copy of the pointed-to struct.
1618+
localParams := params
1619+
if v := reflect.ValueOf(params); v.Kind() == reflect.Pointer {
1620+
cp := reflect.New(v.Type().Elem())
1621+
cp.Elem().Set(v.Elem())
1622+
localParams = cp.Interface().(P)
1623+
}
1624+
var seen map[string]bool
1625+
pages := 0
1626+
// 0 means use default (64); negative means unlimited.
1627+
effectiveMax := maxPages
1628+
if effectiveMax == 0 {
1629+
effectiveMax = DefaultListMaxPages
1630+
}
1631+
16001632
for {
1601-
res, err := listFunc(ctx, params)
1633+
pages++
1634+
res, err := listFunc(ctx, localParams)
16021635
if err != nil {
16031636
yield(nil, err)
16041637
return
@@ -1612,7 +1645,21 @@ func paginate[P listParams, R listResult[T], T any](ctx context.Context, params
16121645
if nextCursorVal == nil || *nextCursorVal == "" {
16131646
return
16141647
}
1615-
*params.cursorPtr() = *nextCursorVal
1648+
// Check max pages limit.
1649+
if effectiveMax > 0 && pages >= effectiveMax {
1650+
yield(nil, fmt.Errorf("mcp: pagination exceeded maximum page limit of %d", effectiveMax))
1651+
return
1652+
}
1653+
// Detect cursor cycles to prevent infinite loops.
1654+
if seen == nil {
1655+
seen = make(map[string]bool)
1656+
}
1657+
if seen[*nextCursorVal] {
1658+
yield(nil, fmt.Errorf("mcp: pagination detected cursor cycle: %q", *nextCursorVal))
1659+
return
1660+
}
1661+
seen[*nextCursorVal] = true
1662+
*localParams.cursorPtr() = *nextCursorVal
16161663
}
16171664
}
16181665
}

‎mcp/client_test.go‎

Lines changed: 184 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,188 @@ 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("DefaultLimit64", func(t *testing.T) {
233+
// With maxPages=0 (default 64), 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("UnlimitedWithNegative", func(t *testing.T) {
275+
results = generatePaginatedResults(all, 1)
276+
var got []*Item
277+
var iterErr error
278+
// maxPages=-1: unlimited, should get all items.
279+
seq := paginate(ctx, &ListTestParams{}, -1, 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+
t.Run("DefaultCapAt64", func(t *testing.T) {
296+
// maxPages=0 defaults to DefaultListMaxPages (64). Create 100 pages; should stop at 64.
297+
bigAll := make([]*Item, 100)
298+
for i := range bigAll {
299+
bigAll[i] = &Item{Name: fmt.Sprintf("item-%d", i)}
300+
}
301+
results := generatePaginatedResults(bigAll, 1)
302+
listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
303+
if len(results) == 0 {
304+
t.Fatal("listFunc called more times than expected")
305+
}
306+
res := results[0]
307+
results = results[1:]
308+
return res, nil
309+
}
310+
var got []*Item
311+
var iterErr error
312+
seq := paginate(ctx, &ListTestParams{}, 0, listFunc, func(r *ListTestResult) []*Item { return r.Items })
313+
for item, err := range seq {
314+
if err != nil {
315+
iterErr = err
316+
break
317+
}
318+
got = append(got, item)
319+
}
320+
if iterErr == nil {
321+
t.Fatal("expected pagination error at default cap, got nil")
322+
}
323+
if len(got) != DefaultListMaxPages {
324+
t.Fatalf("got %d items, want %d (DefaultListMaxPages)", len(got), DefaultListMaxPages)
325+
}
326+
})
327+
}
328+
329+
func TestPaginateCursorCycle(t *testing.T) {
330+
ctx := context.Background()
331+
callCount := 0
332+
// Server returns a cursor cycle: page1 → cursorA, page2 → cursorB, page3 → cursorA again.
333+
listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
334+
callCount++
335+
switch callCount {
336+
case 1:
337+
return &ListTestResult{Items: []*Item{{Name: "a"}}, NextCursor: "cursorA"}, nil
338+
case 2:
339+
return &ListTestResult{Items: []*Item{{Name: "b"}}, NextCursor: "cursorB"}, nil
340+
case 3:
341+
// Cycle: cursorA was already seen.
342+
return &ListTestResult{Items: []*Item{{Name: "c"}}, NextCursor: "cursorA"}, nil
343+
default:
344+
t.Fatal("listFunc called after cycle should have been detected")
345+
return nil, fmt.Errorf("unreachable")
346+
}
347+
}
348+
349+
var got []*Item
350+
var iterErr error
351+
seq := paginate(ctx, &ListTestParams{}, -1, listFunc, func(r *ListTestResult) []*Item { return r.Items })
352+
for item, err := range seq {
353+
if err != nil {
354+
iterErr = err
355+
break
356+
}
357+
got = append(got, item)
358+
}
359+
if iterErr == nil {
360+
t.Fatal("expected cursor cycle error, got nil")
361+
}
362+
if len(got) != 3 {
363+
t.Fatalf("got %d items, want 3 (items on cycling page are yielded before detection)", len(got))
364+
}
365+
}
366+
367+
func TestPaginateNoMutation(t *testing.T) {
368+
ctx := context.Background()
369+
params := &ListTestParams{Cursor: "initial"}
370+
results := generatePaginatedResults(allItems[:3], 1)
371+
372+
listFunc := func(ctx context.Context, p *ListTestParams) (*ListTestResult, error) {
373+
if len(results) == 0 {
374+
return &ListTestResult{Items: nil, NextCursor: ""}, nil
375+
}
376+
res := results[0]
377+
results = results[1:]
378+
return res, nil
379+
}
380+
381+
var got []*Item
382+
seq := paginate(ctx, params, -1, listFunc, func(r *ListTestResult) []*Item { return r.Items })
383+
for item, err := range seq {
384+
if err != nil {
385+
t.Fatalf("unexpected error: %v", err)
386+
}
387+
got = append(got, item)
388+
}
389+
390+
// The original params should not have been mutated.
391+
if params.Cursor != "initial" {
392+
t.Errorf("params.Cursor was mutated to %q, want %q", params.Cursor, "initial")
393+
}
394+
if len(got) != 3 {
395+
t.Fatalf("got %d items, want 3", len(got))
396+
}
397+
}
398+
217399
func TestClientCapabilities(t *testing.T) {
218400
testCases := []struct {
219401
name string

0 commit comments

Comments
 (0)