fix(guardrails): wire _unpack_sequence_ into the custom-code sandbox

RestrictedPython rewrites every tuple-unpacking target that is not a for-loop
target into a call to _unpack_sequence_, and ships no default for it -- the
same way it ships no _inplacevar_, which this module already supplies. The
sandbox globals never defined it, so ordinary guardrail code compiles cleanly
and then raises NameError on its first run:

    a, b = pair
    a, (b, c) = nested
    a, *rest = seq
    with ctx as (a, b):

RestrictedPython.Guards.guarded_unpack_sequence is the helper the rewritten
bytecode expects, and it applies the same _getiter_ guard to each element that
guarded_iter_unpack_sequence already applies for `for a, b in x`.

Rebuilt on `main`: this PR was opened against `litellm_internal_staging`,
which has since been deleted, and `main` diverged from it rather than
containing it. The change is re-applied on top of current `main`, where the
guard is still absent from build_sandbox_globals(); the test moved upstream to
tests/unit/proxy/guardrails/ and the new cases are appended there.

tests/unit/proxy/guardrails/test_custom_code_security.py -> 61 passed. With
the one globals entry removed, the four new cases fail with exactly
`NameError: name '_unpack_sequence_' is not defined`.
This commit is contained in:
lzhan011 2026-10-05 04:57:51 -05:00
parent 1238cfe90f
commit 433e78070d
2 changed files with 36 additions and 0 deletions

View file

@ -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
@ -137,6 +138,12 @@ 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,
"_write_": full_write_guard,
"_inplacevar_": _inplacevar_,
"_budget_ok_": budget_ok,

View file

@ -846,3 +846,32 @@ 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