Restrict unpickler in async optimizer checkpoint loading (fixes #6471)

colossalai/utils/safetensors.py

@@ -1,5 +1,6 @@
 # a python safetensors serializer modified from https://github.com/huggingface/safetensors/blob/41bd1acf38ad28ac559522d40596c6c802f79453/safetensors/src/tensor.rs#L214
 import json
+import pickle
 import warnings
 from dataclasses import asdict, dataclass
 from typing import Dict, List, Optional, Tuple
@@ -20,7 +21,7 @@
 _TYPES_INV = {v: k for k, v in _TYPES.items()}
 import io
 
-from torch.distributed.distributed_c10d import _pickler, _unpickler
+from torch.distributed.distributed_c10d import _pickler
 
 ASYNC_WRITE_ENTRIES = 32
 
@@ -31,6 +32,51 @@ def _require_async_file_writer():
     return AsyncFileWriter
 
 
+class _RestrictedUnpickler(pickle.Unpickler):
+    """Unpickler for the non-tensor values stored by :func:`_flatten_optim_state_dict`.
+
+    Those values are plain optimizer data (``param_groups`` and non-tensor state
+    entries such as steps and hyperparameters). Only the small allow-list of
+    globals below may be referenced by the pickle stream; anything else -- for
+    example ``os.system`` referenced by a crafted checkpoint file -- is rejected
+    before it can execute. See issue #6471: previously the set of keys to unpickle
+    came from the checkpoint's own metadata and the unpickler was unrestricted,
+    so a malicious ``.safetensors`` optimizer checkpoint achieved arbitrary code
+    execution on load.
+    """
+
+    _ALLOWED_GLOBALS = {
+        ("builtins", "bytearray"),
+        ("builtins", "complex"),
+        ("builtins", "frozenset"),
+        ("builtins", "range"),
+        ("builtins", "set"),
+        ("builtins", "slice"),
+        ("collections", "Counter"),
+        ("collections", "OrderedDict"),
+        ("collections", "defaultdict"),
+        # Tensors nested inside a non-tensor container are rebuilt with the same
+        # constructors that torch.load(weights_only=True) permits.
+        ("torch", "Tensor"),
+        ("torch._utils", "_rebuild_meta_tensor_no_storage"),
+        ("torch._utils", "_rebuild_parameter"),
+        ("torch._utils", "_rebuild_parameter_with_state"),
+        ("torch._utils", "_rebuild_tensor"),
+        ("torch._utils", "_rebuild_tensor_v2"),
+        ("torch.storage", "TypedStorage"),
+        ("torch.storage", "UntypedStorage"),
+        ("torch.storage", "_load_from_bytes"),
+    }
+
+    def find_class(self, module, name):
+        if (module, name) in self._ALLOWED_GLOBALS:
+            return super().find_class(module, name)
+        raise pickle.UnpicklingError(
+            f"global '{module}.{name}' is forbidden when loading optimizer checkpoints; "
+            "only plain data and tensors may be restored from non-tensor entries"
+        )
+
+
 def _object_to_tensor(obj, device):
     f = io.BytesIO()
     _pickler(f).dump(obj)
@@ -45,7 +91,7 @@ def _object_to_tensor(obj, device):
 def _tensor_to_object(tensor, tensor_size):
     tensor = tensor.cpu()
     buf = tensor.numpy().tobytes()[:tensor_size]
-    return _unpickler(io.BytesIO(buf)).load()
+    return _RestrictedUnpickler(io.BytesIO(buf)).load()
 
 
 @dataclass
@@ -107,6 +153,15 @@ def _unflatten_optim_state_dict(flat_dict: dict, metadata: Optional[dict] = None
         non_tensor_keys = json.loads(metadata["non_tensor_keys"])
     else:
         non_tensor_keys = []
+    # The metadata comes from the checkpoint file itself, so only accept the key
+    # shapes _flatten_optim_state_dict() can produce: "param_groups" and
+    # "state<seperator><idx><seperator><field>".
+    for key in non_tensor_keys:
+        if key == "param_groups":
+            continue
+        parts = key.split(seperator)
+        if len(parts) != 3 or parts[0] != "state" or not parts[1].isdigit():
+            raise ValueError(f"invalid non-tensor key in checkpoint metadata: {key!r}")
     flat_dict = {k: _cast_to_object(v) if k in non_tensor_keys else v for k, v in flat_dict.items()}
     if "param_groups" in flat_dict:
         # 3-level dict

tests/test_checkpoint_io/test_safetensors_security.py

@@ -0,0 +1,87 @@
+"""Regression tests for issue #6471.
+
+The async optimizer-checkpoint loader used to unpickle whichever entries a
+checkpoint file's own metadata named, with an unrestricted unpickler -- a
+crafted .safetensors file could execute arbitrary code on load. These tests
+check that (a) legitimate optimizer state dicts still round-trip and
+(b) a malicious checkpoint is rejected without executing its payload.
+"""
+
+import json
+import os
+import pickle
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from colossalai.utils.safetensors import _flatten_optim_state_dict, _unflatten_optim_state_dict, load_flat
+
+
+class _Payload:
+    def __reduce__(self):
+        return (os.system, ("echo pwned > PWNED.txt",))
+
+
+def _write_malicious_checkpoint(path, key):
+    payload = pickle.dumps(_Payload(), protocol=2)
+    tensors = {
+        "param_groups": torch.tensor([0], dtype=torch.uint8),
+        key: torch.tensor(list(payload), dtype=torch.uint8),
+    }
+    save_file(tensors, str(path), metadata={"non_tensor_keys": json.dumps([key])})
+
+
+def test_malicious_checkpoint_does_not_execute(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+    checkpoint = tmp_path / "optimizer.safetensors"
+    # A key with the exact shape the writer produces, carrying pickle payload bytes.
+    _write_malicious_checkpoint(checkpoint, "state.0.candidate")
+    with pytest.raises(pickle.UnpicklingError):
+        load_flat(str(checkpoint))
+    assert not (tmp_path / "PWNED.txt").exists()
+
+
+def test_off_schema_non_tensor_key_rejected(tmp_path):
+    checkpoint = tmp_path / "optimizer.safetensors"
+    _write_malicious_checkpoint(checkpoint, "totally.bogus.key.extra")
+    with pytest.raises(ValueError):
+        load_flat(str(checkpoint))
+
+
+def test_optimizer_state_dict_round_trip(tmp_path):
+    model = torch.nn.Linear(4, 2)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01, betas=(0.9, 0.95))
+    for _ in range(2):
+        loss = model(torch.randn(3, 4)).sum()
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+
+    state_dict = optimizer.state_dict()
+    flat, metadata = _flatten_optim_state_dict(state_dict)
+    restored = _unflatten_optim_state_dict(dict(flat), {"non_tensor_keys": json.dumps(metadata["non_tensor_keys"])})
+    assert restored["param_groups"] == state_dict["param_groups"]
+    for idx, fields in state_dict["state"].items():
+        for name, value in fields.items():
+            restored_value = restored["state"][idx][name]
+            if isinstance(value, torch.Tensor):
+                assert torch.equal(restored_value, value)
+            else:
+                assert restored_value == value
+
+    checkpoint = tmp_path / "optimizer.safetensors"
+    save_file(flat, str(checkpoint), metadata={"non_tensor_keys": json.dumps(metadata["non_tensor_keys"])})
+    assert load_flat(str(checkpoint))["param_groups"] == state_dict["param_groups"]
+
+
+def test_tensor_nested_in_non_tensor_value_round_trips():
+    state_dict = {
+        "state": {0: {"step": 3, "bundle": {"vec": torch.tensor([1.0, 2.0]), "name": "x"}}},
+        "param_groups": [{"lr": 0.1, "params": [0]}],
+    }
+    flat, metadata = _flatten_optim_state_dict(state_dict)
+    restored = _unflatten_optim_state_dict(dict(flat), {"non_tensor_keys": json.dumps(metadata["non_tensor_keys"])})
+    bundle = restored["state"][0]["bundle"]
+    assert bundle["name"] == "x"
+    assert torch.equal(bundle["vec"], torch.tensor([1.0, 2.0]))