fix: confine nnUNet postprocessing pickle loads to the results directory (GHSA-8f32)

monai/apps/nnunet/nnunetv2_runner.py

@@ -40,6 +40,38 @@
 DATASET_ID_FORMAT = r"Dataset[0-9]{3}|[0-9]+"  # regex format for a valid nnUnet dataset name
 
 
+def _confine_to_dir(candidate: str, allowed_dir: str, description: str) -> str:
+    """
+    Resolve ``candidate`` and require that it stays inside ``allowed_dir``.
+
+    ``inference_information.json`` is produced by nnU-Net inside the results directory, but its
+    contents are plain JSON: whoever can write that file chooses the paths this runner then loads.
+    Confining the path means a tampered ``inference_information.json`` cannot redirect a load to an
+    attacker-planted file elsewhere on the filesystem.
+
+    Args:
+        candidate: the path read from the configuration file.
+        allowed_dir: the directory the path must resolve inside.
+        description: name of the field, used in the error message.
+
+    Returns:
+        The resolved absolute path.
+
+    Raises:
+        ValueError: if the resolved path escapes ``allowed_dir``.
+    """
+    resolved = os.path.realpath(candidate)
+    allowed_root = os.path.realpath(allowed_dir)
+    if os.path.commonpath([resolved, allowed_root]) != allowed_root:
+        raise ValueError(
+            f"refusing to load {description} from '{candidate}': it resolves to '{resolved}', outside the "
+            f"expected directory '{allowed_root}'. This path is read from inference_information.json; a value "
+            "pointing outside the results directory indicates that file has been tampered with "
+            "(see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-8f32-8649-rv87)."
+        )
+    return resolved
+
+
 class nnUNetV2Runner:  # noqa: N801
     """
     ``nnUNetV2Runner`` provides an interface in MONAI to use `nnU-Net` V2 library to analyze, train, and evaluate
@@ -1004,22 +1036,37 @@ def predict_ensemble_postprocessing(
 
         # apply postprocessing
         if run_postprocessing:
-            postprocessing_file = self.best_configuration["best_model_or_ensemble"]["postprocessing_file"]
+            results_root = os.path.join(self.nnunet_results, self.dataset_name)
+            postprocessing_file = _confine_to_dir(
+                self.best_configuration["best_model_or_ensemble"]["postprocessing_file"],
+                results_root,
+                "postprocessing_file",
+            )
+            plans_file = _confine_to_dir(
+                self.best_configuration["best_model_or_ensemble"]["some_plans_file"], results_root, "some_plans_file"
+            )
             warnings.warn(
                 f"unpickling postprocessing_file {postprocessing_file}: this path is read from "
                 "inference_information.json and is loaded with Python pickle without any allow list, "
-                "which gives whoever controls that file arbitrary code execution. Only proceed if the "
-                "inference_information.json is from a source you trust "
+                "which gives whoever controls that file arbitrary code execution. The path is confined to "
+                "the dataset's results directory, but anyone able to write inside that directory can still "
+                "supply a malicious pickle. Only proceed if the results directory is from a source you trust "
                 "(see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-8f32-8649-rv87).",
                 stacklevel=2,
             )
+            warnings.warn(
+                "loading nnU-Net postprocessing via Python pickle will require the environment variable "
+                "MONAI_ALLOW_PICKLE=1 from MONAI 1.7. Set it now to keep this call working after the change.",
+                FutureWarning,
+                stacklevel=2,
+            )
             pp_fns, pp_fn_kwargs = load_pickle(postprocessing_file)
             apply_postprocessing_to_folder(
                 folder_for_pp,
                 join(target_dir_base, "ensemble_predictions_postprocessed"),
                 pp_fns,
                 pp_fn_kwargs,
-                plans_file_or_dict=self.best_configuration["best_model_or_ensemble"]["some_plans_file"],
+                plans_file_or_dict=plans_file,
             )
 
     def run(

tests/apps/nnunet/test_nnunetv2_runner_command.py

@@ -13,6 +13,7 @@
 
 import os
 import sys
+import tempfile
 import threading
 import types
 import unittest
@@ -148,19 +149,14 @@ def _fake_popen(cmd, *args, **kwargs):
 
 
 class TestPredictEnsemblePostprocessingWarnings(unittest.TestCase):
-    def test_postprocessing_pickle_warns_on_untrusted_file(self):
-        runner = _make_runner()
-        runner.dataset_name = "Dataset001_Test"
-        runner.nnunet_raw = "/tmp/nnunet_raw"
-        runner.nnunet_results = "/tmp/nnunet_results"
-        runner.best_configuration = {
-            "best_model_or_ensemble": {
-                "selected_model_or_models": [{"configuration": "3d_fullres"}],
-                "postprocessing_file": "/tmp/attacker_controlled_postprocessing.pkl",
-                "some_plans_file": "/tmp/plans.json",
-            }
-        }
-
+    def _run_postprocessing(self, runner, events, load_pickle):
+        """Drive ``predict_ensemble_postprocessing`` with nnU-Net's modules stubbed out.
+
+        Args:
+            runner: the runner under test.
+            events: list that records ``warn``/``load_pickle`` calls in order.
+            load_pickle: the mock standing in for ``load_pickle``.
+        """
         ensemble_mod = types.ModuleType("nnunetv2.ensembling.ensemble")
         ensemble_mod.ensemble_folders = mock.MagicMock()
         pp_mod = types.ModuleType("nnunetv2.postprocessing.remove_connected_components")
@@ -174,20 +170,6 @@ def test_postprocessing_pickle_warns_on_untrusted_file(self):
             "nnunetv2.utilities.file_path_utilities": fp_mod,
         }
 
-        events = []
-
-        def _load_pickle(path):
-            """Record a ``load_pickle`` call and return an empty postprocessing pipeline.
-
-            Args:
-                path: path to the pickle file (unused).
-
-            Returns:
-                A tuple of ``(postprocessing_fns, postprocessing_kwargs)``.
-            """
-            events.append("load_pickle")
-            return [], {}
-
         def _warn(*args, **kwargs):
             """Record a ``warnings.warn`` call.
 
