fix(security): harden untrusted serialized inputs

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/data/energon/base_energon_datamodule.py

@@ -13,14 +13,55 @@
 # limitations under the License.
 
 import logging
-from typing import Any, Literal, Optional
+from typing import Any, Callable, Literal, Optional
 
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.energon import WorkerConfig, get_savable_loader, get_train_dataset
+from megatron.energon.epathlib import EPath
+from megatron.energon.flavors.webdataset.default_generic_webdataset import DefaultGenericWebdatasetFactory
 
 
 logger = logging.getLogger(__name__)
 
+_energon_factory_init = DefaultGenericWebdatasetFactory.__init__
+
+
+def _secure_energon_factory_init(
+    self: DefaultGenericWebdatasetFactory,
+    path: EPath,
+    *,
+    subflavors: dict[str, Any] | None = None,
+    field_map: dict[str, str] | None = None,
+    sample_loader: str | Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    part_filter: str | list[str] | Callable[[str], bool] | None = None,
+    **kwargs: Any,
+) -> None:
+    """Reject dataset-local Python hooks before Energon resolves their files."""
+    executable_fields = [
+        name
+        for name, value in (("sample_loader", sample_loader), ("part_filter", part_filter))
+        if isinstance(value, str)
+    ]
+    if executable_fields:
+        raise ValueError(
+            "Energon dataset metadata cannot load Python files through "
+            f"{', '.join(executable_fields)}. Use a declarative field_map instead."
+        )
+    _energon_factory_init(
+        self,
+        path,
+        subflavors=subflavors,
+        field_map=field_map,
+        sample_loader=sample_loader,
+        part_filter=part_filter,
+        **kwargs,
+    )
+
+
+# Energon constructs this factory internally after reading dataset.yaml, so
+# Bridge must install the guard before calling get_train_dataset().
+DefaultGenericWebdatasetFactory.__init__ = _secure_energon_factory_init
+
 
 class EnergonMultiModalDataModule:
     """

src/megatron/bridge/diffusion/data/common/diffusion_task_encoder_with_sp.py

@@ -24,6 +24,7 @@
 from megatron.bridge.data.energon.metadata import sample_metadata_kwargs
 from megatron.bridge.data.packing.algorithms import first_fit_decreasing
 from megatron.bridge.diffusion.data.common.diffusion_sample import DiffusionSample
+from megatron.bridge.diffusion.data.common.safe_decoder import SafeDiffusionSampleDecoder
 from megatron.bridge.diffusion.data.common.sequence_packing_utils import packing_length
 
 
@@ -51,6 +52,8 @@ def cook(sample: dict) -> dict:
 
 
 class DiffusionTaskEncoderWithSequencePacking(DefaultTaskEncoder, ABC):  # noqa: D101
+    decoder = SafeDiffusionSampleDecoder()
+
     cookers = [
         Cooker(cook),
     ]

src/megatron/bridge/diffusion/data/common/safe_decoder.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Safe decoding for Bridge diffusion WebDataset fields."""
+
+import io
+import re
+from typing import Any
+
+import torch
+from megatron.energon.flavors.webdataset.sample_decoder import SampleDecoder
+from webdataset.autodecode import basichandlers
+
+from megatron.bridge.utils.safe_pickle import safe_torch_tensor_pickle_loads
+
+
+_PICKLE_EXTENSIONS = frozenset({"pickle", "pkl", "pyd"})
+
+
+def _safe_basic_handler(key: str, data: bytes) -> Any:
+    """Decode tensor fields safely and delegate non-executable formats."""
+    extension = re.sub(r".*[.]", "", key).lower()
+    if extension in _PICKLE_EXTENSIONS:
+        return safe_torch_tensor_pickle_loads(data)
+    if extension == "pth":
+        return torch.load(io.BytesIO(data), map_location="cpu", weights_only=True)
+    return basichandlers(key, data)
+
+
+class SafeDiffusionSampleDecoder(SampleDecoder):
+    """Energon decoder that never uses WebDataset's unrestricted pickle handlers."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        replaced = False
+        for index, handler in enumerate(self._decoder.handlers):
+            if handler is basichandlers:
+                self._decoder.handlers[index] = _safe_basic_handler
+                replaced = True
+        if not replaced:
+            raise RuntimeError("Energon SampleDecoder no longer exposes the expected WebDataset basic handler.")

src/megatron/bridge/utils/instantiate_utils.py

@@ -105,6 +105,15 @@
     },
 }
 
+# Resolve aliases before enforcing these entries. For example,
+# ``numpy.lib.npyio.load`` is the same callable as ``numpy.load`` and
+# ``torch.serialization.load`` is the same callable as ``torch.load``.
+_DISALLOWED_CANONICAL_TARGETS: set[str] = {
+    "numpy.load",
+    "torch.serialization.load",
+    "transformers.dynamic_module_utils.get_class_from_dynamic_module",
+}
+
 _DISALLOWED_CALLABLE_FIELD_NAMES: set[str] = {
     "collate_impl",
     "hf_filter_lambda",
@@ -182,7 +191,17 @@ def _resolve_target(
     """Resolve target string, type, or callable after Bridge validation."""
     if isinstance(target, str):
         _validate_target_prefix(target=target, full_key=full_key)
-    return _mcore_resolve_target(target, full_key, check_callable)
+    resolved_target = _mcore_resolve_target(target, full_key, check_callable)
+    if isinstance(target, str):
+        module = getattr(resolved_target, "__module__", None)
+        qualname = getattr(resolved_target, "__qualname__", None)
+        canonical_target = f"{module}.{qualname}" if isinstance(module, str) and isinstance(qualname, str) else None
+        if canonical_target in _DISALLOWED_CANONICAL_TARGETS:
+            raise InstantiationException(
+                f"Instantiation of '{target}' is not allowed because it resolves to the unsafe target "
+                f"'{canonical_target}'." + (f"\nfull_key: {full_key}" if full_key else "")
+            )
+    return resolved_target
 
 
 _mcore_instantiate_utils._resolve_target = _resolve_target

src/megatron/bridge/utils/safe_pickle.py

@@ -239,6 +239,37 @@ def find_class(self, module: str, name: str) -> object:
         )
 
 
+def _safe_torch_load_from_bytes(data: bytes) -> object:
+    """Reconstruct tensor storage bytes without enabling arbitrary pickle globals."""
+    import torch
+
+    return torch.load(io.BytesIO(data), map_location="cpu", weights_only=True)
+
+
+class _TorchTensorRestrictedUnpickler(_RestrictedUnpickler):
+    """Restricted unpickler for plain tensors and containers of tensors."""
+
+    _SAFE_MODULES = MappingProxyType(
+        {
+            **_RestrictedUnpickler._SAFE_MODULES,
+            "torch._utils": frozenset({"_rebuild_tensor", "_rebuild_tensor_v2"}),
+        }
+    )
+
+    def find_class(self, module: str, name: str) -> object:
+        if (module, name) in {
+            ("megatron.core.safe_globals", "safe_load_from_bytes"),
+            ("torch.storage", "_load_from_bytes"),
+        }:
+            return _safe_torch_load_from_bytes
+        if module in self._SAFE_MODULES and name in self._SAFE_MODULES[module]:
+            return pickle.Unpickler.find_class(self, module, name)
+        raise pickle.UnpicklingError(
+            f"Restricted unpickler refused to load '{module}.{name}'. "
+            "Only safe built-in types and plain torch tensors are allowed."
+        )
+
+
 class _NumpyRestrictedUnpickler(pickle.Unpickler):
     """Unpickler that allows safe builtins and the narrow set of numpy types needed for object array reconstruction.
 
