fix(z3): statically gate solver code to the declared z3/sympy subset

optillm/z3_solver.py

@@ -15,14 +15,77 @@
 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.
 
     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]:

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', '<s>', '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"]))