@@ -197,7 +179,6 @@ def _warn(*args, **kwargs):
             """
             events.append("warn")
 
-        load_pickle = mock.MagicMock(side_effect=_load_pickle)
         with mock.patch.dict(sys.modules, fake_modules):
             with mock.patch.object(ConfigParser, "load_config_file", return_value=runner.best_configuration):
                 with mock.patch.object(nnunetv2_runner, "join", os.path.join):
@@ -207,8 +188,137 @@ def _warn(*args, **kwargs):
                                 run_predict=False, run_ensemble=False, run_postprocessing=True
                             )
 
-        load_pickle.assert_called_once_with("/tmp/attacker_controlled_postprocessing.pkl")
-        self.assertEqual(events, ["warn", "load_pickle"])
+    def test_postprocessing_pickle_warns_on_untrusted_file(self):
+        """A pickle inside the results directory is loaded, but only after warning."""
+        with tempfile.TemporaryDirectory() as tempdir:
+            results_root = os.path.join(tempdir, "Dataset001_Test")
+            os.makedirs(results_root)
+            pp_file = os.path.join(results_root, "postprocessing.pkl")
+            plans_file = os.path.join(results_root, "plans.json")
+            open(pp_file, "w").close()
+            open(plans_file, "w").close()
+
+            runner = _make_runner()
+            runner.dataset_name = "Dataset001_Test"
+            runner.nnunet_raw = "/tmp/nnunet_raw"
+            runner.nnunet_results = tempdir
+            runner.best_configuration = {
+                "best_model_or_ensemble": {
+                    "selected_model_or_models": [{"configuration": "3d_fullres"}],
+                    "postprocessing_file": pp_file,
+                    "some_plans_file": plans_file,
+                }
+            }
+
+            events = []
+
+            def _load_pickle(path):
+                """Record a ``load_pickle`` call and return an empty postprocessing pipeline.
+
+                Args:
+                    path: path to the pickle file (unused).
+
+                Returns:
+                    A tuple of ``(postprocessing_fns, postprocessing_kwargs)``.
+                """
+                events.append("load_pickle")
+                return [], {}
+
+            load_pickle = mock.MagicMock(side_effect=_load_pickle)
+            self._run_postprocessing(runner, events, load_pickle)
+
+            load_pickle.assert_called_once_with(os.path.realpath(pp_file))
+            # Two warnings now precede the load: the trust warning and the MONAI 1.7 FutureWarning.
+            self.assertEqual(events, ["warn", "warn", "load_pickle"])
+
+    def test_postprocessing_file_outside_results_dir_is_rejected(self):
+        """Regression test for GHSA-8f32-8649-rv87.
+
+        ``postprocessing_file`` is read from ``inference_information.json``; a value pointing
+        outside the dataset's results directory means that file has been tampered with, so the
+        pickle must never be opened.
+        """
+        with tempfile.TemporaryDirectory() as tempdir:
+            results_root = os.path.join(tempdir, "results", "Dataset001_Test")
+            os.makedirs(results_root)
+            evil = os.path.join(tempdir, "evil.pkl")
+            open(evil, "w").close()
+
+            runner = _make_runner()
+            runner.dataset_name = "Dataset001_Test"
+            runner.nnunet_raw = "/tmp/nnunet_raw"
+            runner.nnunet_results = os.path.join(tempdir, "results")
+            runner.best_configuration = {
+                "best_model_or_ensemble": {
+                    "selected_model_or_models": [{"configuration": "3d_fullres"}],
+                    "postprocessing_file": evil,
+                    "some_plans_file": os.path.join(results_root, "plans.json"),
+                }
+            }
+
+            events = []
+            load_pickle = mock.MagicMock()
+            with self.assertRaisesRegex(ValueError, r"GHSA-8f32-8649-rv87"):
+                self._run_postprocessing(runner, events, load_pickle)
+            load_pickle.assert_not_called()
+
+    def test_postprocessing_traversal_is_rejected(self):
+        """A ``..`` traversal that escapes the results directory is rejected."""
+        with tempfile.TemporaryDirectory() as tempdir:
+            results_root = os.path.join(tempdir, "results", "Dataset001_Test")
+            os.makedirs(results_root)
+            evil = os.path.join(tempdir, "evil.pkl")
+            open(evil, "w").close()
+
+            runner = _make_runner()
+            runner.dataset_name = "Dataset001_Test"
+            runner.nnunet_raw = "/tmp/nnunet_raw"
+            runner.nnunet_results = os.path.join(tempdir, "results")
+            runner.best_configuration = {
+                "best_model_or_ensemble": {
+                    "selected_model_or_models": [{"configuration": "3d_fullres"}],
+                    "postprocessing_file": os.path.join(results_root, "..", "..", "evil.pkl"),
+                    "some_plans_file": os.path.join(results_root, "plans.json"),
+                }
+            }
+
+            load_pickle = mock.MagicMock()
+            with self.assertRaisesRegex(ValueError, r"GHSA-8f32-8649-rv87"):
+                self._run_postprocessing(runner, [], load_pickle)
+            load_pickle.assert_not_called()
+
+    def test_plans_file_outside_results_dir_is_rejected(self):
+        """Regression test for GHSA-8f32-8649-rv87 (``some_plans_file``).
+
+        ``some_plans_file`` is read from ``inference_information.json`` and handed to
+        ``apply_postprocessing_to_folder``; a value pointing outside the dataset's results
+        directory would let a tamperer name an arbitrary file on disk, so it must be rejected
+        even when ``postprocessing_file`` itself is valid and in-scope.
+        """
+        with tempfile.TemporaryDirectory() as tempdir:
+            results_root = os.path.join(tempdir, "results", "Dataset001_Test")
+            os.makedirs(results_root)
+            pp_file = os.path.join(results_root, "postprocessing.pkl")
+            open(pp_file, "w").close()
+            evil_plans = os.path.join(tempdir, "malicious_plans.json")
+            open(evil_plans, "w").close()
+
+            runner = _make_runner()
+            runner.dataset_name = "Dataset001_Test"
+            runner.nnunet_raw = "/tmp/nnunet_raw"
+            runner.nnunet_results = os.path.join(tempdir, "results")
+            runner.best_configuration = {
+                "best_model_or_ensemble": {
+                    "selected_model_or_models": [{"configuration": "3d_fullres"}],
+                    "postprocessing_file": pp_file,
+                    "some_plans_file": evil_plans,
+                }
+            }
+
+            load_pickle = mock.MagicMock()
+            with self.assertRaisesRegex(ValueError, r"some_plans_file.*GHSA-8f32-8649-rv87"):
+                self._run_postprocessing(runner, [], load_pickle)
+            load_pickle.assert_not_called()
 
 
 if __name__ == "__main__":