T09 · Insecure Skill Coding Practices
Error
- Location
- scripts/infer_depth.py:31
- Finding
- Unsafe Deserialization of User-Supplied PyTorch Model in Image Depth Inference## Vulnerability Details **File Location**: `scripts/infer_depth.py:31` **Vulnerability Type**: Unsafe PyTorch model deserialization **Risk Level**: High **Complete Code Snippet**: ```python model = DepthAnythingV2(**MODEL_CONFIGS[encoder]) model.load_state_dict(torch.load(model_path, map_location='cpu')) model = model.to(DEVICE).eval() ``` ### Technical Analysis The script accepts `model_path` from a command-line argument and passes it directly to `torch.load`: ```python model_path = sys.argv[1] if len(sys.argv) > 1 else 'depth_anything_v2_vitb.pth' ``` PyTorch model files can use Python pickle-based serialization. On PyTorch versions or configurations where unrestricted deserialization is used, loading an untrusted `.pth` file may invoke attacker-controlled pickle reduction functions and execute arbitrary Python code. This occurs during `torch.load`, before `load_state_dict` can validate whether the result is a legitimate state dictionary. The risk is increased by documentation instructing users to download model weights from an unspecified `hf-mirror` source without an exact immutable URL, revision, checksum, signature, or provenance-verification procedure. The model loading itself is necessary for depth inference, but unrestricted deserialization is not the minimum privilege or safest mechanism required for that functionality. ### Attack Path 1. An attacker creates a malicious `.pth` file containing a pickle payload that invokes an operating-system command during deserialization. 2. The attacker supplies the file directly, replaces a downloaded model, or compromises the unspecified mirror or distribution channel. 3. A user or Agent invokes: ```bash python scripts/infer_depth.py malicious.pth input.png depth.png vitb ``` 4. The script passes `malicious.pth` to `torch.load`. 5. The pickle payload executes before `model.load_state_dict` validates the loaded object. 6. The payload runs with ...[truncated 725 chars]
- Remediation
- ## Remediation Suggestions 1. Use a supported PyTorch release and explicitly enable restricted weight-only loading: ```python state_dict = torch.load( model_path, map_location="cpu", weights_only=True, ) model.load_state_dict(state_dict) ``` 2. Prefer a non-executable tensor serialization format such as `safetensors`. 3. Distribute an exact official model URL pinned to an immutable revision rather than referring generically to a mirror. 4. Publish an expected SHA-256 digest and verify it before loading: ```python import hashlib def sha256_file(path): digest = hashlib.sha256() with open(path, "rb") as source: for chunk in iter(lambda: source.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() ``` 5. Reject the model unless its checksum matches an allowlisted value. 6. Validate that the loaded value is a dictionary containing only expected parameter keys and tensor values. 7. Treat all caller-supplied model files as untrusted. If legacy pickle-based checkpoints must be supported, load them only in an isolated, unprivileged environment with minimal filesystem and network access.