Treat BundleGen algos mapping values as literal data

monai/apps/auto3dseg/bundle_gen.py

@@ -35,6 +35,7 @@
     _prepare_cmd_bcprun,
     _prepare_cmd_default,
     _prepare_cmd_torchrun,
+    _reject_non_algo_target,
     _run_cmd_bcprun,
     _run_cmd_torchrun,
     algo_to_json,
@@ -531,13 +532,66 @@ def _copy_algos_folder(folder, at_path):
     return algos_all
 
 
+def _algo_from_params(algo_name: Any, algo_params: Any) -> Algo:
+    """Instantiate one algorithm from an entry of the ``algos`` mapping, using literal values only.
+
+    The mapping is user input: ``AutoRunner`` reads it from ``input.yaml`` and ``BundleGen``/its CLI accept it
+    directly. Its values are therefore plain data, and they are handed to the ``Algo`` subclass as literals.
+    They are never passed through ``ConfigParser``, which would treat strings starting with ``$``, ``@`` or
+    ``%`` as configuration directives and evaluate them.
+
+    The class named by ``_target_`` is still imported, from ``template_path`` if the module is not installed,
+    so ``template_path`` must point to a directory you trust: importing a module runs its top-level code.
+
+    Args:
+        algo_name: the key of the entry, used in error messages.
+        algo_params: the mapping with a ``_target_`` key naming an ``Algo`` subclass and optional keyword
+            arguments for its constructor, such as ``template_path``.
+
+    Returns:
+        the instantiated ``Algo``.
+
+    Raises:
+        ValueError: if the entry is not a mapping, has no ``_target_``, or the target is not an ``Algo``
+            subclass.
+    """
+    if not isinstance(algo_params, dict):
+        raise ValueError(
+            f"Invalid algos entry {algo_name!r}: expected a dict with a '_target_' key, "
+            f"got {type(algo_params).__name__}."
+        )
+
+    params = dict(algo_params)
+    target = params.pop("_target_", None)
+    if target is None:
+        raise ValueError(f"Invalid algos entry {algo_name!r}: missing '_target_' key.")
+
+    template_path = params.get("template_path", ".")
+    if not isinstance(template_path, str):
+        raise ValueError(
+            f"Invalid template_path {template_path!r} for algos entry {algo_name!r}: expected a string, "
+            f"got {type(template_path).__name__}."
+        )
+    if template_path and template_path not in sys.path:
+        sys.path.append(template_path)
+
+    # the allowlist resolves the class and checks issubclass(..., Algo) before anything is instantiated
+    algo_cls: Any = _reject_non_algo_target(target, f"algos entry {algo_name!r}")
+    return algo_cls(**params)
+
+
 class BundleGen(AlgoGen):
     """
     This class generates a set of bundles according to the cross-validation folds, each of them can run independently.
 
     Args:
         algo_path: the directory path to save the algorithm templates. Default is the current working dir.
-        algos: If dictionary, it outlines the algorithm to use. If a list or a string, defines a subset of names of
+        algos: If dictionary, it outlines the algorithm to use. The values are treated as plain data: each entry
+            needs a ``_target_`` key naming an ``Algo`` subclass (a full module path or a monai component name)
+            and may pass further keyword arguments for its constructor (for example ``template_path``).
+            Configuration expressions (``$``, ``@``, ``%``) are not evaluated. The target class is imported
+            from ``template_path`` when the module is not installed, so that directory must be trusted.
+            If a list or a string, defines a subset of names of
             the algorithms to use, e.g. ('segresnet', 'dints') out of the full set of algorithm templates provided
             by templates_path_or_url. Defaults to None - to use all available algorithms.
         templates_path_or_url: the folder with the algorithm templates or a url. If None provided, the default template
@@ -592,12 +646,8 @@ def __init__(
         self.algos: Any = []
         if isinstance(algos, dict):
             for algo_name, algo_params in sorted(algos.items()):
-                template_path = algo_params.get("template_path", ".")
-                if len(template_path) > 0 and template_path not in sys.path:
-                    sys.path.append(template_path)
-
                 try:
-                    onealgo = ConfigParser(algo_params).get_parsed_content()
+                    onealgo: Any = _algo_from_params(algo_name, algo_params)
                     onealgo.name = algo_name
                     self.algos.append(onealgo)
                 except RuntimeError as e:

monai/auto3dseg/utils.py

@@ -27,6 +27,7 @@
 import torch
 
 from monai.auto3dseg import Algo
+from monai.bundle.config_item import ComponentLocator
 from monai.bundle.config_parser import ConfigParser
 from monai.bundle.utils import ID_SEP_KEY
 from monai.config import PathLike
@@ -63,6 +64,17 @@ def _require_pickle_allowed() -> None:
         raise RuntimeError(_PICKLE_DISABLED_MSG)
 
 
+_COMPONENT_LOCATOR: ComponentLocator | None = None
+
+
+def _component_locator() -> ComponentLocator:
+    """Return the process-wide locator used to expand short component names into module paths."""
+    global _COMPONENT_LOCATOR
+    if _COMPONENT_LOCATOR is None:
+        _COMPONENT_LOCATOR = ComponentLocator()
+    return _COMPONENT_LOCATOR
+
+
 def _reject_non_algo_target(target: Any, filename: str) -> type[Algo]:
     """
     Require that ``target`` names an ``Algo`` subclass before it is instantiated.
