Fix: AST-based sandbox validation (#62) and Windows test compatibility

tests/test_ast_validator.py

@@ -0,0 +1,134 @@
+"""Tests for AST-based static code validation."""
+
+import ast
+import re
+
+import pytest
+from triton_kernel_agent.worker_util import validate_kernel_ast
+
+
+def naive_string_validator(code_str: str) -> bool:
+    """A regex-based validator for baseline comparison."""
+    if re.search(r'\bimport\s+(os|sys|importlib)\b', code_str):
+        return False
+    if re.search(r'\beval\s*\(', code_str):
+        return False
+    if re.search(r'\bexec\s*\(', code_str):
+        return False
+    return True
+
+
+def test_ast_validator_allows_safe_code() -> None:
+    """Tests standard Python code."""
+    safe_code = '''
+def kernel_function(x):
+    return x + 1
+'''
+    assert naive_string_validator(safe_code) is True
+    assert validate_kernel_ast(safe_code) is True
+
+
+def test_ast_validator_allows_safe_triton_code() -> None:
+    """Tests complex legitimate Triton kernels."""
+    triton_code = '''
+import triton
+import triton.language as tl
+
+@triton.jit
+def kernel_function(
+    x_ptr,
+    y_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(x_ptr + offsets, mask=mask)
+    output = x * 2.0
+    tl.store(y_ptr + offsets, output, mask=mask)
+'''
+    assert naive_string_validator(triton_code) is True
+    assert validate_kernel_ast(triton_code) is True
+
+
+def test_ast_validator_catches_import_os() -> None:
+    """Tests direct module imports."""
+    code = "import os\nos.system('ls')"
+    assert naive_string_validator(code) is False
+    assert validate_kernel_ast(code) is False
+
+
+def test_ast_validator_catches_dynamic_import_bypass() -> None:
+    """Tests string concatenated __import__ calls."""
+    code = "__import__('o' + 's').system('ls')"
+    assert naive_string_validator(code) is True
+    assert validate_kernel_ast(code) is False
+
+
+def test_ast_validator_catches_getattr_bypass() -> None:
+    """Tests dynamic method retrieval via getattr."""
+    code = "getattr(__import__('os'), 'system')('ls')"
+    assert naive_string_validator(code) is True
+    assert validate_kernel_ast(code) is False
+
+
+def test_ast_validator_catches_eval_bypass() -> None:
+    """Tests aliased eval calls."""
+    code = "f = eval; f('1 + 1')"
+    assert naive_string_validator(code) is True
+    assert validate_kernel_ast(code) is False
+
+
+def test_ast_validator_catches_importlib() -> None:
+    """Tests dynamic module imports via importlib."""
+    code = "import importlib\nos = importlib.import_module('os')\nos.system('ls')"
+    assert naive_string_validator(code) is False
+    assert validate_kernel_ast(code) is False
+
+
+def test_ast_validator_catches_sys_modules() -> None:
+    """Tests accessing sys.modules."""
+    code = "import sys\nsys.modules['os'].system('ls')"
+    assert naive_string_validator(code) is False
+    assert validate_kernel_ast(code) is False
+
+
+def test_ast_validator_catches_builtins_dict() -> None:
+    """Tests accessing builtins.__dict__."""
+    code = "import builtins\nbuiltins.__dict__['eval']('1+1')"
+    assert naive_string_validator(code) is True
+    assert validate_kernel_ast(code) is False
+
+
+def test_ast_validator_catches_nested_eval() -> None:
+    """Tests nested eval/exec calls."""
+    code = "exec(eval('\"import os; os.system(\\'ls\\')\"'))"
+    assert naive_string_validator(code) is False
+    assert validate_kernel_ast(code) is False
+
+
+def test_ast_validator_catches_string_constructed_getattr() -> None:
+    """Tests getattr calls constructed with strings."""
+    code = "f = getattr; f(sys, 'modules')"
+    assert naive_string_validator(code) is True
+    assert validate_kernel_ast(code) is False
+
+
+def test_curve_ball_1_syntax_error() -> None:
+    """Tests handling of invalid syntax."""
+    code = "def kernel_function( :"
+    assert validate_kernel_ast(code) is False
+
+
+def test_curve_ball_2_dictionary_lookup() -> None:
+    """Tests dictionary lookup on builtins."""
+    code = "__builtins__['__import__']('os')"
+    assert validate_kernel_ast(code) is False
+
+
+def test_curve_ball_2_function_aliasing() -> None:
+    """Tests function aliasing."""
+    code = "sneaky_get = getattr; sneaky_get(tl, 'system')"
+    assert validate_kernel_ast(code) is False

tests/test_config_injectable.py

@@ -32,24 +32,25 @@ def test_config_injectable_all_args_provided():
 
 def test_config_injectable_from_yaml():
     """All arguments filled from a YAML config file."""
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("a: 5\nb: 10\nc: 20\n")
-        f.flush()
-        try:
-            assert sample_func(config=f.name) == 35
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("a: 5\nb: 10\nc: 20\n")
+    # Explicitly close the file before unlinking to prevent Windows file-locking issues (PermissionError: [WinError 32])
+    f.close()
+    try:
+        assert sample_func(config=f.name) == 35
+    finally:
+        os.unlink(f.name)
 
 
 def test_config_injectable_partial_override():
     """Positional args take precedence; config fills the rest."""
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("b: 100\nc: 200\n")
-        f.flush()
-        try:
-            assert sample_func(1, config=f.name) == 301
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("b: 100\nc: 200\n")
+    f.close()
+    try:
+        assert sample_func(1, config=f.name) == 301
+    finally:
+        os.unlink(f.name)
 
 
 def test_config_injectable_missing_required():
@@ -67,13 +68,13 @@ def add(self, a, b, c=10):
 def test_config_injectable_class_method():
     """Decorator works on instance methods; self is bound normally."""
     obj = SampleClass()
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("b: 100\nc: 200\n")
-        f.flush()
-        try:
-            assert obj.add(1, config=f.name) == 301
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("b: 100\nc: 200\n")
+    f.close()
+    try:
+        assert obj.add(1, config=f.name) == 301
+    finally:
+        os.unlink(f.name)
 
 
 @config_injectable
@@ -94,42 +95,42 @@ def test_config_injectable_class_init_no_config():
 
 def test_config_injectable_class_init():
     """Decorator on a class injects config into __init__."""
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("a: 5\nb: 10\nc: 20\n")
-        f.flush()
-        try:
-            obj = SampleClassDecorated(config=f.name)
-            assert obj.a == 5
-            assert obj.b == 10
-            assert obj.c == 20
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("a: 5\nb: 10\nc: 20\n")
+    f.close()
+    try:
+        obj = SampleClassDecorated(config=f.name)
+        assert obj.a == 5
+        assert obj.b == 10
+        assert obj.c == 20
+    finally:
+        os.unlink(f.name)
 
 
 def test_config_injectable_class_init_partial():
     """Explicit args override config values during class init."""
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("a: 999\nb: 100\nc: 200\n")
-        f.flush()
-        try:
-            obj = SampleClassDecorated(1, config=f.name)
-            assert obj.a == 1  # explicit overrides yaml's 999
-            assert obj.b == 100
-            assert obj.c == 200
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("a: 999\nb: 100\nc: 200\n")
+    f.close()
+    try:
+        obj = SampleClassDecorated(1, config=f.name)
+        assert obj.a == 1  # explicit overrides yaml's 999
+        assert obj.b == 100
+        assert obj.c == 200
+    finally:
+        os.unlink(f.name)
 
 
 def test_config_injectable_missing_required_from_yaml():
     """Raises TypeError when YAML config only provides some required arguments."""
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("a: 5\n")
-        f.flush()
-        try:
-            with pytest.raises(TypeError, match="Missing required arguments"):
-                sample_func(config=f.name)
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("a: 5\n")
+    f.close()
+    try:
+        with pytest.raises(TypeError, match="Missing required arguments"):
+            sample_func(config=f.name)
+    finally:
+        os.unlink(f.name)
 
 
 # ---------------------------------------------------------------------------
@@ -155,29 +156,29 @@ def test_var_keyword_no_config():
 
 def test_var_keyword_with_config():
     """Named params from YAML + explicit extra kwargs coexist."""
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("name: from_yaml\nmode: fast\n")
-        f.flush()
-        try:
-            obj = ClassWithVarKw(config=f.name, warmup=25, repeat=100)
-            assert obj.name == "from_yaml"
-            assert obj.mode == "fast"
-            assert obj.extra == {"warmup": 25, "repeat": 100}
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("name: from_yaml\nmode: fast\n")
+    f.close()
+    try:
+        obj = ClassWithVarKw(config=f.name, warmup=25, repeat=100)
+        assert obj.name == "from_yaml"
+        assert obj.mode == "fast"
+        assert obj.extra == {"warmup": 25, "repeat": 100}
+    finally:
+        os.unlink(f.name)
 
 
 def test_var_keyword_no_extras():
     """VAR_KEYWORD param doesn't leak into the result when empty."""
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("name: solo\n")
-        f.flush()
-        try:
-            obj = ClassWithVarKw(config=f.name)
-            assert obj.name == "solo"
-            assert obj.extra == {}
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("name: solo\n")
+    f.close()
+    try:
+        obj = ClassWithVarKw(config=f.name)
+        assert obj.name == "solo"
+        assert obj.extra == {}
+    finally:
+        os.unlink(f.name)
 
 
 @config_injectable
@@ -187,39 +188,39 @@ def func_with_var_kw(a, b=10, **opts):
 
 def test_var_keyword_function():
     """VAR_KEYWORD works correctly on decorated functions too."""
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("a: 1\nb: 2\n")
-        f.flush()
-        try:
-            result = func_with_var_kw(config=f.name, debug=True)
-            assert result == {"a": 1, "b": 2, "opts": {"debug": True}}
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("a: 1\nb: 2\n")
+    f.close()
+    try:
+        result = func_with_var_kw(config=f.name, debug=True)
+        assert result == {"a": 1, "b": 2, "opts": {"debug": True}}
+    finally:
+        os.unlink(f.name)
 
 
 def test_var_keyword_yaml_extras_forwarded():
     """YAML keys that don't match named params are forwarded via **kwargs."""
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("a: 1\nb: 2\nextra_one: 100\nextra_two: 200\n")
-        f.flush()
-        try:
-            result = func_with_var_kw(config=f.name)
-            assert result == {
-                "a": 1,
-                "b": 2,
-                "opts": {"extra_one": 100, "extra_two": 200},
-            }
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("a: 1\nb: 2\nextra_one: 100\nextra_two: 200\n")
+    f.close()
+    try:
+        result = func_with_var_kw(config=f.name)
+        assert result == {
+            "a": 1,
+            "b": 2,
+            "opts": {"extra_one": 100, "extra_two": 200},
+        }
+    finally:
+        os.unlink(f.name)
 
 
 def test_var_keyword_yaml_extras_explicit_override():
     """Explicit kwargs override same-named YAML extras."""
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
-        f.write("a: 1\nextra: from_yaml\n")
-        f.flush()
-        try:
-            result = func_with_var_kw(config=f.name, extra="from_kwarg")
-            assert result == {"a": 1, "b": 10, "opts": {"extra": "from_kwarg"}}
-        finally:
-            os.unlink(f.name)
+    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
+    f.write("a: 1\nextra: from_yaml\n")
+    f.close()
+    try:
+        result = func_with_var_kw(config=f.name, extra="from_kwarg")
+        assert result == {"a": 1, "b": 10, "opts": {"extra": "from_kwarg"}}
+    finally:
+        os.unlink(f.name)