diff --git a/src/smolagents/local_python_executor.py b/src/smolagents/local_python_executor.py index 8e150cd5c..c61d53e58 100644 --- a/src/smolagents/local_python_executor.py +++ b/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 diff --git a/tests/test_local_python_executor.py b/tests/test_local_python_executor.py index aede6cd2d..ed759b2c8 100644 --- a/tests/test_local_python_executor.py +++ b/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={})