From a8eaccd2b660f77c29b21a4cc250030cd569ee6e Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Sat, 20 Jun 2026 15:10:17 -0700 Subject: [PATCH] Copybara import of the project: -- 6ab68a057443f038a3d844807f205b1c99cab944 by Kevin Zakka : Fix pickle code-execution vulnerability in sysid loaders Serialize signal_mapping as JSON so the trajectory and time series loaders can use allow_pickle=False, preventing arbitrary code execution from untrusted .npz files. COPYBARA_INTEGRATE_REVIEW=https://github.com/google-deepmind/mujoco/pull/3353 from kevinzakka:sysid-disable-pickle-load 6ab68a057443f038a3d844807f205b1c99cab944 PiperOrigin-RevId: 935400242 Change-Id: Iecd907174441fbcfd02c106b912a7ff375c9098c --- doc/changelog.rst | 3 ++ python/mujoco/sysid/_src/timeseries.py | 38 ++++++++++++++++++++------ python/mujoco/sysid/_src/trajectory.py | 29 ++++++++++++-------- 3 files changed, 51 insertions(+), 19 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 65cafedf..b7e0154c 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -66,6 +66,9 @@ General Bug fixes ^^^^^^^^^ +- Fixed a vulnerability in the System Identification toolbox where loading a trajectory or time series called + ``np.load`` with ``allow_pickle=True``, allowing arbitrary code execution from a malicious ``.npz`` file. Signal + metadata is now serialized as JSON and loaded with ``allow_pickle=False``. - Fixed a bug in the ``mjz`` :ref:`decoder ` where unnormalized paths would fail to be read. - Fixed a bug where the mesh compiler would produce non-unit convex hull polygon normals. diff --git a/python/mujoco/sysid/_src/timeseries.py b/python/mujoco/sysid/_src/timeseries.py index 8fef0228..73d2204b 100644 --- a/python/mujoco/sysid/_src/timeseries.py +++ b/python/mujoco/sysid/_src/timeseries.py @@ -20,6 +20,7 @@ from __future__ import annotations from collections.abc import Sequence from dataclasses import dataclass from enum import Enum +import json import pathlib from typing import Literal, TypeAlias @@ -39,6 +40,27 @@ class SignalType(Enum): SignalMappingType: TypeAlias = dict[str, tuple[SignalType, np.ndarray]] + +def encode_signal_mapping(signal_mapping: SignalMappingType) -> str: + """Serialize a signal mapping to a JSON string. + + This avoids storing object arrays (which require ``allow_pickle=True`` to + load) on disk. The mapping is plain data: enum values and integer indices. + """ + return json.dumps({ + name: [sig_type.value, np.asarray(indices).astype(int).tolist()] + for name, (sig_type, indices) in signal_mapping.items() + }) + + +def decode_signal_mapping(encoded: str) -> SignalMappingType: + """Inverse of :func:`encode_signal_mapping`.""" + raw = json.loads(encoded) + return { + name: (SignalType(sig_type), np.asarray(indices, dtype=int)) + for name, (sig_type, indices) in raw.items() + } + InterpolationMethod = Literal[ "linear", "cubic", "quadratic", "quintic", "zero_order_hold", "zoh" ] @@ -569,12 +591,12 @@ class TimeSeries: Args: path: Path where the data will be saved. """ - np.savez( - path, - times=self.times, - data=self.data, - signal_mapping=np.array(self.signal_mapping, dtype=object), - ) + save_dict = {"times": self.times, "data": self.data} + if self.signal_mapping: + save_dict["signal_mapping"] = np.asarray( + encode_signal_mapping(self.signal_mapping) + ) + np.savez(path, **save_dict) def save_to_csv(self, path: str | pathlib.Path) -> None: """Save the time series data to a CSV file. @@ -598,11 +620,11 @@ class TimeSeries: Returns: A new TimeSeries object. """ - with np.load(path, allow_pickle=True) as npz: + with np.load(path, allow_pickle=False) as npz: times = npz["times"] data = npz["data"] if "signal_mapping" in npz: - signal_mapping = npz["signal_mapping"].item() + signal_mapping = decode_signal_mapping(str(npz["signal_mapping"])) else: signal_mapping = None diff --git a/python/mujoco/sysid/_src/trajectory.py b/python/mujoco/sysid/_src/trajectory.py index 1719b3c5..dade545f 100644 --- a/python/mujoco/sysid/_src/trajectory.py +++ b/python/mujoco/sysid/_src/trajectory.py @@ -102,18 +102,19 @@ class SystemTrajectory: if self.state is not None: save_dict["state_times"] = self.state.times save_dict["state_data"] = self.state.data - save_dict["state_signal_mapping"] = np.array( - self.state.signal_mapping, dtype=object - ) + if self.state.signal_mapping: + save_dict["state_signal_mapping"] = np.asarray( + timeseries.encode_signal_mapping(self.state.signal_mapping) + ) if self.control.signal_mapping: - save_dict["control_signal_mapping"] = np.array( - self.control.signal_mapping, dtype=object + save_dict["control_signal_mapping"] = np.asarray( + timeseries.encode_signal_mapping(self.control.signal_mapping) ) if self.sensordata.signal_mapping: - save_dict["sensordata_signal_mapping"] = np.array( - self.sensordata.signal_mapping, dtype=object + save_dict["sensordata_signal_mapping"] = np.asarray( + timeseries.encode_signal_mapping(self.sensordata.signal_mapping) ) np.savez(path, **save_dict) # type: ignore @@ -126,7 +127,7 @@ class SystemTrajectory: allow_missing_sensors: bool = False, ) -> SystemTrajectory: """Load a trajectory from a compressed NumPy archive.""" - with np.load(path, allow_pickle=True) as npz: + with np.load(path, allow_pickle=False) as npz: control_times = npz["control_times"] control_data = npz["control_data"] sensordata_times = npz["sensordata_times"] @@ -137,15 +138,21 @@ class SystemTrajectory: control_signal_mapping = None if "control_signal_mapping" in npz: - control_signal_mapping = npz["control_signal_mapping"].item() + control_signal_mapping = timeseries.decode_signal_mapping( + str(npz["control_signal_mapping"]) + ) sensordata_signal_mapping = None if "sensordata_signal_mapping" in npz: - sensordata_signal_mapping = npz["sensordata_signal_mapping"].item() + sensordata_signal_mapping = timeseries.decode_signal_mapping( + str(npz["sensordata_signal_mapping"]) + ) state_signal_mapping = None if "state_signal_mapping" in npz: - state_signal_mapping = npz["state_signal_mapping"].item() + state_signal_mapping = timeseries.decode_signal_mapping( + str(npz["state_signal_mapping"]) + ) predicted_rollout = cls( model=model,