Block uninterruptible big-int operations in the local Python executor

src/smolagents/local_python_executor.py

@@ -58,9 +58,73 @@ class InterpreterError(ValueError):
 MAX_OPERATIONS = 10000000
 MAX_WHILE_ITERATIONS = 1000000
 MAX_EXECUTION_TIME_SECONDS = 30
+# A single `**`, `<<` or `*` on ints beyond this result size runs as one uninterruptible C call
+# holding the GIL, which the thread-based `timeout` cannot stop and which freezes the whole process.
+MAX_INT_BITS = 1000000
+# Sequence repetition (`str`/`bytes`/`list`/`tuple` times an int) is the same freeze in memory rather
+# than CPU: `seq * n` builds the whole result in one C call. Cap the element count of the result.
+MAX_SEQUENCE_LENGTH = 100_000_000
 ALLOWED_DUNDER_METHODS = ["__init__", "__str__", "__repr__"]
 
 
+def check_safe_operation(op: str, left: Any, right: Any) -> None:
+    """
+    Raise InterpreterError if an operation would produce a result too large to compute safely.
+
+    Some operators run as a single uninterruptible C call that holds the GIL, so the thread-based
+    `timeout` cannot stop them and they freeze the whole process. Big-int `**`/`<<`/`*` blow up CPU;
+    sequence repetition (`str`/`bytes`/`list`/`tuple` times an int) blows up memory. Both are blocked
+    here before the operation starts.
+
+    Args:
+        op (`str`): The operator symbol: `"**"`, `"<<"` or `"*"`.
+        left: Left operand.
+        right: Right operand.
+    """
+    if op == "*":
+        # Sequence repetition, in either operand order: seq * count or count * seq.
+        if isinstance(left, (str, bytes, bytearray, list, tuple)) and isinstance(right, int):
+            seq, count = left, right
+        elif isinstance(right, (str, bytes, bytearray, list, tuple)) and isinstance(left, int):
+            seq, count = right, left
+        else:
+            seq = None
+        if seq is not None:
+            if not isinstance(count, bool) and count > 0 and len(seq) * count > MAX_SEQUENCE_LENGTH:
+                raise InterpreterError(
+                    f"Operation '*' would produce a sequence of around {len(seq) * count} elements, "
+                    f"exceeding the maximum of {MAX_SEQUENCE_LENGTH} allowed. Use a smaller repeat count."
+                )
+            return
+    if not isinstance(left, int) or not isinstance(right, int):
+        return
+    if op == "**":
+        if abs(left) <= 1 or right <= 1:
+            return
+        estimated_bits = left.bit_length() * right
+    elif op == "<<":
+        if left == 0 or right <= 0:
+            return
+        estimated_bits = left.bit_length() + right
+    elif op == "*":
+        estimated_bits = left.bit_length() + right.bit_length()
+    else:
+        return
+    if estimated_bits > MAX_INT_BITS:
+        raise InterpreterError(
+            f"Operation '{op}' would produce an integer of around {estimated_bits} bits, "
+            f"exceeding the maximum of {MAX_INT_BITS} bits allowed. "
+            "Use smaller operands, or pow(base, exp, mod) for modular exponentiation."
+        )
+
+
+def safe_pow(base, exp, mod=None):
+    if mod is None:
+        check_safe_operation("**", base, exp)
+        return pow(base, exp)
+    return pow(base, exp, mod)
+
+
 def custom_print(*args):
     return None
 
@@ -97,7 +161,7 @@ def nodunder_getattr(obj, name, default=None):
     "atan2": math.atan2,
     "degrees": math.degrees,
     "radians": math.radians,
-    "pow": pow,
+    "pow": safe_pow,
     "sqrt": math.sqrt,
     "len": len,
     "sum": sum,
@@ -672,12 +736,14 @@ def get_current_value(target: ast.AST) -> Any:
     elif isinstance(expression.op, ast.Sub):
         current_value -= value_to_add
     elif isinstance(expression.op, ast.Mult):
+        check_safe_operation("*", current_value, value_to_add)
         current_value *= value_to_add
     elif isinstance(expression.op, ast.Div):
         current_value /= value_to_add
     elif isinstance(expression.op, ast.Mod):
         current_value %= value_to_add
     elif isinstance(expression.op, ast.Pow):
+        check_safe_operation("**", current_value, value_to_add)
         current_value **= value_to_add
     elif isinstance(expression.op, ast.FloorDiv):
         current_value //= value_to_add
@@ -688,6 +754,7 @@ def get_current_value(target: ast.AST) -> Any:
     elif isinstance(expression.op, ast.BitXor):
         current_value ^= value_to_add
     elif isinstance(expression.op, ast.LShift):
