Validate AI-Generated Code Safely

Validate AI-Generated Code Safely

When developers integrate generative AI into software engineering workflows, automatically generated code snippets can introduce security vulnerabilities (e.g., command injection, dangerous built-in execution) or reliability risks (e.g., deeply nested loops, unauthorized network calls). Before sending generated code to an automated sandbox or recommending it to an engineer, your system must statically validate the code against safety policies.

Represent the code snippet as a hierarchical tree of AST nodes (represented as Python dictionaries). Implement a function validate_code_ast(root: dict, allowed_imports: set[str], max_loop_depth: int) -> tuple[bool, list[str]] that traverses the tree and reports all policy violations.

AST Node Specifications

Each node in the tree is a dictionary with a "type" key and type-specific attributes:

• "Block": contains "body": a list of child nodes.
• "Import": contains "module": a string representing the module name.
• "Call": contains "func": a string of the function name, and "args": a list of argument nodes.
• "Loop": contains "body": a list of child nodes.
• "BinOp": contains "op": an operator string (e.g., "+"), "left": a node, and "right": a node.
• "Format": represents formatted string interpolation; contains "values": a list of nodes.
• "Literal": contains "value": a primitive value (string, integer, etc.).
• "Variable": contains "name": a string.

Validation Policies

1. Forbidden Calls: No Call node may have func in {"eval", "exec", "system", "spawn"}.
2. Import Allowlist: Every Import node must have its module present in allowed_imports.
3. Loop Depth: Nested Loop nodes must not exceed max_loop_depth. A top-level loop has depth 1.
4. Query Injection Risk: If a Call node has func == "db_query", none of its argument nodes (or their descendants) may be of type "BinOp" or "Format" (to prevent unsafe raw SQL string concatenation).

If any policy is violated, the code is invalid (False), and all violations must be returned as descriptive strings in a deterministic order (pre-order traversal). If no policies are violated, return (True, []).

Function Signature

def validate_code_ast(
    root: dict,
    allowed_imports: set[str],
    max_loop_depth: int
) -> tuple[bool, list[str]]:
    pass

Examples

Example 1:

ast = {
    "type": "Block",
    "body": [
        {"type": "Import", "module": "math"},
        {"type": "Call", "func": "eval", "args": [{"type": "Literal", "value": "1 + 1"}]}
    ]
}
allowed_imports = {"math", "json"}
max_loop_depth = 2

is_valid, violations = validate_code_ast(ast, allowed_imports, max_loop_depth)
# is_valid: False
# violations: ["Forbidden call: eval"]

Example 2:

ast = {
    "type": "Loop",
    "body": [
        {
            "type": "Loop",
            "body": [
                {"type": "Call", "func": "print", "args": [{"type": "Literal", "value": "hello"}]}
            ]
        }
    ]
}
allowed_imports = set()
max_loop_depth = 1

is_valid, violations = validate_code_ast(ast, allowed_imports, max_loop_depth)
# is_valid: False
# violations: ["Exceeded maximum loop depth: 2 > 1"]

Constraints

• 1 <= total AST nodes <= 10,000
• 1 <= max_loop_depth <= 100
• AST structure is guaranteed to be acyclic and conform to the node specifications.