diff --git a/optillm/z3_solver.py b/optillm/z3_solver.py index 0f0b940..4351b16 100644 --- a/optillm/z3_solver.py +++ b/optillm/z3_solver.py @@ -15,6 +15,66 @@ class TimeoutException(Exception): pass +# Modules the solver contract allows: the system prompt promises the generated +# code only uses Z3/SymPy (plus basic stdlib math helpers), and these are +# exactly the modules execute_code_in_process() makes available. +_ALLOWED_SOLVER_IMPORTS = {'z3', 'sympy', 'math', 'itertools', 'fractions'} + +# Calls that provide arbitrary execution or file/IO access, or that reach the +# same denied attributes through string arguments (getattr/setattr/delattr); +# the declared solver subset has no use for any of them. +_DENIED_SOLVER_CALLS = { + 'eval', 'exec', 'compile', '__import__', 'open', 'input', 'breakpoint', + 'globals', 'locals', 'vars', 'getattr', 'setattr', 'delattr', +} + +# Dunder attributes used in classic module-object traversal escapes (e.g. +# z3.__loader__, ().__class__.__bases__[0].__subclasses__()). __dict__ and +# __getattribute__ are denied because they turn every other denied attribute +# into a string lookup (z3.__dict__['__loader__'], +# z3.__getattribute__('__spec__')), and a module's __dict__ exposes its real +# __builtins__. A fixed denylist instead of banning every dunder so legitimate +# solver code that touches z3/sympy internals by name is not rejected. +_DENIED_SOLVER_ATTRS = { + '__class__', '__bases__', '__subclasses__', '__mro__', '__globals__', + '__builtins__', '__init__', '__code__', '__closure__', '__loader__', + '__spec__', '__import__', '__reduce__', '__dict__', '__getattribute__', + '__getattr__', '__setattr__', '__delattr__', +} + +def _validate_solver_ast(tree: ast.AST) -> None: + """Statically reject solver code that leaves the declared solver subset. + + The system prompt in analyze_query() promises the generated code will not + import anything beyond Z3/SymPy/math; this makes that documented contract + enforceable before the code reaches exec(). This is defense-in-depth, NOT + a sandbox. + + Raises ValueError listing every violation found. + """ + violations = [] + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + root = alias.name.split('.')[0] + if root not in _ALLOWED_SOLVER_IMPORTS: + violations.append( + f"import of '{alias.name}' (allowed modules: {sorted(_ALLOWED_SOLVER_IMPORTS)})") + elif isinstance(node, ast.ImportFrom): + root = (node.module or '').split('.')[0] + if root not in _ALLOWED_SOLVER_IMPORTS: + violations.append( + f"import from '{node.module}' (allowed modules: {sorted(_ALLOWED_SOLVER_IMPORTS)})") + elif isinstance(node, ast.Call) and isinstance(node.func, ast.Name) \ + and node.func.id in _DENIED_SOLVER_CALLS: + violations.append(f"call to '{node.func.id}()'") + elif isinstance(node, ast.Attribute) and node.attr in _DENIED_SOLVER_ATTRS: + violations.append(f"access to attribute '{node.attr}'") + + if violations: + raise ValueError( + "solver code rejected by static safety check: " + "; ".join(violations)) + def prepare_execution_globals(): """ Prepare globals dictionary for Z3/SymPy code execution. @@ -22,7 +82,10 @@ def prepare_execution_globals(): WARNING: This is NOT a security sandbox. The name "execution_globals" reflects that this simply provides the execution environment for solver code, not a security boundary. The code is executed via exec() with access to z3, sympy, - and math libraries. Only execute trusted code. + and math libraries. Only execute trusted code. A static AST gate + (_validate_solver_ast) additionally rejects code that leaves the declared + solver subset before it reaches this environment; this is defense-in-depth, + not a sandbox. """ execution_globals = { 'print': print, @@ -273,7 +336,21 @@ def standard_llm_inference(self, query: str) -> str: def extract_and_validate_expressions(self, analysis: str) -> str: formulation = re.search(r"```python\n([\s\S]+?)```", analysis) if formulation: - return formulation.group(1).strip() + code = formulation.group(1).strip() + # Static gate: enforce the "z3/sympy only" contract promised in the + # analysis prompt before the code can reach exec(). Raise (instead of + # returning an "Error:" string) so rejected code never enters the + # fix-up loop in solve_with_z3_sympy(); process_query catches this + # and degrades to standard inference. Syntax errors are left to the + # existing execution-time handling, which feeds them back for a + # corrected formulation. + try: + tree = ast.parse(code) + except SyntaxError: + tree = None + if tree is not None: + _validate_solver_ast(tree) + return code raise ValueError("No valid Z3 or SymPy formulation found in the analysis.") def solve_with_z3_sympy(self, formulation: str, max_attempts: int = 3) -> Dict[str, Any]: diff --git a/tests/test_z3_solver_guard.py b/tests/test_z3_solver_guard.py new file mode 100644 index 0000000..faeb38c --- /dev/null +++ b/tests/test_z3_solver_guard.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +""" +Tests for the static AST gate on z3_solver (issue #303). + +The z3 approach exec()s LLM-generated code; the system prompt promises the +code only uses Z3/SymPy/math. These tests verify the AST gate enforces that +contract: dangerous imports/calls/dunder-attribute traversals are rejected +before exec(), and rejected code degrades to standard inference instead of +entering the fix-up loop. + +All tests are offline except the allow-path cases that run the real solver in +the local venv (z3 is a core dependency); no network or LLM is contacted. +""" + +import ast +import sys +import os +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from optillm.z3_solver import ( + Z3SymPySolverSystem, + _validate_solver_ast, + execute_code_in_process, +) + + +def _guard(code: str): + """Run the gate the way extract_and_validate_expressions() does.""" + _validate_solver_ast(ast.parse(code)) + + +def _response(content, completion_tokens=10): + resp = MagicMock() + resp.choices = [MagicMock(message=MagicMock(content=content))] + resp.usage.completion_tokens = completion_tokens + return resp + + +# --------------------------------------------------------------------------- +# Refusal cases: code that leaves the declared solver subset +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("code,violation", [ + ("import os", "import of 'os'"), + ("import subprocess\nsubprocess.run(['id'])", "import of 'subprocess'"), + ("from subprocess import run", "import from 'subprocess'"), + ("import pathlib", "import of 'pathlib'"), + ("__import__('os').system('id')", "__import__"), + ("open('/etc/passwd')", "open()"), + ("eval('1+1')", "eval()"), + ("exec('x = 1')", "exec()"), + ("compile('1', '', 'eval')", "compile()"), + ("name = globals()", "globals()"), + ("().__class__.__bases__[0].__subclasses__()", "__class__"), + ("z3.__loader__", "__loader__"), + ("sympy.__spec__", "__spec__"), + ("f = z3.Bool.__init__", "__init__"), + ("import z3\nz3.__builtins__", "__builtins__"), + # String-based evasions of the attribute denylist (QA finding): getattr + # reaches the same objects via a constant argument, and __dict__ / + # __getattribute__ turn every dunder into a string lookup. A module's + # __dict__ also exposes its real __builtins__. + ("getattr(z3, '__loader__')", "getattr()"), + ("setattr(z3, '__loader__', None)", "setattr()"), + ("delattr(z3, 'Int')", "delattr()"), + ("z3.__dict__['__builtins__']['__import__']('os')", "__dict__"), + ("z3.__getattribute__('__spec__')", "__getattribute__"), + ("sympy.__getattr__('_load')", "__getattr__"), +]) +def test_guard_rejects(code, violation): + with pytest.raises(ValueError) as excinfo: + _guard(code) + assert "static safety check" in str(excinfo.value) + assert violation in str(excinfo.value) + + +# --------------------------------------------------------------------------- +# Allow cases: legitimate solver code passes the gate and really executes +# --------------------------------------------------------------------------- + +SOLVER_OK_CASES = [ + "x = Int('x')\ns = Solver()\ns.add(x > 2)\nprint(s.check())", + "from sympy import Symbol, solve\nx = Symbol('x')\nprint(solve(x**2 - 4, x))", + "from fractions import Fraction\nprint(Fraction(1, 3))", + "import math\nprint(math.sqrt(16))", + "from itertools import permutations\nprint(list(permutations([1, 2])))", +] + + +@pytest.mark.parametrize("code", SOLVER_OK_CASES) +def test_guard_allows_solver_code(code): + _guard(code) # must not raise + + +@pytest.mark.parametrize("code", SOLVER_OK_CASES) +def test_allowed_solver_code_executes(code): + status, output = execute_code_in_process(code) + assert status == "success", output + assert output.strip(), "expected some printed output" + + +# --------------------------------------------------------------------------- +# Wiring: extract_and_validate_expressions enforces the gate +# --------------------------------------------------------------------------- + +ANALYSIS_TEMPLATE = """SOLVER_CAN_BE_APPLIED: True + +SOLVER_FORMULATION: +```python +{code} +``` + +Analysis: some analysis +""" + + +def test_extract_raises_on_malicious_formulation(): + system = Z3SymPySolverSystem("system prompt", client=MagicMock(), model="m") + analysis = ANALYSIS_TEMPLATE.format(code="import os\nos.system('id')") + with pytest.raises(ValueError) as excinfo: + system.extract_and_validate_expressions(analysis) + assert "import of 'os'" in str(excinfo.value) + + +def test_extract_returns_clean_formulation(): + system = Z3SymPySolverSystem("system prompt", client=MagicMock(), model="m") + analysis = ANALYSIS_TEMPLATE.format(code="x = Int('x')\nprint(x)") + assert system.extract_and_validate_expressions(analysis) == "x = Int('x')\nprint(x)" + + +# --------------------------------------------------------------------------- +# Behavior: rejection degrades to standard inference, bypassing the fix loop +# --------------------------------------------------------------------------- + +def test_process_query_degrades_on_malicious_code(): + """Malicious formulation -> straight to standard inference, exec() never runs.""" + malicious = "import os\nos.system('id')" + analyze_resp = _response(ANALYSIS_TEMPLATE.format(code=malicious)) + standard_resp = _response("Direct answer from the model.") + + client = MagicMock() + client.chat.completions.create.side_effect = [analyze_resp, standard_resp] + + with patch("optillm.z3_solver.execute_code_in_process") as exec_mock: + system = Z3SymPySolverSystem("system prompt", client=client, model="m") + result, tokens = system.process_query("solve x > 2") + + assert result == "Direct answer from the model." + exec_mock.assert_not_called() + # Exactly two provider calls: analyze + standard inference. No fix-up calls. + assert client.chat.completions.create.call_count == 2 + for call in client.chat.completions.create.call_args_list[1:]: + sent = str(call) + assert "Fix the Z3 or SymPy code" not in sent, \ + "rejected code must not enter the error-fix loop" + + +def test_process_query_solves_legitimate_formulation(): + """Legitimate formulation still takes the full solver path.""" + code = "x = Int('x')\ns = Solver()\ns.add(x > 2)\nprint(s.check())" + analyze_resp = _response(ANALYSIS_TEMPLATE.format(code=code)) + answer_resp = _response("The constraint is satisfiable: sat.") + + client = MagicMock() + client.chat.completions.create.side_effect = [analyze_resp, answer_resp] + + system = Z3SymPySolverSystem("system prompt", client=client, model="m") + result, tokens = system.process_query("solve x > 2") + + assert "sat" in result + assert client.chat.completions.create.call_count == 2 + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))