T09 · Insecure Skill Coding Practices
Error
- Location
- scripts/utils/predictor/evaluation.py:25
- Finding
- Unsafe PyTorch Checkpoint Deserialization Can Lead to Arbitrary Code Execution<![CDATA[ ## Vulnerability Details **File Location**: `scripts/utils/predictor/evaluation.py:25-40` **Vulnerability Type**: Unsafe deserialization of a potentially untrusted model checkpoint **Risk Level**: High ### Vulnerable Code ```python class NewsModelEvaluator: def __init__(self, model_path=None): self.trainer = AutoSynthesisTrainer() self.device = self.trainer.device if model_path is None: # Try to find the latest model in exports/models model_files = glob.glob(os.path.join(SRC_DIR, "exports/models/*.pt")) if not model_files: logger.warning("⚠️ No trained models found in exports/models/. Using base model (zero-init proj).") else: model_path = max(model_files, key=os.path.getctime) if model_path: self.load_weights(model_path) def load_weights(self, path): logger.info(f"🔄 Loading model weights from {path}...") checkpoint = torch.load(path, map_location=self.device) self.trainer.model.news_proj.load_state_dict(checkpoint['news_proj_state_dict']) logger.success("✅ News projection layer loaded.") ``` ### Technical Analysis The evaluator loads a caller-supplied checkpoint or automatically selects the newest `.pt` file from `exports/models`. It then passes that file to `torch.load()` without explicitly enabling a restricted weights-only mode. PyTorch checkpoint loading has historically used Python pickle-compatible deserialization. On versions or configurations where unrestricted deserialization is active, specially constructed checkpoint objects can execute Python code during loading. This occurs before the subsequent `load_state_dict()` call can validate that the checkpoint contains the expected tensor dictionary. Automatic selection of the newest checkpoint increases the risk: an attacker who can place or replace a file in the model directory does not need to control the c ...[truncated 1500 chars]
- Remediation
- <![CDATA[ ## Remediation Suggestions 1. Explicitly request restricted loading on supported PyTorch versions: ```python checkpoint = torch.load( path, map_location=self.device, weights_only=True, ) ``` 2. Prefer a non-executable weight format such as `safetensors`. 3. Restrict checkpoints to a dedicated trusted directory and resolve the path before use: ```python trusted_dir = (Path(SRC_DIR) / "exports" / "models").resolve() checkpoint_path = Path(path).resolve() if trusted_dir not in checkpoint_path.parents: raise ValueError("Checkpoint is outside the trusted model directory") if checkpoint_path.is_symlink(): raise ValueError("Symbolic-link checkpoints are not permitted") ``` 4. Verify every checkpoint against a trusted SHA-256 digest or digital signature before loading. 5. Validate file ownership and permissions, and ensure the runtime account cannot be tricked into loading files written by untrusted users. 6. Validate the checkpoint structure and tensor shapes after safe parsing. 7. Pin a minimum PyTorch version whose restricted loading behavior has been reviewed, while still specifying `weights_only=True` explicitly. ]]>