Do not evaluate imports when reading an algorithm score

monai/apps/auto3dseg/bundle_gen.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import ast
 import importlib
 import ntpath
 import os
@@ -39,7 +40,9 @@
     _run_cmd_torchrun,
     algo_to_json,
 )
+from monai.bundle.config_item import ConfigExpression
 from monai.bundle.config_parser import ConfigParser
+from monai.bundle.reference_resolver import ReferenceResolver
 from monai.config import PathLike
 from monai.utils import ensure_tuple, look_up_option, run_cmd
 from monai.utils.enums import AlgoKeys
@@ -97,6 +100,103 @@ def _confined_join(root: str, name: str) -> str:
     return path
 
 
+def _resolve_path_value(value: Any, parser: ConfigParser, what: str) -> str:
+    """
+    Resolve a config value that must be a plain path string, without evaluating it with ``eval``.
+
+    The value can come from a work directory this process did not write, so only the shapes the bundle
+    templates use are accepted: a plain string, or a ``$`` expression made of string literals, ``+`` concatenation, ``str(...)`` of a reference, and references to other
+    entries, for example ``$@bundle_root + '/model_fold' + str(@fold)``. Calls other than ``str`` with one
+    positional argument, attribute access, subscripts to anything but the reference table, comprehensions
+    and ``_target_`` mappings are refused.
+
+    Args:
+        value: the config value, as stored in the file.
+        parser: the parser that holds the other entries referenced with ``@``.
+        what: a description of the value, used in the error messages.
+
+    Returns:
+        the resolved path string.
+
+    Raises:
+        ValueError: if ``value`` is not a string, or is not one of the accepted shapes.
+    """
+    if not isinstance(value, str):
+        raise ValueError(f"Invalid {what}: expected a string, got {type(value).__name__}.")
+    if value.startswith("@"):
+        resolved = _resolve_path_reference(value[1:], parser, what)
+        if not isinstance(resolved, str):
+            raise ValueError(f"Invalid {what}: reference {value!r} is not a plain string.")
+        return resolved
+    if not value.startswith(ConfigExpression.prefix):
+        return value
+
+    refs = {
+        ref: _resolve_path_reference(ref, parser, what, allow_integer=True)
+        for ref in ReferenceResolver.match_refs_pattern(value)
+    }
+    expression = ReferenceResolver.update_refs_pattern(value, refs)
+    try:
+        node = ast.parse(expression[len(ConfigExpression.prefix) :], mode="eval").body
+    except SyntaxError as e:
+        raise ValueError(f"Invalid {what} expression: {value!r}.") from e
+    resolved = _path_expression(node, refs, what)
+    if not isinstance(resolved, str):
+        raise ValueError(f"Invalid {what}: the expression does not resolve to a string.")
+    return resolved
+
+
+def _resolve_path_reference(ref: str, parser: ConfigParser, what: str, allow_integer: bool = False) -> str | int:
+    """Return the value of ``@ref``, refusing anything that would need evaluation itself.
+
+    A string is returned as is. An integer is only returned when ``allow_integer`` is set, because the
+    templates use it through ``str()``, as in ``str(@fold)``; it is never concatenated directly.
+    """
+    try:
+        ref_value = parser[ref]
+    except (KeyError, IndexError) as e:
+        raise ValueError(f"Invalid {what}: reference @{ref} does not exist.") from e
+    if isinstance(ref_value, str):
+        if ref_value.startswith(("@", "$")):
+            raise ValueError(f"Invalid {what}: reference @{ref} is not a plain value.")
+        return ref_value
+    if allow_integer and isinstance(ref_value, int) and not isinstance(ref_value, bool):
+        return ref_value
+    raise ValueError(f"Invalid {what}: reference @{ref} is not a plain value.")
+
+
+def _path_expression(node: ast.AST, refs: dict[str, str | int], what: str) -> str | int:
+    """Evaluate the small expression grammar accepted by :func:`_resolve_path_value`."""
+    if isinstance(node, ast.Constant) and isinstance(node.value, str):
+        return node.value
+    if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add):
+        left = _path_expression(node.left, refs, what)
+        right = _path_expression(node.right, refs, what)
+        if isinstance(left, str) and isinstance(right, str):
+            return left + right
+        raise ValueError(f"Invalid {what}: only strings can be concatenated.")
+    if (
+        isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Name)
+        and node.func.id == "str"
+        and len(node.args) == 1
+        and not node.keywords
+    ):
+        return str(_path_expression(node.args[0], refs, what))
+    if (
+        isinstance(node, ast.Subscript)
+        and isinstance(node.value, ast.Name)
+        and node.value.id == ReferenceResolver._vars
+        and isinstance(node.slice, ast.Constant)
+        and isinstance(node.slice.value, str)
+        and node.slice.value in refs
+    ):
+        return refs[node.slice.value]
+    raise ValueError(
+        f"Invalid {what}: only string literals, + concatenation, str() of a reference and @references are allowed."
+    )
+
+
 class BundleAlgo(Algo):
     """
     An algorithm represented by a set of bundle configurations and scripts.
@@ -360,11 +460,29 @@ def train(
     def get_score(self, *args, **kwargs):
         """
         Returns validation scores of the model trained by the current Algo.
