Skip to content

Make tool-output validation O(n) instead of O(n^2) (pre-auth CPU DoS)#547

Open
gigioneggiando wants to merge 1 commit into
antirez:mainfrom
gigioneggiando:fix/quadratic-tool-validate
Open

Make tool-output validation O(n) instead of O(n^2) (pre-auth CPU DoS)#547
gigioneggiando wants to merge 1 commit into
antirez:mainfrom
gigioneggiando:fix/quadratic-tool-validate

Conversation

@gigioneggiando

Copy link
Copy Markdown

Summary

responses_validate_tool_outputs() and its Anthropic sibling anthropic_validate_tool_results() resolved each tool output's call_id with responses_find_prior_call_msg() — a backward scan over all preceding messages. Run once per id per tool message over an attacker-supplied, uncapped message array, this is O(n²): ~10⁶ messages ≈ 10¹² comparisons, pinning a CPU core pre-auth.

Fix

Only assistant messages declare call ids (chat_msg_has_call_id matched an assistant's .calls entries). So build a call_id → nearest-preceding assistant message map incrementally with a rax (already used in this file) while scanning forward, and resolve each id with an O(1) lookup.

Registering assistant messages as they are passed and looking up on tool messages reproduces the previous "nearest assistant call before index i" result exactly — including the case where the declaring assistant message appears after the tool message (still not matched, since it isn't in the map yet). The now-unused responses_find_prior_call_msg() and chat_msg_has_call_id() are removed.

Verification

Timed the real responses_validate_tool_outputs() on synthetic message arrays, and confirmed the accept/reject result (ok) is unchanged at every size:

      n     seconds    ratio
   4000     0.0032        -
   8000     0.0052     1.63
  16000     0.0094     1.80
  32000     0.0191     2.03
  64000     0.0399     2.09
  • before: ~4× per doubling of n (O(n²)).
  • after: ~2× per doubling (linear), results identical.

Found during a coordinated security review of ds4; a standalone offline PoC is available on request.

responses_validate_tool_outputs() and anthropic_validate_tool_results()
resolved each tool output's call_id with responses_find_prior_call_msg(),
a backward scan over all preceding messages. Run once per id per tool
message over an attacker-supplied, uncapped message array, that is O(n^2):
~10^6 messages is ~10^12 comparisons, pinning a CPU core pre-auth.

Only assistant messages declare call ids (chat_msg_has_call_id matched
assistant .calls entries), so build a call_id -> nearest-preceding
assistant message map incrementally with a rax while scanning forward, and
resolve each id with an O(1) lookup. Registering assistant messages as they
are passed and looking up on tool messages reproduces the previous
nearest-preceding-before-i result exactly. Drop the now-unused
responses_find_prior_call_msg() and chat_msg_has_call_id().

Verified: the validator's accept/reject results are unchanged, and wall
clock now scales ~linearly with n instead of ~4x per doubling.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant