T09 · Insecure Skill Coding Practices
Error
- Location
- scripts/internalize.py:67
- Finding
- Arbitrary Code Execution Through Unsafe PyTorch Checkpoint Deserialization<![CDATA[ ## Vulnerability Details **File Location**: `scripts/internalize.py:67` **Vulnerability Type**: Unsafe deserialization of a user-selectable checkpoint **Risk Level**: High ### Vulnerable Code ```python # weights_only=False is needed: checkpoint contains config dataclasses # (AggregatorConfig, LoraConfig, HypernetConfig) not just tensors. # Only load from trusted sources. state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=False) ``` The checkpoint path is exposed through a command-line argument: ```python parser.add_argument( "--checkpoint", default="trained_d2l/gemma_demo/checkpoint-80000/pytorch_model.bin", help="Path to D2L checkpoint (only load from trusted sources)", ) ``` ### Technical Analysis `torch.load(..., weights_only=False)` supports Python pickle-compatible object deserialization. Pickle is not a data-only serialization format: specially constructed objects can invoke attacker-controlled functions while being deserialized. The `--checkpoint` option allows the caller to select an arbitrary local checkpoint. The warning that checkpoints must be trusted is documentation only and does not enforce provenance, file integrity, an approved directory, or a cryptographic digest. The default checkpoint is also obtained from an external model repository without an immutable revision or an application-level hard-coded digest. Consequently, local checkpoint replacement, a malicious checkpoint supplied by another user, or compromise of the model distribution chain can reach this dangerous deserialization operation. ### Attack Path 1. An attacker creates a malicious PyTorch checkpoint containing a serialized object with a code-execution reduction method. 2. The attacker convinces a user or agent to run the Skill with `--checkpoint /path/to/malicious.bin`, replaces the checkpoint at the default path, or compromises an upstream checkpoint source. 3. `load_model()` passes the file to `torch.load()` with `weights_ ...[truncated 784 chars]
- Remediation
- <![CDATA[ ## Remediation Suggestions 1. Replace pickle-based checkpoints with a data-only format such as Safetensors. 2. Use `torch.load(..., weights_only=True)` wherever possible. 3. If legacy configuration objects are unavoidable, explicitly allowlist only the minimum required classes instead of permitting unrestricted object reconstruction. 4. Pin the official checkpoint to an immutable repository revision and verify it against a hard-coded SHA-256 digest before deserialization. 5. Restrict checkpoint loading to an approved model directory and reject symbolic links or unexpected file types. 6. Treat integrity verification as mandatory and fail closed if the revision, digest, or expected checkpoint structure does not match. 7. Consider converting the trusted legacy checkpoint to a safe tensor-only representation in an isolated environment and distributing only the converted artifact. ]]>
