-
Notifications
You must be signed in to change notification settings - Fork 490
mcp: harden paginate against infinite loops and param mutation #1110
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -204,8 +204,19 @@ type ClientOptions struct { | |
| // reset" guidance, letting a transient miss pass without tearing down an | ||
| // otherwise live session. Has no effect unless KeepAlive is non-zero. | ||
| KeepAliveFailureThreshold int | ||
| // ListMaxPages is the maximum number of pages to fetch during automatic | ||
| // pagination of list operations (Tools, Resources, ResourceTemplates, | ||
| // Prompts). A value of 0 uses the default of [DefaultListMaxPages] (64). | ||
| // A negative value means unlimited. A positive value caps the number of | ||
| // pages. This prevents runaway pagination loops caused by server-side | ||
| // cursor cycles. | ||
| ListMaxPages int | ||
| } | ||
|
|
||
| // DefaultListMaxPages is the default value for [ClientOptions.ListMaxPages], | ||
| // matching the TypeScript SDK's DEFAULT_LIST_MAX_PAGES. | ||
| const DefaultListMaxPages = 64 | ||
|
|
||
| // toolContextKeyType is the context key type for passing tool definitions | ||
| // from CallTool to the transport layer. | ||
| type toolContextKeyType struct{} | ||
|
|
@@ -1550,7 +1561,7 @@ func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) ite | |
| if params == nil { | ||
| params = &ListToolsParams{} | ||
| } | ||
| return paginate(ctx, params, cs.ListTools, func(res *ListToolsResult) []*Tool { | ||
| return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListTools, func(res *ListToolsResult) []*Tool { | ||
| return res.Tools | ||
| }) | ||
| } | ||
|
|
@@ -1563,7 +1574,7 @@ func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesPar | |
| if params == nil { | ||
| params = &ListResourcesParams{} | ||
| } | ||
| return paginate(ctx, params, cs.ListResources, func(res *ListResourcesResult) []*Resource { | ||
| return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListResources, func(res *ListResourcesResult) []*Resource { | ||
| return res.Resources | ||
| }) | ||
| } | ||
|
|
@@ -1576,7 +1587,7 @@ func (cs *ClientSession) ResourceTemplates(ctx context.Context, params *ListReso | |
| if params == nil { | ||
| params = &ListResourceTemplatesParams{} | ||
| } | ||
| return paginate(ctx, params, cs.ListResourceTemplates, func(res *ListResourceTemplatesResult) []*ResourceTemplate { | ||
| return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListResourceTemplates, func(res *ListResourceTemplatesResult) []*ResourceTemplate { | ||
| return res.ResourceTemplates | ||
| }) | ||
| } | ||
|
|
@@ -1589,16 +1600,38 @@ func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams) | |
| if params == nil { | ||
| params = &ListPromptsParams{} | ||
| } | ||
| return paginate(ctx, params, cs.ListPrompts, func(res *ListPromptsResult) []*Prompt { | ||
| return paginate(ctx, params, cs.client.opts.ListMaxPages, cs.ListPrompts, func(res *ListPromptsResult) []*Prompt { | ||
| return res.Prompts | ||
| }) | ||
| } | ||
|
|
||
| // paginate is a generic helper function to provide a paginated iterator. | ||
| 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] { | ||
| // | ||
| // It fetches pages by calling listFunc until the result has no NextCursor, | ||
| // maxPages is exceeded (if non-zero), or a cursor cycle is detected. | ||
| // The caller's params struct is not mutated; a local copy is used instead. | ||
| 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] { | ||
| return func(yield func(*T, error) bool) { | ||
| // Copy the underlying struct so we don't mutate the caller's params. | ||
| // P is always a pointer to a struct (e.g. *ListToolsParams). | ||
| // We use reflect to create a shallow copy of the pointed-to struct. | ||
| localParams := params | ||
| if v := reflect.ValueOf(params); v.Kind() == reflect.Pointer { | ||
| cp := reflect.New(v.Type().Elem()) | ||
| cp.Elem().Set(v.Elem()) | ||
| localParams = cp.Interface().(P) | ||
| } | ||
|
Comment on lines
+1618
to
+1623
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. it could be simplified to sth like
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks for the suggestion! Unfortunately this doesn't work because P is a type parameter constrained by the listParams interface. Go doesn't allow comparing a type parameter with nil or dereferencing it directly. All concrete uses of P are pointer types, but the generic constraint is an interface, so we need reflect to create the shallow copy. |
||
| var seen map[string]bool | ||
| pages := 0 | ||
| // 0 means use default (64); negative means unlimited. | ||
| effectiveMax := maxPages | ||
| if effectiveMax == 0 { | ||
| effectiveMax = DefaultListMaxPages | ||
| } | ||
|
|
||
| for { | ||
| res, err := listFunc(ctx, params) | ||
| pages++ | ||
| res, err := listFunc(ctx, localParams) | ||
| if err != nil { | ||
| yield(nil, err) | ||
| return | ||
|
|
@@ -1612,7 +1645,21 @@ func paginate[P listParams, R listResult[T], T any](ctx context.Context, params | |
| if nextCursorVal == nil || *nextCursorVal == "" { | ||
| return | ||
| } | ||
| *params.cursorPtr() = *nextCursorVal | ||
| // Check max pages limit. | ||
| if effectiveMax > 0 && pages >= effectiveMax { | ||
| yield(nil, fmt.Errorf("mcp: pagination exceeded maximum page limit of %d", effectiveMax)) | ||
| return | ||
| } | ||
| // Detect cursor cycles to prevent infinite loops. | ||
| if seen == nil { | ||
| seen = make(map[string]bool) | ||
| } | ||
| if seen[*nextCursorVal] { | ||
| yield(nil, fmt.Errorf("mcp: pagination detected cursor cycle: %q", *nextCursorVal)) | ||
| return | ||
| } | ||
| seen[*nextCursorVal] = true | ||
| *localParams.cursorPtr() = *nextCursorVal | ||
| } | ||
| } | ||
| } | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this is a behavioral change for existing users
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @guglielmo-san
You're right, this is a behavioral change. However, it's a security hardening fix — without a default limit, a misbehaving server can cause unbounded pagination loops (infinite CPU/memory consumption), which is a denial-of-service vector.
The default of 64 aligns with both the TypeScript SDK (
DEFAULT_LIST_MAX_PAGES = 64) and the Python SDK. Users who legitimately need more pages can setListMaxPagesto a higher value or-1for unlimited.I've added a note to the PR description explaining this behavioral change.