diff --git a/tests/test_ast_validator.py b/tests/test_ast_validator.py new file mode 100644 index 00000000..a878d380 --- /dev/null +++ b/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 diff --git a/tests/test_config_injectable.py b/tests/test_config_injectable.py index 29be38a7..39c4ccac 100644 --- a/tests/test_config_injectable.py +++ b/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) diff --git a/triton_kernel_agent/worker.py b/triton_kernel_agent/worker.py index e594d74c..f25817ad 100644 --- a/triton_kernel_agent/worker.py +++ b/triton_kernel_agent/worker.py @@ -32,7 +32,7 @@ from utils.providers import get_model_provider from .prompt_manager import PromptManager -from .worker_util import _run_test_multiprocess +from .worker_util import _run_test_multiprocess, validate_kernel_ast DISALLOWED_TORCH_PATTERNS = [ @@ -281,6 +281,9 @@ def _validate_kernel_candidate(self, kernel_code: str | None) -> str | None: if not has_kernel_function: return "missing required top-level kernel_function definition" + if not validate_kernel_ast(kernel_code): + return "Kernel rejected: forbidden dynamic import or execution pattern detected." + return None def _write_kernel(self, kernel_code: str): diff --git a/triton_kernel_agent/worker_util.py b/triton_kernel_agent/worker_util.py index e382ac24..615a66f9 100644 --- a/triton_kernel_agent/worker_util.py +++ b/triton_kernel_agent/worker_util.py @@ -14,6 +14,7 @@ """Utility functions for the verification and optimization workers.""" +import ast import multiprocessing as mp import os import re @@ -97,6 +98,65 @@ def _extract_history_usage_from_response( return result if result else None +def validate_kernel_ast(code_str: str) -> bool: + """Validate AST nodes against disallowed execution patterns. + + Args: + code_str: The source code to validate. + + Returns: + bool: True if safe, False if disallowed patterns are found or parsing fails. + """ + try: + tree = ast.parse(code_str) + except SyntaxError: + return False + + class SecurityVisitor(ast.NodeVisitor): + """Node visitor for AST validation.""" + def __init__(self) -> None: + self.is_safe = True + self.unsafe_names = { + "eval", "exec", "__import__", "getattr", + "setattr", "delattr", "compile", "globals", + "locals", "vars", "open", "input", "__builtins__" + } + self.forbidden_modules = {"os", "sys", "subprocess", "builtins", "importlib"} + + def visit_Name(self, node: ast.Name) -> None: + if node.id in self.unsafe_names: + self.is_safe = False + self.generic_visit(node) + + def visit_Attribute(self, node: ast.Attribute) -> None: + if node.attr in self.unsafe_names: + self.is_safe = False + if node.attr.startswith("__") and node.attr.endswith("__"): + self.is_safe = False + self.generic_visit(node) + + def visit_Call(self, node: ast.Call) -> None: + self.generic_visit(node) + + def visit_Import(self, node: ast.Import) -> None: + for alias in node.names: + base_module = alias.name.split(".")[0] + if base_module in self.forbidden_modules: + self.is_safe = False + self.generic_visit(node) + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + if node.module: + base_module = node.module.split(".")[0] + if base_module in self.forbidden_modules: + self.is_safe = False + self.generic_visit(node) + + visitor = SecurityVisitor() + visitor.visit(tree) + return visitor.is_safe + + # ------------------------ # File I/O Utilities # ------------------------