fix: enforce toolset/promptset boundary on tools/call and prompts/get#3036
fix: enforce toolset/promptset boundary on tools/call and prompts/get#3036sjhddh wants to merge 1 commit intogoogleapis:mainfrom
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces ContainsPrompt and ContainsTool methods to the Promptset and Toolset structs, respectively, along with corresponding unit tests. These methods are used to verify that a requested prompt or tool belongs to the current set before global resolution in the MCP handlers. The review feedback points out that the current linear search implementation ($O(N)$) for these lookups could become a performance bottleneck as the number of items grows, suggesting a map-based approach for
| func (p Promptset) ContainsPrompt(name string) bool { | ||
| for _, n := range p.PromptNames { | ||
| if n == name { | ||
| return true | ||
| } | ||
| } | ||
| return false | ||
| } |
There was a problem hiding this comment.
The ContainsPrompt method uses a linear search ($O(N)$), which is executed on every prompts/get request. While likely acceptable for small promptsets, this could impact performance as the number of prompts grows. Consider using a map for Initialize method.
| func (t Toolset) ContainsTool(name string) bool { | ||
| for _, n := range t.ToolNames { | ||
| if n == name { | ||
| return true | ||
| } | ||
| } | ||
| return false | ||
| } |
There was a problem hiding this comment.
The ContainsTool method performs a linear search ($O(N)$) on every tools/call. For toolsets with many tools, this may become a performance bottleneck. Using a map for Initialize) would be more efficient and scalable.
duwenxin99
left a comment
duwenxin99
left a comment
There was a problem hiding this comment.
Hi @sjhddh , the change LGTM overall except for the error messages. Please let me know once you fix them and I'll handle the merge process for you. Thank you for contributing!
| // Verify tool belongs to the current toolset before resolving globally. | ||
| if !toolset.ContainsTool(toolName) { | ||
| err = fmt.Errorf("tool %q is not part of the current toolset", toolName) | ||
| return jsonrpc.NewError(id, jsonrpc.INVALID_PARAMS, err.Error(), nil), err |
There was a problem hiding this comment.
We should return the same error message as non-existent tool call (invalid tool name: tool with name %q does not exist) . Otherwise if the tool doesn't exist at all, we will still return "tool %q is not part of the current toolset" which could be misleading.
| ) | ||
|
|
||
| // Verify tool belongs to the current toolset before resolving globally. | ||
| if !toolset.ContainsTool(toolName) { |
There was a problem hiding this comment.
Same for the other errors.
| span.SetAttributes(attribute.String("gen_ai.prompt.name", promptName)) | ||
|
|
||
| // Verify prompt belongs to the current promptset before resolving globally. | ||
| if !promptset.ContainsPrompt(promptName) { |
There was a problem hiding this comment.
Similar for prompt-let's treat it the same way as a non-existent prompt.
|
You may want to squash the commits to get rid of the agent account that's missing CLA. |
tools/call and prompts/get resolved tools and prompts via the global resourceMgr without verifying membership in the current toolset or promptset. This allowed clients connected to a scoped toolset to invoke any tool by name, bypassing toolset access boundaries (IDOR). Add ContainsTool/ContainsPrompt membership checks and enforce them in all four MCP protocol versions before delegating to resourceMgr. Return the same "does not exist" error on boundary violation as on a truly missing tool/prompt so clients cannot distinguish the two cases. Fixes googleapis#2755
sjhddh
force-pushed
the
fix/toolset-boundary-enforcement
branch
from
97edc7f to
edb3e0c
Compare
|
@duwenxin99 Thanks! Fixed both:
Let me know if anything else needs polishing. |
2 participants
Summary
tools/callpreviously resolved tools viaresourceMgr.GetTool()without verifying the requested tool belongs to the current toolset, allowing clients to invoke any tool by name regardless of their toolset scope.prompts/getwithresourceMgr.GetPrompt().ContainsTool()/ContainsPrompt()membership checks and enforce them in all four MCP protocol versions (v20241105, v20250326, v20250618, v20251125) before delegating to the global resource manager.Fixes #2755
Test plan
Toolset.ContainsTool()andPromptset.ContainsPrompt()(both positive and negative cases, including empty sets)go test ./internal/tools/... ./internal/prompts/... ./internal/server/mcp/...)tools/callrejects tools from the other toolset🤖 Generated with Claude Code