wehub-resource-sync
86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
"""Regression guard: every ``torch.load`` call must set ``weights_only=True``.
|
|
|
|
VoxCPM deliberately loads checkpoints with ``weights_only=True`` so that a
|
|
crafted ``.ckpt``/``.pth``/``.bin`` file cannot execute arbitrary code via
|
|
pickle during unpickling (see
|
|
``tests/test_lora_checkpoint_loading.py::test_load_lora_weights_rejects_malicious_pickle_payloads``).
|
|
|
|
The fine-tuning resume path in ``scripts/train_voxcpm_finetune.py`` originally
|
|
called ``torch.load`` without that flag, leaving an arbitrary-code-execution
|
|
gap when resuming from an attacker-supplied checkpoint directory. This test
|
|
statically asserts the flag is present on every ``torch.load`` call across the
|
|
package and scripts so the gap cannot silently reappear.
|
|
"""
|
|
|
|
import ast
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
# Directories whose Python files load checkpoints at runtime / on resume.
|
|
SCANNED_DIRS = [REPO_ROOT / "src", REPO_ROOT / "scripts", REPO_ROOT / "app.py", REPO_ROOT / "lora_ft_webui.py"]
|
|
|
|
|
|
def _python_files():
|
|
for entry in SCANNED_DIRS:
|
|
if entry.is_file() and entry.suffix == ".py":
|
|
yield entry
|
|
elif entry.is_dir():
|
|
yield from entry.rglob("*.py")
|
|
|
|
|
|
def _is_torch_load(node: ast.Call) -> bool:
|
|
func = node.func
|
|
# Matches ``torch.load(...)`` and ``load(...)`` aliased from torch.
|
|
if isinstance(func, ast.Attribute) and func.attr == "load":
|
|
return isinstance(func.value, ast.Name) and func.value.id == "torch"
|
|
return False
|
|
|
|
|
|
def _has_weights_only_true(node: ast.Call) -> bool:
|
|
for kw in node.keywords:
|
|
if kw.arg == "weights_only":
|
|
return isinstance(kw.value, ast.Constant) and kw.value.value is True
|
|
return False
|
|
|
|
|
|
def test_every_torch_load_sets_weights_only_true():
|
|
offenders = []
|
|
checked = 0
|
|
for path in _python_files():
|
|
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Call) and _is_torch_load(node):
|
|
checked += 1
|
|
if not _has_weights_only_true(node):
|
|
offenders.append(f"{path.relative_to(REPO_ROOT)}:{node.lineno}")
|
|
|
|
assert checked > 0, "expected to find at least one torch.load call to verify"
|
|
assert not offenders, (
|
|
"torch.load without weights_only=True (pickle RCE risk):\n "
|
|
+ "\n ".join(offenders)
|
|
)
|
|
|
|
|
|
def test_torch_load_weights_only_blocks_malicious_pickle(tmp_path):
|
|
"""Behavioral check that weights_only=True actually rejects a code-exec payload."""
|
|
torch = pytest.importorskip("torch")
|
|
|
|
marker = tmp_path / "pwned.txt"
|
|
|
|
class Exploit:
|
|
def __reduce__(self):
|
|
import pathlib
|
|
|
|
return (pathlib.Path.write_text, (marker, "executed\n"))
|
|
|
|
ckpt = tmp_path / "optimizer.pth"
|
|
torch.save({"state_dict": {"w": torch.zeros(1)}, "boom": Exploit()}, ckpt)
|
|
|
|
with pytest.raises(Exception):
|
|
torch.load(ckpt, map_location="cpu", weights_only=True)
|
|
|
|
assert not marker.exists(), "malicious pickle executed despite weights_only=True"
|