From 9475d76790365dc0c7947d8e50c66f9ac6a9a1f8 Mon Sep 17 00:00:00 2001 From: M Platypus Date: Tue, 4 Aug 2026 17:46:56 -0400 Subject: [PATCH] fix: unsafe deserialization in load_weights()'s pretrained-checkpoint loading torch.load(pretrained_file, map_location=device, weights_only=False) fully disabled PyTorch's deserialization safety check. pretrained_file can be a local path or a URL fetched via data_utils.download_url with no integrity check, so an untrusted or spoofed checkpoint would get arbitrary code execution during unpickling -- the same bug class as an already-fixed issue in resume_from_checkpoint (train_base.py), but missed here. Fixed by allowlisting only what this function's own inputs legitimately need: torch.serialization.safe_globals([Namespace]), since checkpoints written by this project's own save_checkpoint() carry an argparse.Namespace under the 'args' key, and such a checkpoint is a plausible 'pretrained' input to load_weights() too. Verified the added test actually distinguishes vulnerable from fixed behavior, not just "any exception happened": with weights_only=False restored, the untrusted payload's code genuinely executes during unpickling (captured via stdout) and the function fails later for an unrelated reason (treating the reconstructed object as a dict) -- a bare pytest.raises(Exception) would not have told these apart. Co-Authored-By: Claude Sonnet 5 --- tinyml-tinyverse/tests/__init__.py | 0 ...est_load_weights_unsafe_deserialization.py | 88 +++++++++++++++++++ .../common/utils/load_weights.py | 11 ++- 3 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 tinyml-tinyverse/tests/__init__.py create mode 100644 tinyml-tinyverse/tests/test_load_weights_unsafe_deserialization.py diff --git a/tinyml-tinyverse/tests/__init__.py b/tinyml-tinyverse/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tinyml-tinyverse/tests/test_load_weights_unsafe_deserialization.py b/tinyml-tinyverse/tests/test_load_weights_unsafe_deserialization.py new file mode 100644 index 00000000..f64ba61c --- /dev/null +++ b/tinyml-tinyverse/tests/test_load_weights_unsafe_deserialization.py @@ -0,0 +1,88 @@ +"""Regression tests for load_weights()'s torch.load() deserialization safety. + +load_weights() previously called torch.load(..., weights_only=False), fully disabling +PyTorch's deserialization safety check on a file that can be fetched from a URL with +no integrity check (data_utils.download_url). An untrusted/spoofed "pretrained" +checkpoint would get arbitrary code execution during unpickling. +""" +import argparse +import os +import pickle +import tempfile + +import pytest +import torch +import torch.nn as nn + +from tinyml_tinyverse.common.utils.load_weights import load_weights + + +class _TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.fc = nn.Linear(4, 2) + + +class _NotOnTheSafeGlobalsAllowlist: + """Stand-in for an arbitrary class an attacker's payload might reference. + + Not a real exploit gadget -- the point is only that torch.load() must refuse to + unpickle a type it doesn't know is safe, not that this specific class is harmful. + """ + + def __reduce__(self): + # If this ever actually executes during unpickling, the safety check failed. + return (print, ("UNSAFE DESERIALIZATION EXECUTED",)) + + +def test_load_weights_loads_a_legitimate_state_dict(): + model = _TinyModel() + source = _TinyModel() + with tempfile.TemporaryDirectory() as tmp_dir: + weights_path = os.path.join(tmp_dir, "weights.pth") + torch.save(source.state_dict(), weights_path) + + load_weights(model, weights_path) + + for p1, p2 in zip(model.parameters(), source.parameters()): + assert torch.equal(p1, p2) + + +def test_load_weights_loads_a_checkpoint_with_namespace_under_args(): + """Backward-compat: this project's own save_checkpoint() (train_base.py) writes + checkpoint['args'] as an argparse.Namespace. Such a checkpoint is a plausible + 'pretrained' input to load_weights() too, and must still load -- not just a raw + state_dict file.""" + model = _TinyModel() + source = _TinyModel() + with tempfile.TemporaryDirectory() as tmp_dir: + checkpoint_path = os.path.join(tmp_dir, "checkpoint.pth") + torch.save( + {"state_dict": source.state_dict(), "args": argparse.Namespace(epochs=10)}, + checkpoint_path, + ) + + load_weights(model, checkpoint_path) + + for p1, p2 in zip(model.parameters(), source.parameters()): + assert torch.equal(p1, p2) + + +def test_load_weights_rejects_untrusted_pickle_payload(capsys): + """With weights_only=False (the pre-fix behavior), unpickling this payload + would *succeed* -- the class is importable from this test module -- and only + fail later for an unrelated reason (treating the reconstructed object as a + dict), which a bare `pytest.raises(Exception)` can't tell apart from a real + security rejection. Assert on the actual security property instead: the + reduce-callable (standing in for arbitrary code) must never run, and the + failure must be attributable to the deserialization safety check itself.""" + model = _TinyModel() + with tempfile.TemporaryDirectory() as tmp_dir: + payload_path = os.path.join(tmp_dir, "malicious.pth") + torch.save({"state_dict": _NotOnTheSafeGlobalsAllowlist()}, payload_path) + + with pytest.raises(pickle.UnpicklingError, match="Weights only load failed"): + load_weights(model, payload_path) + + # The strongest possible check: the payload's code must never have run. + assert "UNSAFE DESERIALIZATION EXECUTED" not in capsys.readouterr().out diff --git a/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py b/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py index aca21d19..c1c04dbe 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py +++ b/tinyml-tinyverse/tinyml_tinyverse/common/utils/load_weights.py @@ -32,6 +32,7 @@ import re import torch import copy +from argparse import Namespace from collections import OrderedDict from . import print_utils from . import data_utils @@ -62,7 +63,15 @@ def load_weights(model, pretrained, change_names_dict=None, keep_original_names= else: pretrained_file = pretrained # - data = torch.load(pretrained_file, map_location=device, weights_only=False) + # weights_only=False fully disables PyTorch's deserialization safety check -- + # pretrained_file can be a URL fetched with no integrity check (download_url + # above), so an untrusted/spoofed checkpoint would get arbitrary code + # execution during unpickling. Namespace is allow-listed because checkpoints + # written by this project's own save_checkpoint() (train_base.py) legitimately + # carry one under the 'args' key, and such a checkpoint is a plausible + # 'pretrained' input here too. + with torch.serialization.safe_globals([Namespace]): + data = torch.load(pretrained_file, map_location=device) else: data = pretrained #