Security: Unsafe torch.load with weights_only=False enables arbitrary code execution

[tomaioo]

Summary

Security: Unsafe torch.load with weights_only=False enables arbitrary code execution

Problem

Severity: Critical | File: src/prime_rl/orchestrator/ckpt.py:L56

In CheckpointManager.load_from_path(), torch.load() is called with weights_only=False (line with torch.load(f, weights_only=False)). This is a critical security vulnerability because it allows arbitrary Python code execution during checkpoint deserialization. An attacker who can control the checkpoint file can execute arbitrary code on the system. The weights_only=False parameter disables PyTorch's safe unpickling mode, which was the default behavior in older PyTorch versions but is known to be dangerous.

Solution

Change torch.load(f, weights_only=False) to torch.load(f, weights_only=True). If the checkpoint requires custom objects, implement a safe deserialization mechanism with explicit allowlists or use a different serialization format like safetensors.

Changes

• src/prime_rl/orchestrator/ckpt.py (modified)

---

Note

Medium Risk
Reduces arbitrary-code-execution risk on untrusted checkpoints; resume may fail if legacy progress.pt files relied on unsafe pickle types.

Overview
Hardens orchestrator resume by loading progress.pt with torch.load(..., weights_only=True) instead of weights_only=False, so checkpoint deserialization no longer runs arbitrary pickle code if a checkpoint file is tampered with.

Only CheckpointManager.load_from_path in ckpt.py changes; save path and buffer loading are unchanged.

^{Reviewed by Cursor Bugbot for commit f322a5f. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit f322a5f. Configure here.}

[cursor]

Safe load breaks progress resume

High Severity

load_from_path now calls torch.load with weights_only=True, but save_to_path still writes progress.pt via torch.save containing a live Progress dataclass instance. Restricted unpickling does not allow that custom type unless it is explicitly allowlisted, so resuming orchestrator checkpoints (when progress loading is not skipped) is likely to fail at load time.

 

^{Reviewed by Cursor Bugbot for commit f322a5f. Configure here.}