Harden nnUNetV2Runner against code execution

monai/apps/nnunet/nnunetv2_runner.py

@@ -17,6 +17,7 @@
 import re
 import shlex
 import subprocess
+import warnings
 from typing import Any
 
 import monai
@@ -741,17 +742,19 @@ def train_parallel(
                     f"log '.txt' inside '{os.path.join(self.nnunet_results, self.dataset_name)}'"
                 )
         for stage in all_cmds:
-            processes = []
-            for device_id in stage:
-                if not stage[device_id]:
-                    continue
-                cmd_str = "; ".join(shlex.join(cmd) for cmd, _ in stage[device_id])
-                env = stage[device_id][0][1]
-                logger.info(f"Current running command on GPU device {device_id}:\n{cmd_str}\n")
-                processes.append(subprocess.Popen(cmd_str, shell=True, env=env, stdout=subprocess.DEVNULL))
-            # finish this stage first
-            for p in processes:
-                p.wait()
+            max_cmds_per_gpu = max((len(cmds) for cmds in stage.values()), default=0)
+            for cmd_index in range(max_cmds_per_gpu):
+                processes = []
+                for device_id, gpu_cmds in stage.items():
+                    if cmd_index >= len(gpu_cmds):
+                        continue
+                    cmd, env = gpu_cmds[cmd_index]
+                    cmd_str = shlex.join(cmd)
+                    logger.info(f"Current running command on GPU device {device_id}:\n{cmd_str}\n")
+                    processes.append(subprocess.Popen(cmd, shell=False, env=env, stdout=subprocess.DEVNULL))
+                # finish this command round on all GPUs before starting the next command
+                for p in processes:
+                    p.wait()
 
     def validate_single_model(self, config: str, fold: int, **kwargs: Any) -> None:
         """
@@ -996,7 +999,16 @@ def predict_ensemble_postprocessing(
 
         # apply postprocessing
         if run_postprocessing:
-            pp_fns, pp_fn_kwargs = load_pickle(self.best_configuration["best_model_or_ensemble"]["postprocessing_file"])
+            postprocessing_file = self.best_configuration["best_model_or_ensemble"]["postprocessing_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 "
+                "(see https://github.com/Project-MONAI/MONAI/security/advisories/GHSA-8f32-8649-rv87).",
+                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"),

tests/apps/nnunet/test_nnunetv2_runner_command.py

@@ -11,10 +11,16 @@
 
 from __future__ import annotations
 
+import os
+import sys
+import types
 import unittest
+import warnings
 from unittest import mock
 
+from monai.apps.nnunet import nnunetv2_runner
 from monai.apps.nnunet.nnunetv2_runner import nnUNetV2Runner
+from monai.bundle import ConfigParser
 
 
 def _make_runner(export_validation_probabilities=False):
@@ -74,5 +80,74 @@ def test_validate_emits_bare_val_flag(self):
         self.assertNotIn("True", cmd)
 
 
+class TestTrainParallelCommand(unittest.TestCase):
+    def test_train_parallel_uses_argv_list_without_shell(self):
+        runner = _make_runner()
+        runner.dataset_name = "Dataset001_Test"
+        runner.nnunet_results = "/tmp/nnunet_results"
+
+        all_cmds = [
+            {
+                0: [
+                    (["python", "-m", "train", "--fold", "0"], {"CUDA_VISIBLE_DEVICES": "0"}),
+                    (["python", "-m", "train", "--fold", "1"], {"CUDA_VISIBLE_DEVICES": "0"}),
+                ],
+                1: [(["python", "-m", "train", "--fold", "2"], {"CUDA_VISIBLE_DEVICES": "1"})],
+            }
+        ]
+
+        with mock.patch.object(runner, "train_parallel_cmd", return_value=all_cmds):
+            with mock.patch("monai.apps.nnunet.nnunetv2_runner.subprocess.Popen") as popen:
+                popen.return_value.wait.return_value = None
+                runner.train_parallel()
+
+        self.assertEqual(popen.call_count, 3)
+        for call in popen.call_args_list:
+            self.assertIsInstance(call.args[0], list)
+            self.assertFalse(call.kwargs["shell"])
+
+
+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",
+            }
+        }
+
+        ensemble_mod = types.ModuleType("nnunetv2.ensembling.ensemble")
+        ensemble_mod.ensemble_folders = mock.MagicMock()
+        pp_mod = types.ModuleType("nnunetv2.postprocessing.remove_connected_components")
+        pp_mod.apply_postprocessing_to_folder = mock.MagicMock()
+        fp_mod = types.ModuleType("nnunetv2.utilities.file_path_utilities")
+        fp_mod.get_output_folder = mock.MagicMock(return_value="/tmp/model_folder")
+
+        fake_modules = {
+            "nnunetv2.ensembling.ensemble": ensemble_mod,
+            "nnunetv2.postprocessing.remove_connected_components": pp_mod,
+            "nnunetv2.utilities.file_path_utilities": fp_mod,
+        }
+
+        load_pickle = mock.MagicMock(return_value=([], {}))
+        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):
+                    with mock.patch.object(nnunetv2_runner, "load_pickle", load_pickle):
+                        with warnings.catch_warnings(record=True) as caught:
+                            warnings.simplefilter("always")
+                            runner.predict_ensemble_postprocessing(
+                                run_predict=False, run_ensemble=False, run_postprocessing=True
+                            )
+
+        load_pickle.assert_called_once_with("/tmp/attacker_controlled_postprocessing.pkl")
+        self.assertTrue(any("unpickling postprocessing_file" in str(item.message) for item in caught))
+
+
 if __name__ == "__main__":
     unittest.main()