+
+        The configuration is read from the algorithm output directory, which may belong to a work
+        directory this process did not create (``import_bundle_algo_history`` reads any work dir).
+        ``ckpt_path`` is therefore resolved as data rather than evaluated: string literals,
+        ``+`` concatenation and ``@references`` such as ``$@bundle_root + '/model_fold' + str(@fold)``
+        work, while any other call, an attribute, a comprehension or a ``_target_`` mapping is refused.
+
+        Args:
+            *args: unused, kept for compatibility with the ``Algo`` interface.
+            **kwargs: unused, kept for compatibility with the ``Algo`` interface.
+
+        Returns:
+            the best average Dice score recorded in ``progress.yaml``.
+
+        Raises:
+            ValueError: if ``ckpt_path`` in the saved configuration is not a plain path string.
         """
         config_yaml = os.path.join(self.output_path, "configs", "hyper_parameters.yaml")
-        parser = ConfigParser()
+        parser = ConfigParser(globals=False)
         parser.read_config(config_yaml)
-        ckpt_path = parser.get_parsed_content("ckpt_path", default=self.output_path)
+        # resolved as data, never with eval(): the file can come from a work directory this process
+        # did not write, and a `$` expression there must not be able to call or import anything
+        ckpt_path = _resolve_path_value(parser.get("ckpt_path", default=self.output_path), parser, "ckpt_path")
 
         dict_file = ConfigParser.load_config_file(os.path.join(ckpt_path, "progress.yaml"))
         # dict_file: a list of scores saved in the form of dict in progress.yaml

monai/apps/auto3dseg/utils.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import logging
 import os
 
 from monai.apps.auto3dseg.bundle_gen import BundleAlgo
@@ -59,8 +60,10 @@ def import_bundle_algo_history(
         if best_metric is None:
             try:
                 best_metric = algo.get_score()
-            except Exception:
-                pass
+            except Exception as e:
+                # a missing score is normal for an untrained algo; anything else means the directory
+                # could not be read as one, which is worth surfacing rather than swallowing
+                logging.warning(f"Could not read the validation score of {name!r} in {output_folder}: {e}")
 
         is_trained = best_metric is not None
 

tests/auto3dseg/test_algo_score_config.py

@@ -0,0 +1,156 @@
+# 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 json
+import os
+import tempfile
+import unittest
+
+from monai.apps.auto3dseg.bundle_gen import BundleAlgo
+from monai.apps.auto3dseg.utils import import_bundle_algo_history
+
+
+class TestAlgoScoreConfig(unittest.TestCase):
+    """Regression tests for the score of an algorithm output directory.
+
+    ``import_bundle_algo_history`` reads any work directory, so the directory and the config inside it
+    can come from someone else. Reading the saved score must not evaluate code from that config.
+    """
+
+    def setUp(self):
+        self.tempdir = tempfile.TemporaryDirectory()
+        self.root = self.tempdir.name
+        self.marker = os.path.join(self.root, "PWNED")
+        self.algo_dir = os.path.join(self.root, "algo1")
+        os.makedirs(os.path.join(self.algo_dir, "configs"))
+
+    def tearDown(self):
+        self.tempdir.cleanup()
+
+    def _write_config(self, **values) -> None:
+        with open(os.path.join(self.algo_dir, "configs", "hyper_parameters.yaml"), "w") as f:
+            for key, value in values.items():
+                f.write(f"{key}: {json.dumps(value)}\n")
+
+    def _write_progress(self, ckpt_dir: str, score: float) -> None:
+        os.makedirs(ckpt_dir, exist_ok=True)
+        with open(os.path.join(ckpt_dir, "progress.yaml"), "w") as f:
+            f.write(f"- best_avg_dice_score: {score}\n")
+
+    def _algo(self) -> BundleAlgo:
+        algo = BundleAlgo(template_path=self.root)
+        algo.output_path = self.algo_dir
+        return algo
+
+    def _assert_refused(self, **values) -> None:
+        self._write_config(**values)
+        with self.assertRaises(ValueError):
+            self._algo().get_score()
+        self.assertFalse(os.path.exists(self.marker), "the ckpt_path payload was executed")
+
+    def test_plain_ckpt_path_still_resolves(self):
+        ckpt_dir = os.path.join(self.algo_dir, "ckpt")
+        self._write_progress(ckpt_dir, 0.75)
+        self._write_config(ckpt_path=ckpt_dir)
+
+        self.assertAlmostEqual(self._algo().get_score(), 0.75)
+
+    def test_string_concatenation_still_resolves(self):
+        ckpt_dir = os.path.join(self.algo_dir, "model_fold_0")
+        self._write_progress(ckpt_dir, 0.42)
+        self._write_config(ckpt_path=f"${self.algo_dir!r} + '/model_fold_0'")
+
+        self.assertAlmostEqual(self._algo().get_score(), 0.42)
+
+    def test_reference_expression_still_resolves(self):
+        """The form used by the shipped templates, ``$@bundle_root + '/model'``, keeps working."""
+        ckpt_dir = os.path.join(self.algo_dir, "model")
+        self._write_progress(ckpt_dir, 0.33)
+        self._write_config(bundle_root=self.algo_dir, ckpt_path="$@bundle_root + '/model'")
+
+        self.assertAlmostEqual(self._algo().get_score(), 0.33)
+
+    def test_integer_reference_through_str_still_resolves(self):
+        """The ``dints`` and ``swinunetr`` templates use ``$@bundle_root + '/model_fold' + str(@fold)``."""
+        ckpt_dir = os.path.join(self.algo_dir, "model_fold0")
+        self._write_progress(ckpt_dir, 0.61)
+        self._write_config(bundle_root=self.algo_dir, fold=0, ckpt_path="$@bundle_root + '/model_fold' + str(@fold)")
+
+        self.assertAlmostEqual(self._algo().get_score(), 0.61)
+
+    def test_integer_reference_cannot_be_concatenated(self):
+        """An integer reference is only usable through str(); concatenating it directly is refused."""
+        self._assert_refused(bundle_root=self.algo_dir, fold=0, ckpt_path="$@bundle_root + @fold")
+
+    def test_missing_reference_is_refused(self):
+        """A reference to an id that is not in the file raises ValueError with the documented message."""
+        self._write_config(ckpt_path="$@no_such_entry + '/model'")
+        with self.assertRaisesRegex(ValueError, "does not exist"):
+            self._algo().get_score()
+
+    def test_unknown_reference_key_is_refused(self):
+        """A subscript that does not come from a matched @reference raises ValueError, not KeyError."""
+        self._assert_refused(ckpt_path="$__local_refs['missing']")
+
+    def test_whole_value_reference_still_resolves(self):
+        ckpt_dir = os.path.join(self.algo_dir, "ckpt")
+        self._write_progress(ckpt_dir, 0.5)
+        self._write_config(ckpt=ckpt_dir, ckpt_path="@ckpt")
+
+        self.assertAlmostEqual(self._algo().get_score(), 0.5)
+
+    def test_expression_in_ckpt_path_does_not_run(self):
+        """A ``$`` expression in ``ckpt_path`` cannot import or call anything."""
+        payload = f"$__import__('pathlib').Path({self.marker!r}).write_text('pwned') or {self.algo_dir!r}"
+
+        self._assert_refused(ckpt_path=payload)
+
+    def test_module_names_are_not_in_scope(self):
+        self._assert_refused(ckpt_path="$str(torch.tensor(1))")
+
+    def test_subclasses_escape_is_refused(self):
+        """The usual ``__subclasses__()`` sandbox escape has no way in: no attribute or call is accepted."""
+        payload = (
+            f"$[c for c in ().__class__.__base__.__subclasses__() if c.__name__ == 'catch_warnings'][0]()"
+            f"._module.__builtins__['__import__']('pathlib').Path({self.marker!r}).write_text('pwned')"
+        )
+
+        self._assert_refused(ckpt_path=payload)
+
+    def test_target_mapping_is_refused(self):
+        """A ``_target_`` mapping is not instantiated: ``get_score`` never asks for a component."""
+        self._assert_refused(ckpt_path={"_target_": "subprocess.Popen", "args": [["touch", self.marker]]})
+
+    def test_history_import_does_not_run_a_config_payload(self):
+        """End to end: a work directory with a malicious config is imported without executing it."""
+        payload = f"$__import__('pathlib').Path({self.marker!r}).write_text('pwned') or {self.algo_dir!r}"
+        self._write_config(ckpt_path=payload)
+        with open(os.path.join(self.algo_dir, "algo_object.json"), "w") as f:
+            json.dump(
+                {
+                    "_target_": "monai.apps.auto3dseg.bundle_gen.BundleAlgo",
+                    "_state_": {"template_path": self.root, "output_path": self.algo_dir},
+                },
+                f,
+            )
+
+        with self.assertLogs(level="WARNING"):
+            history = import_bundle_algo_history(self.root, only_trained=False)
+
+        self.assertFalse(os.path.exists(self.marker), "the work directory payload was evaluated")
+        self.assertEqual(len(history), 1)
+        self.assertFalse(history[0]["is_trained"])
+
+
+if __name__ == "__main__":
+    unittest.main()