Skip to content

[WRONG BRANCH] fix(google): bound tool schema $ref expansion to prevent exponential DoS #182

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.

Sign up for GitHub

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

Closed
luvs01 wants to merge 1 commit into from
+54 −9
Closed

[WRONG BRANCH] fix(google): bound tool schema $ref expansion to prevent exponential DoS #182

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 28 additions & 9 deletions src/adapters/google-tool-schema.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,12 @@ type Schema = Record<string, unknown>;
const ALLOWED_TYPES = new Set(["string", "integer", "number", "boolean", "array", "object"]);
const MAX_SCHEMA_DEPTH = 24; // Google's documented nesting limit is 32; leave headroom for CCA.
const MAX_DEREF_DEPTH = 16;
const MAX_SCHEMA_NODES = 1_024;

interface SanitizeState {
activeRefs: Set<string>;
remainingNodes: number;
}

function isRecord(value: unknown): value is Schema {
return !!value && typeof value === "object" && !Array.isArray(value);
Expand Down Expand Up @@ -63,11 +69,12 @@ function sanitizeEnum(value: unknown): string[] | undefined {
function normalizeAnyOf(
value: unknown,
defs: Map<string, unknown>,
state: SanitizeState,
depth: number,
refDepth: number,
): Schema {
if (!Array.isArray(value) || value.length === 0) return {};
const schemas = value.map(item => sanitizeSchema(item, defs, depth + 1, refDepth, true));
const schemas = value.map(item => sanitizeSchema(item, defs, state, depth + 1, refDepth, true));

const nonNullSchemas = schemas.filter(schema => schema.type !== "null");
const nullSchemas = schemas.filter(schema => schema.type === "null");
Expand Down Expand Up @@ -98,35 +105,46 @@ function normalizeAnyOf(
function sanitizeProperties(
value: unknown,
defs: Map<string, unknown>,
state: SanitizeState,
depth: number,
refDepth: number,
): Record<string, Schema> | undefined {
if (!isRecord(value)) return undefined;
const properties: Record<string, Schema> = Object.create(null) as Record<string, Schema>;
for (const [name, schema] of Object.entries(value)) {
// Property names form a name bag and must never be interpreted as schema keywords.
properties[name] = sanitizeSchema(schema, defs, depth + 1, refDepth, false);
properties[name] = sanitizeSchema(schema, defs, state, depth + 1, refDepth, false);
}
return properties;
}

function sanitizeSchema(
node: unknown,
defs: Map<string, unknown>,
state: SanitizeState,
depth: number,
refDepth: number,
preserveNullType: boolean,
): Schema {
if (depth >= MAX_SCHEMA_DEPTH || !isRecord(node)) return {};
if (depth >= MAX_SCHEMA_DEPTH || !isRecord(node) || state.remainingNodes-- <= 0) return {};

if (typeof node.$ref === "string" && refDepth < MAX_DEREF_DEPTH) {
if (
typeof node.$ref === "string"
&& refDepth < MAX_DEREF_DEPTH
&& !state.activeRefs.has(node.$ref)
) {
const target = resolveRef(node.$ref, defs);
if (isRecord(target)) {
const merged: Schema = { ...target };
for (const [key, value] of Object.entries(node)) {
if (key !== "$ref") merged[key] = value;
}
return sanitizeSchema(merged, defs, depth, refDepth + 1, preserveNullType);
state.activeRefs.add(node.$ref);
try {
return sanitizeSchema(merged, defs, state, depth, refDepth + 1, preserveNullType);
} finally {
state.activeRefs.delete(node.$ref);
}
}
}

Expand All @@ -140,26 +158,27 @@ function sanitizeSchema(
const enumValues = sanitizeEnum(node.enum ?? (typeof node.const === "string" ? [node.const] : undefined));
if (enumValues) out.enum = enumValues;

const properties = sanitizeProperties(node.properties, defs, depth, refDepth);
const properties = sanitizeProperties(node.properties, defs, state, depth, refDepth);
if (properties) out.properties = properties;

if (isRecord(node.items)) {
out.items = sanitizeSchema(node.items, defs, depth + 1, refDepth, false);
out.items = sanitizeSchema(node.items, defs, state, depth + 1, refDepth, false);
}

if (Array.isArray(node.required)) {
out.required = [...new Set(node.required.filter((item): item is string => typeof item === "string"))];
}

if (node.anyOf !== undefined) Object.assign(out, normalizeAnyOf(node.anyOf, defs, depth, refDepth));
if (node.anyOf !== undefined) Object.assign(out, normalizeAnyOf(node.anyOf, defs, state, depth, refDepth));
return out;
}

export function sanitizeGeminiToolParameters(parameters: unknown): Record<string, unknown> {
try {
const defs = new Map<string, unknown>();
collectDefs(parameters, defs);
const root = sanitizeSchema(parameters, defs, 0, 0, false);
const state: SanitizeState = { activeRefs: new Set(), remainingNodes: MAX_SCHEMA_NODES };
const root = sanitizeSchema(parameters, defs, state, 0, 0, false);

// Function arguments are always an object. Claude additionally rejects root composition and a
// missing root type even when those forms are valid general-purpose JSON Schema.
Expand Down
26 changes: 26 additions & 0 deletions tests/google-tool-schema.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,32 @@ describe("sanitizeGeminiToolParameters", () => {
});
});

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: {} },
},
},
});
});
Comment on lines +299 to +323

@coderabbitai coderabbitai Bot Aug 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 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.


test("never leaks the internal null type used while normalizing unions", () => {
const out = sanitizeGeminiToolParameters({
type: "object",
Expand Down
Loading