[WRONG BRANCH] fix(google): bound tool schema $ref expansion to prevent exponential DoS - #182
[WRONG BRANCH] fix(google): bound tool schema $ref expansion to prevent exponential DoS#182luvs01 wants to merge 1 commit into
Conversation
|
This pull request currently targets @luvs01 Please retarget this PR to Its title has been prefixed with This pull request is being kept as a draft automatically. Once every issue above is resolved, it will be marked ready for review again. |
📝 WalkthroughWalkthroughThe schema sanitizer now bounds recursive expansion with active-reference tracking and a global node budget. Recursive ChangesBounded schema sanitization
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/google-tool-schema.test.ts`:
- Around line 299-323: Extend the recursive-reference coverage near “bounds
expansion of branching recursive refs” with an acyclic layered $defs fan-out
that produces more than 1,024 sanitized schema nodes. Sanitize the generated
schema and assert its resulting node count is at most the configured 1,024-node
budget, ensuring remainingNodes is propagated and enforced independently of
activeRefs.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 02cd524f-c25c-49a2-a69f-222bc6b70c7a
📒 Files selected for processing (2)
src/adapters/google-tool-schema.tstests/google-tool-schema.test.ts
| test("bounds expansion of branching recursive refs", () => { | ||
| const out = sanitizeGeminiToolParameters({ | ||
| type: "object", | ||
| properties: { tree: { $ref: "#/$defs/Tree" } }, | ||
| $defs: { | ||
| Tree: { | ||
| type: "object", | ||
| properties: { | ||
| left: { $ref: "#/$defs/Tree" }, | ||
| right: { $ref: "#/$defs/Tree" }, | ||
| }, | ||
| }, | ||
| }, | ||
| }); | ||
|
|
||
| expect(out).toEqual({ | ||
| type: "object", | ||
| properties: { | ||
| tree: { | ||
| type: "object", | ||
| properties: { left: {}, right: {} }, | ||
| }, | ||
| }, | ||
| }); | ||
| }); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Add coverage that exhausts the node budget.
This test only verifies activeRefs. If remainingNodes is removed or is not propagated, this test still passes because the recursive $ref becomes active after one expansion.
Add an acyclic layered $defs fan-out that exceeds 1,024 sanitized nodes. Assert that the output schema-node count remains at or below the budget.
Proposed regression-test shape
+ test("bounds acyclic shared-definition fan-out by node budget", () => {
+ const defs: Record<string, unknown> = {};
+ for (let index = 0; index < 17; index += 1) {
+ const ref = `#/$defs/Node${index + 1}`;
+ defs[`Node${index}`] = {
+ type: "object",
+ properties: { left: { $ref: ref }, right: { $ref: ref } },
+ };
+ }
+ defs.Node17 = { type: "string" };
+
+ const out = sanitizeGeminiToolParameters({
+ type: "object",
+ properties: { tree: { $ref: "`#/`$defs/Node0" } },
+ $defs: defs,
+ });
+
+ // Count schemas through `properties` and `items`, then require <= 1,024.
+ expect(countSanitizedSchemas(out)).toBeLessThanOrEqual(1_024);
+ });🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/google-tool-schema.test.ts` around lines 299 - 323, Extend the
recursive-reference coverage near “bounds expansion of branching recursive refs”
with an acyclic layered $defs fan-out that produces more than 1,024 sanitized
schema nodes. Sanitize the generated schema and assert its resulting node count
is at most the configured 1,024-node budget, ensuring remainingNodes is
propagated and enforced independently of activeRefs.
Motivation
$refgraphs, allowing a crafted tool schema to exhaust CPU/memory during request construction.Description
MAX_SCHEMA_NODESbudget and aSanitizeStatestruct that tracksremainingNodesandactiveRefsto limit total work and detect active$refcycles.SanitizeStatethroughsanitizeSchema,sanitizeProperties,normalizeAnyOf, and arrayitemshandling so every traversal decrements the budget and re-entrancy is prevented.$reftargets only when not already active and ensure an active-ref is removed on return to avoid infinite recursion while still allowing safe inlining of shared defs.$defsare bounded and produce a small safe schema instead of exponential expansion (added totests/google-tool-schema.test.ts).Testing
bun test tests/google-tool-schema.test.ts, which executed 20 tests and all passed.bun run typecheck(bun x tsc --noEmit) and it succeeded.bun run test) in this environment but the full-run was queued/contended with an existing test runner and could not be completed reliably here; focused tests and typecheck passed and CI should run the full suite.Codex Task
Summary by CodeRabbit
Bug Fixes
Tests