fix: Prevent Python RCE via __import__ on metadata classnames (#2514)

* Fixes

* Consolidate common code

* Rename to secure_import

Author: ranadeepsingh

cognitive/src/main/python/synapse/ml/services/langchain/LangchainTransform.py

@@ -48,6 +48,7 @@
 from pyspark.sql.types import StructType, StructField, StringType
 from typing import cast, Optional, TypeVar, Type
 from synapse.ml.core.platform import running_on_synapse_internal
+from synapse.ml.core.serialize._safe_import import secure_import_class
 
 OPENAI_API_VERSION = "2022-12-01"
 RL = TypeVar("RL", bound="MLReadable")
@@ -82,21 +83,9 @@ def saveImpl(self, path: str) -> None:
 
 
 class LangchainTransformerParamsReader(DefaultParamsReader):
-    @staticmethod
-    def __get_class(clazz: str) -> Type[RL]:
-        """
-        Loads Python class from its name.
-        """
-        parts = clazz.split(".")
-        module = ".".join(parts[:-1])
-        m = __import__(module, fromlist=[parts[-1]])
-        return getattr(m, parts[-1])
-
     def load(self, path: str) -> RL:
         metadata = LangchainTransformerParamsReader.loadMetadata(path, self.sc)
-        py_type: Type[RL] = LangchainTransformerParamsReader.__get_class(
-            metadata["class"]
-        )
+        py_type: Type[RL] = secure_import_class(metadata["class"])
         instance = py_type()
         cast("Params", instance)._resetUid(metadata["uid"])
         # deserialize the chain before setting Params

core/src/main/python/synapse/ml/core/schema/Utils.py

@@ -10,6 +10,7 @@
 from pyspark.ml.wrapper import JavaParams
 from pyspark.ml.common import inherit_doc, _java2py
 from pyspark import SparkContext
+from synapse.ml.core.serialize._safe_import import secure_import_class
 
 
 def from_java(java_stage, stage_name):
@@ -26,25 +27,8 @@ def from_java(java_stage, stage_name):
         object: The python wrapper
     """
 
-    def __get_class(clazz):
-        """
-        Loads a python object from its class
-
-        Args:
-            clazz (str): The name of the class
-
-        Returns:
-            object: The python object
-        """
-        parts = clazz.split(".")
-        module = ".".join(parts[:-1])
-        m = __import__(module)
-        for comp in parts[1:]:
-            m = getattr(m, comp)
-        return m
-
     # Generate a default new instance from the stage_name class.
-    py_type = __get_class(stage_name)
+    py_type = secure_import_class(stage_name)
     if issubclass(py_type, JavaParams):
         # Load information from java_stage to the instance.
         py_stage = py_type()

core/src/main/python/synapse/ml/core/serialize/_safe_import.py

@@ -0,0 +1,39 @@
+# Copyright (C) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See LICENSE in project root for information.
+
+"""
+Secure dynamic class loading for SynapseML model deserialization.
+
+All classname-based imports during model loading MUST go through
+``secure_import_class`` to prevent arbitrary code execution via
+crafted model metadata.
+"""
+
+import importlib
+
+_ALLOWED_MODULE_PREFIXES = ("pyspark.", "synapse.ml.")
+
+
+def secure_import_class(fully_qualified_name):
+    """
+    Import a Python class by its fully-qualified name, restricted to
+    an allowlist of trusted module prefixes.
+
+    Args:
+        fully_qualified_name (str): e.g. ``"pyspark.ml.classification.LogisticRegression"``
+
+    Returns:
+        type: The resolved Python class.
+
+    Raises:
+        ImportError: If the name does not start with an allowed prefix.
+    """
+    if not fully_qualified_name.startswith(_ALLOWED_MODULE_PREFIXES):
+        raise ImportError(
+            f"Refusing to load class '{fully_qualified_name}': "
+            f"module must start with one of {_ALLOWED_MODULE_PREFIXES}"
+        )
+    parts = fully_qualified_name.split(".")
+    module = ".".join(parts[:-1])
+    m = importlib.import_module(module)
+    return getattr(m, parts[-1])

core/src/main/python/synapse/ml/core/serialize/java_params_patch.py

@@ -1,13 +1,13 @@
 from pyspark.ml.wrapper import JavaParams, JavaWrapper
 from py4j.java_gateway import JavaObject
 from pyspark import RDD, SparkContext
-from pyspark.serializers import PickleSerializer, AutoBatchedSerializer
-from pyspark.sql import DataFrame, SQLContext
+from pyspark.serializers import PickleSerializer
+from pyspark.sql import DataFrame
 from pyspark.ml.common import _to_java_object_rdd, _java2py
-import pyspark
 from pyspark.ml import PipelineModel
 from pyspark.ml.util import DefaultParamsReader
 from pyspark.sql.types import DataType
+from synapse.ml.core.serialize._safe_import import secure_import_class
 
 
 @staticmethod
@@ -19,21 +19,10 @@ def _mml_from_java(java_stage):
     Meta-algorithms such as Pipeline should override this method as a classmethod.
     """
 
-    def __get_class(clazz):
-        """
-        Loads Python class from its name.
-        """
-        parts = clazz.split(".")
-        module = ".".join(parts[:-1])
-        m = __import__(module)
-        for comp in parts[1:]:
-            m = getattr(m, comp)
-        return m
-
     stage_name = java_stage.getClass().getName().replace("org.apache.spark", "pyspark")
     stage_name = stage_name.replace("com.microsoft.azure.synapse.ml", "synapse.ml")
     # Generate a default new instance from the stage_name class.
-    py_type = __get_class(stage_name)
+    py_type = secure_import_class(stage_name)
     if issubclass(py_type, JavaParams):
         # Load information from java_stage to the instance.
         py_stage = py_type()
@@ -59,17 +48,6 @@ def _mml_loadParamsInstance(path, sc):
     This assumes the instance inherits from :py:class:`MLReadable`.
     """
 
-    def __get_class(clazz):
-        """
-        Loads Python class from its name.
-        """
-        parts = clazz.split(".")
-        module = ".".join(parts[:-1])
-        m = __import__(module)
-        for comp in parts[1:]:
-            m = getattr(m, comp)
-        return m
-
     metadata = DefaultParamsReader.loadMetadata(path, sc)
     if DefaultParamsReader.isPythonParamsInstance(metadata):
         pythonClassName = metadata["class"]
@@ -78,7 +56,7 @@ def __get_class(clazz):
         pythonClassName = pythonClassName.replace(
             "com.microsoft.azure.synapse.ml", "synapse.ml"
         )
-    py_type = __get_class(pythonClassName)
+    py_type = secure_import_class(pythonClassName)
     instance = py_type.load(path)
     return instance