Block unsafe PyYAML deserialization in Studio tools

[Imagineer99]

Summary

• detect unsafe PyYAML deserialization in Studio's Python static safety check
• block yaml.unsafe_load(...) and yaml.load(...) unless the loader is clearly SafeLoader/CSafeLoader/BaseLoader/CBaseLoader
• add regression coverage for unsafe PyYAML execution payloads and safe YAML parsing

Validation

• Reproduced before fix: _check_code_safety(...) returned None, auto classifier returned False, and _python_exec(...) printed a harmless YAML_RCE_MARKER via yaml.load(..., Loader=yaml.Loader).
• Verified after fix: the same payload returns Error: unsafe code detected (...), auto classifier returns True, and _python_exec(...) refuses execution.
• PYTHONPATH=studio/backend .venv/bin/python -m pytest studio/backend/tests/test_sandbox_tools.py studio/backend/tests/test_permission_mode.py::test_python_classifier -> 430 passed.
• PYTHONPATH=studio/backend .venv/bin/python -m py_compile studio/backend/core/inference/tools.py studio/backend/tests/test_permission_mode.py studio/backend/tests/test_sandbox_tools.py
• git diff --check

Note

A broader local run of studio/backend/tests/test_sandbox_tools.py studio/backend/tests/test_permission_mode.py reached unrelated route import tests and failed only because the trimmed local venv was missing optional/full backend dependencies such as datasets; the security-relevant sandbox and classifier coverage passed.

[gemini-code-assist]

Code Review

This pull request introduces AST-based safety checks to detect and block unsafe PyYAML deserialization (such as yaml.load with unsafe loaders or yaml.unsafe_load) while permitting safe loaders. It also adds corresponding unit tests. The reviewer identified a potential security bypass where users could alias the load functions themselves (e.g., my_load = yaml.load) to evade detection. To address this, the reviewer suggested updating visit_Assign to track load function aliases and adding test cases to verify that these aliased calls are correctly blocked.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

I am having trouble creating individual review comments. Click here to see my feedback.

studio/backend/core/inference/tools.py (4547-4554)

The current implementation of 'visit_Assign' only tracks aliases for safe loaders (e.g., 'Safe = yaml.SafeLoader'). It does not track assignments of the load functions themselves (e.g., 'my_load = yaml.load' or 'my_unsafe_load = yaml.unsafe_load'). This allows users to bypass the safety check by aliasing the load functions before calling them. We should update 'visit_Assign' to also track assignments of 'load' and 'unsafe_load' attributes.

        def visit_Assign(self, node):
            if _pyyaml_loader_is_safe(
                node.value, self.yaml_aliases, self.yaml_safe_loader_aliases
            ):
                for target in node.targets:
                    if isinstance(target, ast.Name):
                        self.yaml_safe_loader_aliases.add(target.id)
            if isinstance(node.value, ast.Attribute) and isinstance(node.value.value, ast.Name):
                if node.value.value.id in self.yaml_aliases:
                    if node.value.attr == "load":
                        for target in node.targets:
                            if isinstance(target, ast.Name):
                                self.yaml_load_aliases.add(target.id)
                    elif node.value.attr == "unsafe_load":
                        for target in node.targets:
                            if isinstance(target, ast.Name):
                                self.yaml_unsafe_load_aliases.add(target.id)
            self.generic_visit(node)

studio/backend/tests/test_sandbox_tools.py (30-57)

Add test cases to verify that aliased 'yaml.load' and 'yaml.unsafe_load' calls are correctly blocked by the safety check.

    @pytest.mark.parametrize(
        "code",
        [
            (
                "import yaml\n"
                "yaml.load("
                "'!!python/object/apply:os.system [\"echo pwned\"]', "
                "Loader=yaml.Loader"
                ")"
            ),
            (
                "import yaml as y\n"
                "y.load("
                "'!!python/object/apply:os.system [\"echo pwned\"]', "
                "Loader=y.Loader"
                ")"
            ),
            (
                "from yaml import load, Loader\n"
                "load('!!python/object/apply:os.system [\"echo pwned\"]', Loader=Loader)"
            ),
            (
                "import yaml\n"
                "yaml.unsafe_load('!!python/object/apply:os.system [\"echo pwned\"]')"
            ),
            "import yaml\nloader = get_loader()\nyaml.load('a: 1', Loader=loader)",
            (
                "import yaml\n"
                "my_load = yaml.load\n"
                "my_load('!!python/object/apply:os.system [\"echo pwned\"]')"
            ),
            (
                "import yaml\n"
                "my_unsafe_load = yaml.unsafe_load\n"
                "my_unsafe_load('!!python/object/apply:os.system [\"echo pwned\"]')"
            ),
        ],
    )