Restrict exception-node deserialization to known classes without importing the stored name

airflow-core/src/airflow/serialization/serialized_objects.py

@@ -41,7 +41,7 @@
 from dateutil import relativedelta
 from pendulum.tz.timezone import FixedTimezone, Timezone
 
-from airflow._shared.module_loading import import_string, qualname
+from airflow._shared.module_loading import qualname
 from airflow._shared.timezones.timezone import from_timestamp, parse_timezone, utcnow
 from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest
 from airflow.exceptions import AirflowException, DeserializationError, SerializationError
@@ -242,6 +242,31 @@ def _decode_priority_weight_strategy(var: str) -> PriorityWeightStrategy:
     return priority_weight_strategy_class()
 
 
+# Builtin exceptions a BASE_EXC_SER node can rebuild into. A user defined subclass
+# (e.g. a custom KeyError) serializes to a name absent here and won't round-trip --
+# unchanged from before, since builtins never held it either.
+_DESERIALIZABLE_BUILTIN_EXCEPTIONS: dict[str, type[BaseException]] = {
+    "KeyError": KeyError,
+    "AttributeError": AttributeError,
+}
+
+
+def _resolve_airflow_exception(exc_cls_name: str) -> type[AirflowException]:
+    """
+    Resolve a stored ``AirflowException`` name without importing it.
+
+    Read via ``vars()`` rather than ``getattr`` -- a module's deprecation-shim
+    ``__getattr__`` can still import on access -- which also lets pre-3.2.0 blobs
+    naming the old ``airflow.exceptions`` path keep resolving.
+    """
+    module_name, _, attr_name = exc_cls_name.rpartition(".")
+    module = sys.modules.get(module_name)
+    exc_cls = vars(module).get(attr_name) if module is not None else None
+    if not (isinstance(exc_cls, type) and issubclass(exc_cls, AirflowException)):
+        raise DeserializationError(f"Refusing to deserialize unknown exception class {exc_cls_name!r}")
+    return exc_cls
+
+
 def _encode_start_trigger_args(var: StartTriggerArgs) -> dict[str, Any]:
     """Encode a StartTriggerArgs."""
 
@@ -664,9 +689,14 @@ def deserialize(cls, encoded_var: Any) -> Any:
             kwargs = deser["kwargs"]
             del deser
             if type_ == DAT.AIRFLOW_EXC_SER:
-                exc_cls = import_string(exc_cls_name)
+                exc_cls: type[BaseException] = _resolve_airflow_exception(exc_cls_name)
             else:
-                exc_cls = import_string(f"builtins.{exc_cls_name}")
+                builtin_exc_cls = _DESERIALIZABLE_BUILTIN_EXCEPTIONS.get(exc_cls_name)
+                if builtin_exc_cls is None:
+                    raise DeserializationError(
+                        f"Refusing to deserialize unsupported builtin exception {exc_cls_name!r}"
+                    )
+                exc_cls = builtin_exc_cls
             return exc_cls(*args, **kwargs)
         elif type_ == DAT.SET:
             return {cls.deserialize(v) for v in var}

airflow-core/tests/unit/serialization/test_dag_serialization.py

@@ -52,6 +52,7 @@
 from airflow.dag_processing.dagbag import DagBag
 from airflow.exceptions import (
     AirflowException,
+    DeserializationError,
     ParamValidationError,
     SerializationError,
 )
@@ -77,7 +78,7 @@
 from airflow.serialization.definitions.param import SerializedParam
 from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg
 from airflow.serialization.encoders import ensure_serialized_asset
-from airflow.serialization.enums import Encoding
+from airflow.serialization.enums import DagAttributeTypes, Encoding
 from airflow.serialization.json_schema import load_dag_schema_dict
 from airflow.serialization.serialized_objects import (
     BaseSerialization,
@@ -519,6 +520,10 @@ def timetable_plugin(monkeypatch: pytest.MonkeyPatch):
     )
 
 
+class PluginContributedError(AirflowException):
+    """Defined outside airflow.exceptions on purpose: resolution must not be limited to exceptions Airflow itself ships."""
+
+
 class TestStringifiedDAGs:
     """Unit tests for stringified DAGs."""
 
@@ -2806,6 +2811,120 @@ def test_create_dagrun_accepts_partition_key_for_partitioned_at_runtime_dag(self
         dr = dag_maker.create_dagrun(partition_key="runtime-key")
         assert dr.partition_key == "runtime-key"
 
+    @pytest.mark.parametrize(
+        ("module_name", "attr_name"),
+        [
+            pytest.param("subprocess", "check_output", id="callable_in_loaded_module"),
+            pytest.param("os", "system", id="another_callable"),
+            pytest.param("builtins", "eval", id="builtin_callable"),
+            pytest.param("airflow.exceptions", "NoSuchThing", id="missing_attr"),
+        ],
+    )
+    def test_airflow_exc_deserialization_rejects_a_name_in_a_loaded_module(self, module_name, attr_name):
+        """Having the module loaded is not enough -- the name must be an AirflowException subclass.
+
+        The module is imported by the test first, so the rejection cannot be an artefact of the
+        module simply being absent.
+        """
+        importlib.import_module(module_name)
+        encoded = BaseSerialization._encode(
+            BaseSerialization.serialize(
+                {"exc_cls_name": f"{module_name}.{attr_name}", "args": [], "kwargs": {}}
+            ),
+            type_=DagAttributeTypes.AIRFLOW_EXC_SER,
+        )
+        with pytest.raises(DeserializationError, match="Refusing to deserialize unknown exception class"):
+            BaseSerialization.deserialize(encoded)
+
+    @pytest.mark.parametrize(
+        "exc_cls_name",
+        [
+            pytest.param("not.a.loaded.module.Thing", id="unloaded_module"),
+            pytest.param("NoModulePart", id="no_module_part"),
+            pytest.param("", id="empty"),
+        ],
+    )
+    def test_airflow_exc_deserialization_rejects_an_unresolvable_name(self, exc_cls_name):
+        encoded = BaseSerialization._encode(
+            BaseSerialization.serialize({"exc_cls_name": exc_cls_name, "args": [], "kwargs": {}}),
+            type_=DagAttributeTypes.AIRFLOW_EXC_SER,
+        )
+        with pytest.raises(DeserializationError, match="Refusing to deserialize unknown exception class"):
+            BaseSerialization.deserialize(encoded)
+
+    def test_airflow_exc_deserialization_does_not_import_the_named_module(self):
+        """Resolution reads ``sys.modules``; it never imports what the blob names."""
+        module_name = "airflow_exc_module_that_must_not_be_imported"
+        encoded = BaseSerialization._encode(
+            BaseSerialization.serialize({"exc_cls_name": f"{module_name}.Boom", "args": [], "kwargs": {}}),
+            type_=DagAttributeTypes.AIRFLOW_EXC_SER,
+        )
+        with mock.patch.object(
+            importlib, "import_module", side_effect=AssertionError("imported"), autospec=True
+        ):
+            with pytest.raises(DeserializationError):
+                BaseSerialization.deserialize(encoded)
+        assert module_name not in sys.modules
+
+    def test_airflow_exc_deserialization_roundtrips_airflow_exception(self):
+        result = BaseSerialization.deserialize(BaseSerialization.serialize(AirflowException("boom")))
+        assert isinstance(result, AirflowException)
+        assert result.args == ("boom",)
+
+    def test_airflow_exc_deserialization_resolves_a_subclass_outside_airflow(self):
+        """Resolves as long as its module is loaded, registration order doesn't matter."""
+        result = BaseSerialization.deserialize(BaseSerialization.serialize(PluginContributedError("boom")))
+        assert isinstance(result, PluginContributedError)
+        assert result.args == ("boom",)
+
+    @pytest.mark.parametrize(
+        "exc_cls_name",
+        [
+            "airflow.exceptions.AirflowException",
+            "airflow.exceptions.AirflowNotFoundException",
+            "airflow.exceptions.ParamValidationError",
+            "airflow.sdk.exceptions.AirflowException",
+        ],
+    )
+    def test_airflow_exc_deserialization_accepts_the_pre_3_2_module_spelling(self, exc_cls_name):
+        """A blob written before these exceptions moved to ``airflow.sdk.exceptions`` still reads.
+
+        3.0/3.1 stored ``airflow.exceptions.<Name>``; 3.2.0 moved the classes and left a re-export
+        behind. Both spellings have to resolve, or upgrading strands every stored blob that carries
+        an exception node.
+        """
+        encoded = BaseSerialization._encode(
+            BaseSerialization.serialize({"exc_cls_name": exc_cls_name, "args": ["boom"], "kwargs": {}}),
+            type_=DagAttributeTypes.AIRFLOW_EXC_SER,
+        )
+        result = BaseSerialization.deserialize(encoded)
+        assert isinstance(result, AirflowException)
+
+    def test_base_exc_deserialization_rejects_non_allowlisted_builtin(self):
+        """A BASE_EXC_SER name outside the {KeyError, AttributeError} the encoder emits is rejected."""
+        # ``eval`` is the weaponisable case; ``ValueError`` is a harmless builtin the encoder
+        # never emits as BASE_EXC_SER -- both must be rejected.
+        for name in ("eval", "ValueError"):
+            encoded = BaseSerialization._encode(
+                BaseSerialization.serialize({"exc_cls_name": name, "args": ["1"], "kwargs": {}}),
+                type_=DagAttributeTypes.BASE_EXC_SER,
+            )
+            with pytest.raises(
+                DeserializationError, match="Refusing to deserialize unsupported builtin exception"
+            ):
+                BaseSerialization.deserialize(encoded)
+
+    @pytest.mark.parametrize("exc_type", [KeyError, AttributeError])
+    def test_base_exc_serialize_deserialize_round_trip(self, exc_type):
+        """Pins the allow-list to the encode branch, by going through ``serialize`` rather than
+        hand-building the node: if that branch ever accepts another builtin, this fails."""
+        result = BaseSerialization.deserialize(BaseSerialization.serialize(exc_type("boom")))
+
+        assert isinstance(result, exc_type)
+        # The encode branch stores ``[var.args]``, so the args arrive nested. Pre-existing shape,
+        # asserted as it is rather than as it ought to be.
+        assert result.args == (("boom",),)
+
 
 def test_kubernetes_optional():
     """Test that serialization module loads without kubernetes, but deserialization of PODs requires it"""