+        check_safe_operation("<<", current_value, value_to_add)
         current_value <<= value_to_add
     elif isinstance(expression.op, ast.RShift):
         current_value >>= value_to_add
@@ -744,12 +811,14 @@ def evaluate_binop(
     elif isinstance(binop.op, ast.Sub):
         return left_val - right_val
     elif isinstance(binop.op, ast.Mult):
+        check_safe_operation("*", left_val, right_val)
         return left_val * right_val
     elif isinstance(binop.op, ast.Div):
         return left_val / right_val
     elif isinstance(binop.op, ast.Mod):
         return left_val % right_val
     elif isinstance(binop.op, ast.Pow):
+        check_safe_operation("**", left_val, right_val)
         return left_val**right_val
     elif isinstance(binop.op, ast.FloorDiv):
         return left_val // right_val
@@ -760,6 +829,7 @@ def evaluate_binop(
     elif isinstance(binop.op, ast.BitXor):
         return left_val ^ right_val
     elif isinstance(binop.op, ast.LShift):
+        check_safe_operation("<<", left_val, right_val)
         return left_val << right_val
     elif isinstance(binop.op, ast.RShift):
         return left_val >> right_val

tests/test_local_python_executor.py

@@ -421,6 +421,76 @@ def func():
         evaluate_python_code(code, {"range": range}, state=state)
         assert state["_operations_count"]["counter"] == 5
 
+    @pytest.mark.parametrize(
+        "code",
+        [
+            "10 ** 10 ** 8",
+            "1 << 10 ** 9",
+            "(10 ** 300000) * (10 ** 300000) * (10 ** 300000)",
+            "x = 2\nx **= 10 ** 9",
+            "x = 1\nx <<= 10 ** 9",
+            "x = 10 ** 300000\nx *= 10 ** 300000\nx *= 10 ** 300000",
+            "pow(10, 10 ** 8)",
+            "True << 10 ** 9",
+            "class BigInt(int):\n    pass\n\nBigInt(10) ** (10 ** 8)",
+        ],
+    )
+    def test_huge_int_operations_are_blocked(self, code):
+        # Uninterruptible C-level big-int operations would freeze the thread-based timeout: see issue #2473
+        with pytest.raises(InterpreterError, match="exceeding the maximum"):
+            evaluate_python_code(code, BASE_PYTHON_TOOLS, state={})
+
+    @pytest.mark.parametrize(
+        "code, expected",
+        [
+            ("2 ** 100", 2**100),
+            ("10 ** 1000", 10**1000),
+            ("1 << 20", 1 << 20),
+            ("pow(7, 2 ** 64, 97)", pow(7, 2**64, 97)),
+            ("2.5 ** 100", 2.5**100),
+            ("1 ** 10 ** 9", 1),
+            ("0 ** 10 ** 9", 0),
+            ("2 ** -2", 0.25),
+            ("True * 5", 5),
+            ("True ** 10 ** 9", 1),
+        ],
+    )
+    def test_reasonable_int_operations_still_work(self, code, expected):
+        result, _ = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={})
+        assert result == expected
+
+    @pytest.mark.parametrize(
+        "code",
+        [
+            "('a' * 10 ** 6) * 10 ** 4",
+            "[0] * 10 ** 9",
+            "10 ** 4 * ('ab' * 10 ** 6)",
+            "b'x' * 10 ** 9",
+            "(1,) * 10 ** 9",
+            "x = [1]\nx *= 10 ** 9",
+        ],
+    )
+    def test_huge_sequence_repetitions_are_blocked(self, code):
+        # Sequence repetition also builds the result in one uninterruptible C call: see issue #2473 follow-up
+        with pytest.raises(InterpreterError, match="exceeding the maximum"):
+            evaluate_python_code(code, BASE_PYTHON_TOOLS, state={})
+
+    @pytest.mark.parametrize(
+        "code, expected",
+        [
+            ("'ab' * 3", "ababab"),
+            ("3 * 'ab'", "ababab"),
+            ("[1, 2] * 4", [1, 2, 1, 2, 1, 2, 1, 2]),
+            ("(0,) * 5", (0, 0, 0, 0, 0)),
+            ("b'ab' * 2", b"abab"),
+            ("'a' * True", "a"),
+            ("[1] * False", []),
+        ],
+    )
+    def test_reasonable_sequence_repetitions_still_work(self, code, expected):
+        result, _ = evaluate_python_code(code, BASE_PYTHON_TOOLS, state={})
+        assert result == expected
+
     def test_evaluate_string_methods(self):
         code = "'hello'.replace('h', 'o').split('e')"
         result, _ = evaluate_python_code(code, {}, state={})