diff --git a/monai/apps/auto3dseg/bundle_gen.py b/monai/apps/auto3dseg/bundle_gen.py index 680de6c99f..8a6d92eab0 100644 --- a/monai/apps/auto3dseg/bundle_gen.py +++ b/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: diff --git a/monai/auto3dseg/utils.py b/monai/auto3dseg/utils.py index f6c656947e..d7b3841394 100644 --- a/monai/auto3dseg/utils.py +++ b/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)): diff --git a/tests/auto3dseg/test_algo_target_allowlist.py b/tests/auto3dseg/test_algo_target_allowlist.py index 7ddddd35f4..e153136c0d 100644 --- a/tests/auto3dseg/test_algo_target_allowlist.py +++ b/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"): diff --git a/tests/auto3dseg/test_bundle_gen_algos_literals.py b/tests/auto3dseg/test_bundle_gen_algos_literals.py new file mode 100644 index 0000000000..fbef46638f --- /dev/null +++ b/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()