@@ -428,6 +459,11 @@ def safe_pickle_loads(data: bytes) -> object:
     return _RestrictedUnpickler(io.BytesIO(data)).load()
 
 
+def safe_torch_tensor_pickle_loads(data: bytes) -> object:
+    """Deserialize raw pickle data containing only safe containers and plain torch tensors."""
+    return _TorchTensorRestrictedUnpickler(io.BytesIO(data)).load()
+
+
 def safe_load_npy(data: bytes):
     """Load a ``.npy`` file from raw bytes without enabling unrestricted pickle.
 

tests/unit_tests/data/energon/test_safe_dataset_config.py

@@ -0,0 +1,55 @@
+import pytest
+from megatron.energon.dataset_config import load_config
+from megatron.energon.epathlib import EPath
+from megatron.energon.flavors.webdataset.default_generic_webdataset import DefaultGenericWebdatasetFactory
+
+from megatron.bridge.data.energon import base_energon_datamodule
+
+
+@pytest.mark.parametrize("field", ["sample_loader", "part_filter"])
+def test_dataset_yaml_rejects_python_hooks_before_import(field, tmp_path):
+    marker = tmp_path / "dataset-python-executed"
+    metadata_dir = tmp_path / ".nv-meta"
+    metadata_dir.mkdir()
+    module_path = metadata_dir / "evil.py"
+    module_path.write_text(f"from pathlib import Path\nPath({str(marker)!r}).touch()\n")
+    config_path = metadata_dir / "dataset.yaml"
+    config_path.write_text(
+        "sample_loader: evil.py\npart_filter:\n  - json\n"
+        if field == "sample_loader"
+        else "part_filter: evil.py\n"
+    )
+    default_kwargs = {"path": EPath(tmp_path)}
+    if field == "part_filter":
+        default_kwargs["sample_loader"] = lambda sample: sample
+
+    with pytest.raises(ValueError, match="cannot load Python files"):
+        load_config(
+            EPath(config_path),
+            default_type=DefaultGenericWebdatasetFactory,
+            default_kwargs=default_kwargs,
+        )
+
+    assert not marker.exists()
+
+
+def test_factory_guard_preserves_callable_hooks(monkeypatch, tmp_path):
+    captured = {}
+
+    def original_init(self, path, **kwargs):
+        captured.update(kwargs)
+
+    monkeypatch.setattr(base_energon_datamodule, "_energon_factory_init", original_init)
+
+    def sample_loader(sample):
+        return sample
+
+    def part_filter(_part):
+        return True
+
+    base_energon_datamodule._secure_energon_factory_init(
+        object(), EPath(tmp_path), sample_loader=sample_loader, part_filter=part_filter
+    )
+
+    assert captured["sample_loader"] is sample_loader
+    assert captured["part_filter"] is part_filter

tests/unit_tests/diffusion/data/common/test_safe_decoder.py

@@ -0,0 +1,59 @@
+import io
+import os
+import pickle
+
+import pytest
+import torch
+from webdataset.autodecode import DecodingError
+
+from megatron.bridge.diffusion.data.common.safe_decoder import SafeDiffusionSampleDecoder
+
+
+class _WriteMarkerPayload:
+    def __init__(self, marker: str) -> None:
+        self.marker = marker
+
+    def __reduce__(self):
+        return os.system, (f"touch {self.marker}",)
+
+
+@pytest.mark.parametrize("extension", ["pickle", "pkl", "pyd"])
+def test_safe_decoder_loads_plain_tensor_pickles(extension):
+    decoder = SafeDiffusionSampleDecoder()
+    value = {"embedding": torch.arange(4, dtype=torch.float32), "shape": [1, 4]}
+
+    restored = decoder.decode(f"sample.{extension}", pickle.dumps(value))
+
+    assert restored["shape"] == value["shape"]
+    assert torch.equal(restored["embedding"], value["embedding"])
+
+
+def test_safe_decoder_loads_weights_only_pth():
+    decoder = SafeDiffusionSampleDecoder()
+    value = {"embedding": torch.arange(4, dtype=torch.float32), "shape": [1, 4]}
+    buffer = io.BytesIO()
+    torch.save(value, buffer)
+
+    restored = decoder.decode("sample.pth", buffer.getvalue())
+
+    assert restored["shape"] == value["shape"]
+    assert torch.equal(restored["embedding"], value["embedding"])
+
+
+@pytest.mark.parametrize("extension", ["pickle", "pth"])
+def test_safe_decoder_rejects_executable_payloads(extension, tmp_path):
+    decoder = SafeDiffusionSampleDecoder()
+    marker = tmp_path / f"{extension}-executed"
+    payload = _WriteMarkerPayload(str(marker))
+    if extension == "pickle":
+        data = pickle.dumps(payload)
+    else:
+        buffer = io.BytesIO()
+        torch.save(payload, buffer)
+        data = buffer.getvalue()
+
+    with pytest.raises(DecodingError) as error:
+        decoder.decode(f"sample.{extension}", data)
+
+    assert isinstance(error.value.__cause__, pickle.UnpicklingError)
+    assert not marker.exists()

tests/unit_tests/utils/test_instantiate_utils.py

@@ -810,6 +810,29 @@ def test_instantiate_rejects_deserialization_and_native_loader_targets(self, tar
         with pytest.raises(InstantiationException, match="bypass target validation"):
             instantiate(config)
 
+    @pytest.mark.parametrize(
+        "target",
+        [
+            "numpy.lib.npyio.load",
+            "torch.serialization.load",
+            "transformers.dynamic_module_utils.get_class_from_dynamic_module",
+        ],
+    )
+    def test_instantiate_rejects_canonical_aliases_and_dynamic_code_helpers(self, target):
+        """Test that aliases and dynamic-code helpers cannot bypass exact target checks."""
+        with pytest.raises(InstantiationException, match="resolves to the unsafe target"):
+            instantiate({"_target_": target, "_call_": False})
+
+    def test_instantiate_rejects_unsafe_alias_hidden_in_unused_kwarg(self):
+        """Test that recursive instantiation rejects unsafe targets before filtering kwargs."""
+        config = {
+            "_target_": "tests.unit_tests.utils.test_instantiate_utils.TestClass",
+            "name": "safe",
+            "unused": {"_target_": "numpy.lib.npyio.load", "_call_": False},
+        }
+        with pytest.raises(InstantiationException, match="resolves to the unsafe target"):
+            instantiate(config)
+
     @pytest.mark.parametrize(
         "target",
         [