@@ -72,6 +84,12 @@ def _reject_non_algo_target(target: Any, filename: str) -> type[Algo]:
     resolving the name and checking it here rejects payloads such as
     ``subprocess.call`` or ``os.system`` before they are ever invoked.
 
+    A ``_target_`` may name a monai component without its module path, e.g. ``"BundleAlgo"``.
+    ``ComponentLocator`` maps such a name to the module that defines it, only for classes and
+    functions defined in ``monai.*`` modules, so expanding the name does not widen the allowlist.
+    Bare names go through the locator instead of ``locate``, so a module of the same name on
+    ``sys.path`` (for example a file in ``template_path``) cannot shadow the real component.
+
     Importing the named module still runs that module's top-level code, so this narrows the sink
     rather than closing it: an untrusted ``algo_object.json`` combined with an attacker-controlled
     template directory remains dangerous.
@@ -84,13 +102,26 @@ def _reject_non_algo_target(target: Any, filename: str) -> type[Algo]:
         the resolved ``Algo`` subclass.
 
     Raises:
-        ValueError: if ``target`` is not a string, or does not resolve to a concrete ``Algo`` subclass.
+        ValueError: if ``target`` is not a string, does not resolve to a concrete ``Algo`` subclass, or is
+            a short name that more than one monai module defines.
         ModuleNotFoundError: if the module cannot be imported, so callers can try the next
             template path.
     """
     if not isinstance(target, str):
         raise ValueError(f"Invalid `_target_` in {filename}: expected a string, got {type(target).__name__}.")
-    resolved = locate(target)
+    if "." in target:
+        resolved = locate(target)
+    else:
+        # a bare component name, e.g. "BundleAlgo" instead of "monai.apps.auto3dseg.BundleAlgo".
+        # It is resolved through ComponentLocator rather than ``locate`` so that a module of the same
+        # name on sys.path, for example one in template_path, cannot shadow it.
+        module = _component_locator().get_component_module_name(target)
+        if isinstance(module, list):
+            raise ValueError(
+                f"Multiple components have name `{target}`: {module}. "
+                f"Set its full module path in `_target_` to specify which one is used."
+            )
+        resolved = locate(f"{module}.{target}") if module is not None else None
     if resolved is None:
         raise ModuleNotFoundError(f"Cannot resolve `_target_` {target!r} from {filename}.")
     if not (isinstance(resolved, type) and resolved is not Algo and issubclass(resolved, Algo)):

tests/auto3dseg/test_algo_target_allowlist.py

@@ -13,8 +13,10 @@
 
 import json
 import os
+import sys
 import tempfile
 import unittest
+from unittest import mock
 
 from monai.auto3dseg.utils import _reject_non_algo_target, algo_from_json
 
@@ -48,6 +50,32 @@ def test_algo_subclass_is_accepted(self):
         """A real Algo subclass passes the check."""
         _reject_non_algo_target("monai.apps.auto3dseg.BundleAlgo", "algo_object.json")
 
+    def test_short_component_name_is_resolved(self):
+        """A bare monai component name resolves to the module that defines it, as ConfigParser does."""
+        from monai.apps.auto3dseg import BundleAlgo
+
+        self.assertIs(_reject_non_algo_target("BundleAlgo", "algo_object.json"), BundleAlgo)
+
+    def test_bare_base_class_name_is_rejected(self):
+        """Resolving a short name does not admit the abstract Algo base class."""
+        with self.assertRaisesRegex(ValueError, r"GHSA-2wx3-8x3w-r8qv"):
+            _reject_non_algo_target("Algo", "algo_object.json")
+
+    def test_bare_name_is_not_shadowed_by_a_module_on_the_path(self):
+        """A module file named like a monai component does not take precedence over the component."""
+        with (
+            tempfile.TemporaryDirectory() as tempdir,
+            mock.patch("sys.path", [tempdir, *sys.path]),
+            mock.patch.dict(sys.modules),
+        ):
+            marker = os.path.join(tempdir, "SHADOWED")
+            with open(os.path.join(tempdir, "BundleAlgo.py"), "w") as f:
+                f.write(f"open({marker!r}, 'w').close()\n")
+            from monai.apps.auto3dseg import BundleAlgo
+
+            self.assertIs(_reject_non_algo_target("BundleAlgo", "algo_object.json"), BundleAlgo)
+            self.assertFalse(os.path.exists(marker), "the same-named module in template_path was imported")
+
     def test_non_string_target_is_rejected(self):
         """A JSON ``null`` ``_target_`` raises the documented ``ValueError``, not ``AttributeError``."""
         with self.assertRaisesRegex(ValueError, r"expected a string"):

tests/auto3dseg/test_bundle_gen_algos_literals.py

@@ -0,0 +1,145 @@
+# Copyright (c) MONAI Consortium
+# 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.
+
+from __future__ import annotations
+
+import os
+import sys
+import tempfile
+import textwrap
+import unittest
+import uuid
+
+from monai.apps.auto3dseg.bundle_gen import BundleGen
+
+_ALGO_TEMPLATE = """
+from monai.auto3dseg.algo_gen import Algo
+
+
+class LiteralAlgo(Algo):
+    def __init__(self, template_path=None, **kwargs):
+        self.template_path = template_path
+        self.extra = kwargs
+"""
+
+
+class TestBundleGenAlgosLiterals(unittest.TestCase):
+    """Regression tests for the ``algos`` mapping of ``BundleGen``.
+
+    The mapping is user input: ``AutoRunner`` reads it from ``input.yaml`` and ``BundleGen``/its CLI accept it
+    directly. Its values are plain data, so they must reach the ``Algo`` constructor verbatim instead of being
+    parsed by ``ConfigParser``, which evaluates strings starting with ``$``, ``@`` or ``%``.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.tempdir = tempfile.TemporaryDirectory()
+        cls.root = cls.tempdir.name
+        # A uniquely named package so repeated runs and other test modules never share a module object.
+        cls.package = f"literal_algos_{uuid.uuid4().hex}"
+        package_dir = os.path.join(cls.root, cls.package)
+        os.makedirs(os.path.join(package_dir, "scripts"))
+        for init in (os.path.join(package_dir, "__init__.py"), os.path.join(package_dir, "scripts", "__init__.py")):
+            with open(init, "w") as f:
+                f.write("")
+        with open(os.path.join(package_dir, "scripts", "algo.py"), "w") as f:
+            f.write(textwrap.dedent(_ALGO_TEMPLATE))
+        cls.target = f"{cls.package}.scripts.algo.LiteralAlgo"
+        cls.marker = os.path.join(cls.root, "PWNED")
+
+    @classmethod
+    def tearDownClass(cls):
+        for name in list(sys.modules):
+            if name == cls.package or name.startswith(f"{cls.package}."):
+                del sys.modules[name]
+        if cls.root in sys.path:
+            sys.path.remove(cls.root)
+        cls.tempdir.cleanup()
+
+    def setUp(self):
+        if os.path.exists(self.marker):
+            os.unlink(self.marker)
+        # Start from a clean import state so each test is independent of the ones before it.
+        if self.root in sys.path:
+            sys.path.remove(self.root)
+        for name in list(sys.modules):
+            if name == self.package or name.startswith(f"{self.package}."):
+                del sys.modules[name]
+
+    def _payload(self, value: str = "pwned") -> str:
+        return f"$__import__('pathlib').Path({self.marker!r}).write_text({value!r})"
+
+    def test_mapping_values_reach_the_constructor_literally(self):
+        """A ``$`` value in the mapping is passed through as a string and never evaluated."""
+        payload = self._payload()
+        gen = BundleGen(algos={"algo1": {"_target_": self.target, "template_path": self.root, "note": payload}})
+
+        self.assertFalse(os.path.exists(self.marker), "the algos mapping payload was evaluated")
+        self.assertEqual(len(gen.algos), 1)
+        algo = gen.algos[0]
+        self.assertEqual(type(algo).__name__, "LiteralAlgo")
+        self.assertEqual(algo.name, "algo1")
+        self.assertEqual(algo.template_path, self.root)
+        self.assertEqual(algo.extra["note"], payload)
+
+    def test_nested_config_directive_is_kept_as_plain_data(self):
+        """A nested ``_target_`` mapping is not resolved; it stays a dict and nothing is instantiated from it."""
+        nested = {"_target_": "os.system", "args": [f"touch {self.marker}"]}
+        gen = BundleGen(algos={"algo1": {"_target_": self.target, "template_path": self.root, "note": nested}})
+
+        self.assertFalse(os.path.exists(self.marker), "a nested _target_ in the algos mapping was resolved")
+        self.assertEqual(gen.algos[0].extra["note"], nested)
+
+    def test_reference_and_file_strings_are_literal(self):
+        """``@`` and ``%`` strings in the mapping reach the constructor as plain strings."""
+        for value in ("@some_id", "%other.json#key"):
+            with self.subTest(value=value):
+                gen = BundleGen(algos={"algo1": {"_target_": self.target, "template_path": self.root, "note": value}})
+                self.assertEqual(gen.algos[0].extra["note"], value)
+
+    def test_short_target_name_is_resolved(self):
+        """A bare monai component name, as bundle syntax allows, still instantiates the class."""
+        gen = BundleGen(algos={"algo1": {"_target_": "BundleAlgo", "template_path": self.root}})
+
+        self.assertEqual(len(gen.algos), 1)
+        self.assertEqual(type(gen.algos[0]).__name__, "BundleAlgo")
+        self.assertEqual(gen.algos[0].template_path, self.root)
+
+    def test_expression_as_template_path_is_not_evaluated(self):
+        """A ``$`` expression used as ``template_path`` is not evaluated, even though the module then fails."""
+        payload = self._payload()
+        with self.assertRaises(ModuleNotFoundError):
+            BundleGen(algos={"algo1": {"_target_": self.target, "template_path": payload}})
+        self.assertFalse(os.path.exists(self.marker), "the template_path payload was evaluated")
+
+    def test_non_algo_target_is_rejected(self):
+        """``_target_`` must name an ``Algo`` subclass, not an arbitrary callable."""
+        with self.assertRaisesRegex(ValueError, r"GHSA-2wx3-8x3w-r8qv"):
+            BundleGen(algos={"algo1": {"_target_": "os.system", "args": [f"touch {self.marker}"]}})
+        self.assertFalse(os.path.exists(self.marker), "a non-Algo _target_ was instantiated")
+
+    def test_string_entry_raises_value_error(self):
+        """The documented mapping shape is a dict per algorithm; a bare string is rejected cleanly."""
+        with self.assertRaisesRegex(ValueError, r"expected a dict"):
+            BundleGen(algos={"algo1": self._payload()})
+        self.assertFalse(os.path.exists(self.marker))
+
+    def test_missing_target_is_rejected(self):
+        with self.assertRaisesRegex(ValueError, r"missing '_target_' key"):
+            BundleGen(algos={"algo1": {"template_path": self.root}})
+
+    def test_non_string_template_path_is_rejected(self):
+        with self.assertRaisesRegex(ValueError, r"expected a string"):
+            BundleGen(algos={"algo1": {"_target_": self.target, "template_path": ["not", "a", "path"]}})
+
+
+if __name__ == "__main__":
+    unittest.main()