Experiment with expression rewriting

Signed-off-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>

Author: ericspod

monai/bundle/scripts.py

@@ -162,7 +162,7 @@ def _get_fake_spatial_shape(shape: Sequence[str | int], p: int = 1, n: int = 1,
                 for c in _get_var_names(i):
                     if c not in ["p", "n"]:
                         raise ValueError(f"only support variables 'p' and 'n' so far, but got: {c}.")
-                ret.append(safe_eval(i, {"p": p, "n": n}))
+                ret.append(safe_eval(i, {"p": p, "n": n}, rewrite_np=True))
         else:
             raise ValueError(f"spatial shape items must be int or string, but got: {type(i)} {i}.")
     return tuple(ret)

monai/utils/safeeval.py

@@ -15,6 +15,8 @@
 from collections.abc import Mapping, Sequence
 from typing import Any
 
+import numpy as np
+
 __all__ = ["SAFE_TYPES", "safe_eval"]
 
 # default set of safe AST node types
@@ -37,11 +39,28 @@
 )
 
 
+class _RewriteConstNp(ast.NodeTransformer):
+    """Replaces int and float constants in the tree with those wrapped in Numpy types."""
+    def __init__(self, int_type_str: str, float_type_str: str):
+        self.int_type_str = int_type_str
+        self.float_type_str = float_type_str
+
+    def visit_Constant(self, node):
+        if isinstance(node.value, (int, float)):
+            type_str = self.int_type_str if isinstance(node.value, int) else self.float_type_str
+            return ast.parse(f"{type_str}({node.value})")
+
+        return node
+
+
 def safe_eval(
     expr: str,
     globals_vars: Mapping[str, Any] | None = None,
     locals_vars: Mapping[str, object] | None = None,
     allowed_types: Sequence[type] = SAFE_TYPES,
+    rewrite_np: bool = False,
+    int_type_str: str = "np.int32",
+    float_type_str: str = "np.float32",
 ) -> Any:
     """
     Evaluate the Python expression `expr` using `eval`, but only if it is a safe expression in that its parsed AST
@@ -50,11 +69,20 @@ def safe_eval(
     expressions with constants and names can be evaluated, so excludes attribute access, indexing, and calls. Code
     injection is infeasible through such expressions, so this is a safe and secure way of evaluating simple expressions.
 
+    If `rewrite_np` is True, int and float constants in the given expression will be wrapped with Numpy types as given
+    by `int_type_str` and `float_type_str`. These are expected to be constructor names prefixed with `np.` as Numpy
+    will be present in the expression global variables under that name. The values can be changed to other types if
+    needed, such as "int64". One advantage of doing this is to avoid denial-of-service attacks by attempting to evaluate
+    an expressoini which is incredibly slow under native Python but fast (though potentially erroneous) under Numpy.
+
     Args:
         expr: expression to evaluate, this will be stripped before parsing to avoid indentation complaints
         globals_vars: global variable mapping, this will be treated as read-only for this function, unlike `eval`
         locals_vars: local variable mapping
         allowed_types: sequence of allowed AST types which can be found in `expr` when parsed
+        rewrite_np: if True, wrap int or float literals in Numpy types
+        int_type_str: int Numpy wrapping type string
+        float_type_str: float Numpy wrapping type string
 
     Raises:
         ValueError: raised when any node in the AST parsed from `expr` has a type not in `allowed_types`
@@ -70,4 +98,8 @@ def safe_eval(
     if disallowed:
         raise ValueError(f"Unsafe expression `{expr}` not evaluated, contains disallowed components: {disallowed}")
 
+    if rewrite_np:
+        parsed = _RewriteConstNp(int_type_str, float_type_str).visit(parsed)
+        locals_vars = {"np": np, **(locals_vars or {})}
+
     return eval(expr, dict(globals_vars) if globals_vars else None, locals_vars)

tests/utils/test_safe_eval.py

@@ -38,12 +38,21 @@ def test_good_exprs(self, expr, globals_vars, locals_vars, expected):
         result = safe_eval(expr, globals_vars, locals_vars)
         self.assertEqual(result, expected)
 
+    @parameterized.expand(GOOD_EXPRS)
+    def test_good_exprs_np(self, expr, globals_vars, locals_vars, expected):
+        """Test valid expressions with globals/locals evaluate to correct values with Numpy wrapping."""
+        result = safe_eval(expr, globals_vars, locals_vars, rewrite_np=True)
+        self.assertEqual(result, expected)
+
     @parameterized.expand(BAD_EXPRS)
     def test_bad_exprs(self, expr):
         """Test bad expressions correctly raise ValueError."""
         with self.assertRaises(ValueError):
             safe_eval(expr)
 
+        with self.assertRaises(ValueError):
+            safe_eval(expr, rewrite_np=True)
+
     def test_allowed_types(self):
         """Test restricting the allowed list of types."""
         allowed = [ast.Expression, ast.Constant, ast.BinOp, ast.Add]