diff --git a/test/test_serialization.py b/test/test_serialization.py index 66513ba062e7..bcf480b66020 100644 --- a/test/test_serialization.py +++ b/test/test_serialization.py @@ -919,6 +919,10 @@ class serialization_method: Point = namedtuple('Point', ['x', 'y']) +class ListSubclass(list): + pass + + class ClassThatUsesBuildInstruction: def __init__(self, num): self.num = num @@ -1373,6 +1377,15 @@ class TestSerialization(TestCase, SerializationMixin): with self.assertRaisesRegex(pickle.UnpicklingError, error_msg): torch.load(f, weights_only=True) + def test_weights_only_list_subclass_append(self): + # APPEND/APPENDS would call overridden methods of list subclasses. + with BytesIOContext() as f: + torch.save(ListSubclass([1, 2]), f) + f.seek(0) + with safe_globals([ListSubclass]): + with self.assertRaisesRegex(pickle.UnpicklingError, "Can only extend lists"): + torch.load(f, weights_only=True) + @parametrize("module", ["traceback", "linecache"]) def test_weights_only_traceback_blocklist(self, module): # traceback.FrameSummary and linecache.getline read arbitrary files. diff --git a/torch/_utils.py b/torch/_utils.py index 04cb737a0217..a36dbf4cf7a6 100644 --- a/torch/_utils.py +++ b/torch/_utils.py @@ -1,5 +1,4 @@ # mypy: allow-untyped-defs -import _compat_pickle import copyreg import functools import importlib @@ -1172,16 +1171,6 @@ NAME_MAPPING = { ("UserDict", "UserDict"): ("collections", "UserDict"), } -# Protocol 2 pickle (torch.save's default) maps builtin exceptions to the -# Python 2 "exceptions" module via REVERSE_NAME_MAPPING; map them back so -# allowlisted exception types resolve under their builtins.* names. -NAME_MAPPING.update( - { - ("exceptions", name): ("builtins", name) - for name in _compat_pickle.PYTHON2_EXCEPTIONS - } -) - def _chunk_or_narrow_cat( tensor: "torch.Tensor", diff --git a/torch/_weights_only_unpickler.py b/torch/_weights_only_unpickler.py index 00b819667d25..228d8b9a06dc 100644 --- a/torch/_weights_only_unpickler.py +++ b/torch/_weights_only_unpickler.py @@ -2,7 +2,7 @@ # Unpickler restricted to loading only state dicts # Restrict constructing types to a list defined in _get_allowed_globals() # Restrict BUILD operation to `Tensor`, `Parameter` and `OrderedDict` types only -# Restrict APPEND/APPENDS to `list` (and allowlisted list subclasses) +# Restrict APPEND/APPENDS to `list` # In `GLOBALS` operation do not do class lookup by name, but rather rely on dictionary # defined by `_get_allowed_globals()` method, that contains: # - torch types (Storage, dtypes, Tensor, `torch.Size`), @@ -467,12 +467,18 @@ class Unpickler: elif key[0] == APPEND[0]: item = self.stack.pop() list_obj = self.stack[-1] - self._check_append_target(list_obj) + if type(list_obj) is not list: + raise UnpicklingError( + f"Can only append to lists, but got {type(list_obj)}" + ) list_obj.append(item) elif key[0] == APPENDS[0]: items = self.pop_mark() list_obj = self.stack[-1] - self._check_append_target(list_obj) + if type(list_obj) is not list: + raise UnpicklingError( + f"Can only extend lists, but got {type(list_obj)}" + ) list_obj.extend(items) elif key[0] == SETITEM[0]: (v, k) = (self.stack.pop(), self.stack.pop()) @@ -595,15 +601,6 @@ class Unpickler: f"collections.Counter, but got {type(self.stack[-1])}" ) - def _check_append_target(self, list_obj): - # list subclasses allowlisted via add_safe_globals can be appended to - # (e.g. traceback.StackSummary); plain list is always allowed - if type(list_obj) is not list and not ( - isinstance(list_obj, list) - and type(list_obj) in _get_user_allowed_globals().values() - ): - raise UnpicklingError(f"Can only append to lists, but got {type(list_obj)}") - def persistent_load(self, pid): raise UnpicklingError("unsupported persistent id encountered") diff --git a/torch/distributed/distributed_c10d.py b/torch/distributed/distributed_c10d.py index 3c91e8e7fe06..b847ebf186dd 100644 --- a/torch/distributed/distributed_c10d.py +++ b/torch/distributed/distributed_c10d.py @@ -1,7 +1,6 @@ # mypy: allow-untyped-defs """Distributed Collective Communication (c10d).""" -import builtins import collections.abc import contextlib import ctypes @@ -319,17 +318,6 @@ def _pg_options_to_hints(pg_options: object) -> dict[str, str] | None: _pickler = pickle.Pickler _unpickler = pickle.Unpickler -# Object collectives can serialize via torch.save/torch.load(weights_only=True) -# by passing weights_only=True. Exceptions are commonly transmitted for error -# propagation, so allowlist them. -torch.serialization.add_safe_globals( - [ - t - for t in vars(builtins).values() - if isinstance(t, type) and issubclass(t, BaseException) - ] -) - GroupName = NewType("GroupName", str)