Merge 907fb9e850 into ade17902a0

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
 
@@ -31,6 +31,7 @@ from RestrictedPython.Eval import default_guarded_getitem
 from RestrictedPython.Guards import (
     full_write_guard,
     guarded_iter_unpack_sequence,
+    guarded_unpack_sequence,
     safer_getattr,
 )
 from RestrictedPython.transformer import copy_locations
@@ -46,8 +47,15 @@ class AsyncAwareTransformer(RestrictingNodeTransformer):
     has the same ``_fields`` as ``FunctionDef`` and the same security
     semantics, so we delegate to ``visit_FunctionDef`` — name check, argument
     check, print-scope wrapping, and any future additions to that method are
-    inherited automatically. ``AsyncFor``/``AsyncWith``/``Await`` delegate to
-    ``node_contents_visit`` so their children still get transformed.
+    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`` 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;
@@ -74,15 +82,71 @@ 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.node_contents_visit(node)
+        return self.visit_With(node)
 
     def visit_Await(self, node: ast.Await) -> ast.AST:
         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,
@@ -137,6 +201,14 @@ def build_sandbox_globals() -> dict[str, object]:
         "_getitem_": default_guarded_getitem,
         "_getiter_": budgeted_iter,
         "_iter_unpack_sequence_": guarded_iter_unpack_sequence,
+        # RestrictedPython emits _unpack_sequence_ for every tuple-unpacking
+        # target that is not a for-loop target -- ``a, b = pair``,
+        # ``a, *rest = seq``, ``with x as (a, b)`` -- and, like _inplacevar_,
+        # 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,

tests/unit/proxy/guardrails/test_custom_code_security.py

@@ -846,3 +846,160 @@ def test_config_without_timeout_uses_the_default():
         _initialize_from_config("custom-code-default-timeout", {}).execution_timeout
         == DEFAULT_EXECUTION_TIMEOUT_SECONDS
     )
+
+
+@pytest.mark.parametrize(
+    ("body", "expected"),
+    [
+        # Every one of these compiles fine and then raises
+        # "NameError: name '_unpack_sequence_' is not defined" at call time when
+        # the guard is missing from the sandbox globals.
+        ("    a, b = inputs['pair']\n    return a + b\n", 3),
+        ("    a, (b, c) = inputs['nested']\n    return a + b + c\n", 6),
+        ("    a, *rest = inputs['seq']\n    return rest\n", [2, 3]),
+        ("    with inputs['ctx'] as (a, b):\n        return a + b\n", 15),
+    ],
+)
+def test_tuple_unpacking_runs(body: str, expected):
+    """Ordinary tuple unpacking must work inside a guardrail, not NameError."""
+
+    class _Ctx:
+        def __enter__(self):
+            return (7, 8)
+
+        def __exit__(self, *args):
+            return False
+
+    guardrail = _compile("def apply_guardrail(inputs, request_data, input_type):\n" + body)
+    fn = guardrail._compiled_function
+    assert fn is not None
+    inputs = {"pair": (1, 2), "nested": (1, (2, 3)), "seq": [1, 2, 3], "ctx": _Ctx()}
+    assert fn(inputs, {}, "request") == expected
+
+
+def _guard_names(source: str) -> set[str]:
+    """Names of the RestrictedPython guards the compiled bytecode calls."""
+    import types
+
+    from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import (
+        compile_sandboxed,
+    )
+
+    found: set[str] = set()
+    stack = [compile_sandboxed(source)]
+    while stack:
+        code = stack.pop()
+        found |= {n for n in code.co_names + code.co_varnames if n.startswith("_") and n.endswith("_")}
+        stack += [c for c in code.co_consts if isinstance(c, types.CodeType)]
+    return found
+
+
+def test_async_with_gets_the_same_guards_as_with():
+    """`async with` is the async spelling of `with`; it must not enforce less."""
+    sync_guards = _guard_names("def f(x):\n    with x as (a, b):\n        pass\n")
+    async_guards = _guard_names("async def f(x):\n    async with x as (a, b):\n        pass\n")
+
+    assert "_unpack_sequence_" in sync_guards, "precondition: `with` unpacking is guarded"
+    assert sync_guards <= async_guards
+
+
+@pytest.mark.asyncio
+async def test_async_with_still_executes():
+    """Guarding `async with` must not break it."""
+    from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import (
+        build_sandbox_globals,
+        compile_sandboxed,
+    )
+
+    class _ACtx:
+        def __init__(self, value):
+            self.value = value
+
+        async def __aenter__(self):
+            return self.value
+
+        async def __aexit__(self, *args):
+            return False
+
+    sandbox_globals = build_sandbox_globals()
+    exec(  # noqa: S102
+        compile_sandboxed(
+            "async def f(ctx):\n    async with ctx as (a, b):\n        return a + b\n"
+        ),
+        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))