From 907fb9e8509ef2e8aed955d0867edf549445f1e4 Mon Sep 17 00:00:00 2001 From: lzhan011 <35493221+lzhan011@users.noreply.github.com> Date: Mon, 5 Oct 2026 05:01:57 -0500 Subject: [PATCH] fix(guardrails): apply the `for` guards to `async for` in the sandbox `visit_AsyncFor` routed to `node_contents_visit`, which only recurses into children, so `async for` compiled with none of the guards `for` gets: `visit_For` is what wraps the loop iterator, and the async spelling therefore enforced less than the sync one. Delegating to `visit_For` is not quite enough on its own. It inserts synchronous wrappers -- `_getiter_(y)`, or `_iter_unpack_sequence_(y, spec, _getiter_)` for a tuple target -- and both are plain generators that `async for` cannot consume. `_use_async_iter_guards` therefore swaps just the wrapper name for an async-capable twin: async for a in y -> async for a in _agetiter_(y) async for a, b in y -> async for (a, b) in _aiter_unpack_sequence_(y, spec, _getiter_) `_agetiter_` is `budgeted_iter` over an async iterator, so an `async for` that never yields is still stopped at the execution deadline. `_aiter_unpack_sequence_` budgets the iteration and then applies `guarded_unpack_sequence` per element; unpacking an element is synchronous, so the `_getiter_` handed in is the synchronous one and is passed through untouched. Adapted while rebuilding on `main`, worth flagging: this PR was written when `_getiter_` was `default_guarded_getiter`. `main` now binds it to `budgeted_iter`, which adds the per-item budget check -- so the async twins have to carry that check too, which is why `_agetiter_` exists rather than the async helper simply calling `_getiter_`. (This PR was opened against the since-deleted `litellm_internal_staging`; `main` still routes `visit_AsyncFor` to `node_contents_visit`.) Tests: the compiled bytecode of each `async for` shape must carry an async-capable wrapper, with the sync forms' guards asserted as preconditions; both shapes must still run under a real async generator; and a non-sequence element must still raise `TypeError` through the unpack guard rather than pass silently. The guard-parity test fails with `visit_AsyncFor` put back on `node_contents_visit` (`1 failed | 3 passed`); all 67 in the file pass here. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_016SxQ2jq4B4KLGwPoLgfHye --- .../guardrail_hooks/custom_code/sandbox.py | 70 +++++++++++++++++- .../guardrails/test_custom_code_security.py | 74 +++++++++++++++++++ 2 files changed, 140 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py index 5681438d8d9..5bc1b9df23a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py @@ -16,7 +16,7 @@ restriction intact. import ast import operator -from collections.abc import Callable, Mapping +from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping from types import CodeType from typing import Final @@ -50,8 +50,12 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): inherited automatically. ``AsyncWith`` gets the same treatment for the same reason: ``node_contents_visit`` only recurses into children, so routing it there left ``async with x as (a, b)`` without the unpack guard that - ``with x as (a, b)`` gets. ``AsyncFor``/``Await`` delegate to - ``node_contents_visit`` so their children still get transformed. + ``with x as (a, b)`` gets. ``AsyncFor`` likewise delegates to ``visit_For``, + which is what wraps the loop iterator in ``_getiter_``; the wrapper names it + inserts are then swapped for async-capable ones, since ``budgeted_iter`` and + ``_iter_unpack_sequence_`` are plain generators that ``async for`` cannot + consume. ``Await`` has no synchronous counterpart and only wraps an + expression, so it stays on ``node_contents_visit``. ``visit_While`` rewrites ``while test:`` to ``while _budget_ok_() and test:`` so a loop that never yields is still stopped at the execution deadline; @@ -78,7 +82,9 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): return bounded def visit_AsyncFor(self, node: ast.AsyncFor) -> ast.AST: - return self.node_contents_visit(node) + transformed: Final = self.visit_For(node) + _use_async_iter_guards(transformed) + return transformed def visit_AsyncWith(self, node: ast.AsyncWith) -> ast.AST: return self.visit_With(node) @@ -87,6 +93,60 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): return self.node_contents_visit(node) +_GETITER_NAME: Final = "_getiter_" +_ASYNC_GETITER_NAME: Final = "_agetiter_" +_ITER_UNPACK_NAME: Final = "_iter_unpack_sequence_" +_ASYNC_ITER_UNPACK_NAME: Final = "_aiter_unpack_sequence_" + + +def _use_async_iter_guards(node: ast.AST) -> None: + """Swap the iterator wrapper ``visit_For`` inserted for its async twin. + + ``visit_For`` rewrites ``for x in y`` into ``for x in _getiter_(y)`` and + ``for a, b in y`` into + ``for (a, b) in _iter_unpack_sequence_(y, spec, _getiter_)``. Both wrappers + are synchronous generators, which ``async for`` cannot consume, so the + async form needs the async-generator equivalents under their own names. + Only the wrapper name changes: the guarding each one performs is the same. + """ + iter_node: Final = getattr(node, "iter", None) + if not (isinstance(iter_node, ast.Call) and isinstance(iter_node.func, ast.Name)): + return + if iter_node.func.id == _ITER_UNPACK_NAME: + iter_node.func.id = _ASYNC_ITER_UNPACK_NAME + elif iter_node.func.id == _GETITER_NAME: + iter_node.func.id = _ASYNC_GETITER_NAME + + +async def _agetiter_(iterable: AsyncIterable[object]) -> AsyncIterator[object]: + """``budgeted_iter`` for ``async for``: the same per-item budget check. + + Bound to the name ``_use_async_iter_guards`` substitutes, so an ``async for`` + that never yields is still stopped at the execution deadline, exactly as a + ``for`` loop is. + """ + async for item in iterable: + budget_ok() + yield item + + +async def _aiter_unpack_sequence_( + it: AsyncIterable[object], + spec: object, + getiter: Callable[[object], object], +) -> AsyncIterator[object]: + """``guarded_iter_unpack_sequence`` for ``async for`` targets. + + Same contract as the RestrictedPython helper -- budget the iteration, then + guard each element's sequence unpacking -- over an async iterator. Unpacking + an element is synchronous, so the ``_getiter_`` handed in is the synchronous + one and is passed through untouched. + """ + async for ob in it: + budget_ok() + yield guarded_unpack_sequence(ob, spec, getiter) + + _INPLACE_OPS: Final[Mapping[str, Callable[[object, object], object]]] = { "+=": operator.iadd, "-=": operator.isub, @@ -147,6 +207,8 @@ def build_sandbox_globals() -> dict[str, object]: # ships no default. Without it those statements compile and then raise # NameError the first time the guardrail runs. "_unpack_sequence_": guarded_unpack_sequence, + _ASYNC_GETITER_NAME: _agetiter_, + _ASYNC_ITER_UNPACK_NAME: _aiter_unpack_sequence_, "_write_": full_write_guard, "_inplacevar_": _inplacevar_, "_budget_ok_": budget_ok, diff --git a/tests/unit/proxy/guardrails/test_custom_code_security.py b/tests/unit/proxy/guardrails/test_custom_code_security.py index 03db0098030..7ad2f96363f 100644 --- a/tests/unit/proxy/guardrails/test_custom_code_security.py +++ b/tests/unit/proxy/guardrails/test_custom_code_security.py @@ -929,3 +929,77 @@ async def test_async_with_still_executes(): sandbox_globals, ) assert await sandbox_globals["f"](_ACtx((5, 6))) == 11 + + +def test_async_for_gets_the_same_guards_as_for(): + """`async for` is the async spelling of `for`; it must not enforce less.""" + sync_plain = _guard_names("def f(x):\n for a in x:\n pass\n") + async_plain = _guard_names("async def f(x):\n async for a in x:\n pass\n") + sync_unpack = _guard_names("def f(x):\n for a, b in x:\n pass\n") + async_unpack = _guard_names("async def f(x):\n async for a, b in x:\n pass\n") + + # Preconditions: the sync forms are guarded at all. + assert "_getiter_" in sync_plain + assert "_iter_unpack_sequence_" in sync_unpack + + # The async forms are guarded by the async-capable twins, under their own + # names — `budgeted_iter` and `_iter_unpack_sequence_` are plain generators + # that `async for` cannot consume. + assert "_agetiter_" in async_plain + assert "_aiter_unpack_sequence_" in async_unpack + + # And neither async form is left with no wrapper at all. + assert async_plain & {"_agetiter_", "_aiter_unpack_sequence_"} + assert async_unpack & {"_agetiter_", "_aiter_unpack_sequence_"} + + +async def _arange(*values): + for v in values: + yield v + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("source", "arg", "expected"), + [ + ( + "async def f(it):\n total = 0\n async for a in it:\n total += a\n return total\n", + (1, 2, 3), + 6, + ), + ( + "async def f(it):\n total = 0\n async for a, b in it:\n total += a * b\n return total\n", + ((1, 2), (3, 4)), + 14, + ), + ], +) +async def test_async_for_still_executes(source, arg, expected): + """Guarding `async for` must not break it, in either target shape.""" + from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( + build_sandbox_globals, + compile_sandboxed, + ) + + sandbox_globals = build_sandbox_globals() + exec(compile_sandboxed(source), sandbox_globals) # noqa: S102 + assert await sandbox_globals["f"](_arange(*arg)) == expected + + +@pytest.mark.asyncio +async def test_async_for_unpacking_still_rejects_a_non_sequence(): + """The unpack guard is what makes the element shape an error, not a crash.""" + from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( + build_sandbox_globals, + compile_sandboxed, + ) + + sandbox_globals = build_sandbox_globals() + exec( # noqa: S102 + compile_sandboxed( + "async def f(it):\n async for a, b in it:\n pass\n" + ), + sandbox_globals, + ) + with pytest.raises(TypeError): + await sandbox_globals["f"](_arange(1, 2))