T09 · Insecure Skill Coding Practices
- Location
engine/emotion_model/scripts/train.py:112- Finding
Unsafe PyTorch Checkpoint Deserialization Permits Arbitrary Code Execution
- Content
View full analysis
Vulnerability Details
File Location:
engine/emotion_model/scripts/train.py:112-126
Vulnerability Type: Unsafe deserialization of an attacker-controlled PyTorch checkpoint
Risk Level: HighVulnerable Code
python def _load_training_checkpoint(path: Path) -> dict: """Load a rich training checkpoint, validating dimensions.""" checkpoint = torch.load(path, map_location="cpu", weights_only=False) # Validate dimension compatibility saved_dims = checkpoint.get("config_num_dims") if saved_dims is not None and saved_dims != config.NUM_EMOTION_DIMS: raise RuntimeError( f"Checkpoint was trained with {saved_dims} emotion dimensions " f"but current config has {config.NUM_EMOTION_DIMS}. " f"You must retrain from scratch or update your config." ) return checkpointThe checkpoint path can be supplied through the training command:
python elif args.resume is not None: resume_from = args.resumeTechnical Analysis
torch.load()withweights_only=Falsepermits pickle-based object deserialization. Pickle is not a safe format for untrusted input because serialized objects can define reduction behavior that invokes arbitrary Python callables during loading.The
--resumecommand-line argument accepts a user-selected checkpoint path. Consequently, a checkpoint obtained from an untrusted source can execute code as soon as_load_training_checkpoint()callstorch.load(). The dimension validation occurs after deserialization and therefore cannot prevent execution.Attack Path
- An attacker creates a malicious
.ptcheckpoint containing a pickle reduction payload. - The attacker distributes the checkpoint as a pretrained or resumable Emoclaw training artifact.
- A user runs:
bash python -m emotion_model.scripts.train --resume attacker.pt - `_load_training_checkpoint ...[truncated 683 chars]
- An attacker creates a malicious
- Remediation
View remediation
Remediation Suggestions
- Replace the unsafe load with
torch.load(path, map_location="cpu", weights_only=True). - Store non-tensor resume metadata, such as epoch numbers and validation metrics, in JSON or another non-executable format.
- If optimizer state cannot be loaded safely, document that only inference weights can be imported from external sources.
- Restrict resumable checkpoints to a user-owned checkpoint directory with restrictive permissions.
- Cryptographically verify trusted checkpoints before loading them.
- Clearly warn users never to load untrusted PyTorch or pickle-based artifacts.
- Replace the unsafe load with
