amoghrajesh commented on code in PR #68511:
URL: https://github.com/apache/airflow/pull/68511#discussion_r3701474129


##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -242,6 +242,54 @@ def _decode_priority_weight_strategy(var: str) -> 
PriorityWeightStrategy:
     return priority_weight_strategy_class()
 
 
+# Builtin exceptions the serializer emits as ``BASE_EXC_SER``. Only these are 
ever
+# serialized (see the encode side), so deserialization resolves the stored 
name against
+# this fixed map instead of importing it -- ``builtins.eval`` / 
``builtins.exec`` and any
+# other name are rejected without importing anything.
+_DESERIALIZABLE_BUILTIN_EXCEPTIONS: dict[str, type[BaseException]] = {
+    "KeyError": KeyError,
+    "AttributeError": AttributeError,
+}
+
+
+def _iter_subclasses(cls: type) -> Iterator[type]:
+    """Yield every (transitive) subclass of ``cls``."""
+    for sub in cls.__subclasses__():
+        yield sub
+        yield from _iter_subclasses(sub)
+
+
+@cache
+def _serializable_airflow_exceptions() -> dict[str, type[AirflowException]]:
+    """
+    Map ``"<module>.<name>" -> AirflowException subclass``, used to resolve 
``AIRFLOW_EXC_SER`` nodes.
+
+    Built once, from the in-memory ``AirflowException`` subclass tree (never 
from the
+    attacker-controlled stored name), and never rebuilt -- a name absent from 
it is rejected, not
+    imported. ``airflow.exceptions`` is imported by this module, so every 
built-in Airflow exception
+    is registered by the time this is first called; exceptions defined later 
are not added.
+    """
+    return {
+        f"{cls.__module__}.{cls.__name__}": cls
+        for cls in (AirflowException, *_iter_subclasses(AirflowException))
+    }

Review Comment:
   This never rebuilds the list, and it is quite common to have a plugin / 
provider subclass AirflowException lazily, this will not be able to deser them. 
This isn't just "blocked at launch" (per @uranusjr's comment above), it is 
blocked forever in that process due to the cache never invalidating.
   
   Maybe we can have a two tier lookup? When a cache miss happens, clear 
`_serializable_airflow_exceptions.cache_clear()` and rebuild once before 
rejecting, so a class imported later in the lifecycle still resolves.



##########
airflow-core/tests/unit/serialization/test_dag_serialization.py:
##########
@@ -2806,6 +2806,64 @@ 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"
 
+    def test_airflow_exc_deserialization_rejects_unknown_class(self):
+        """An AIRFLOW_EXC_SER name that is not a loaded AirflowException 
subclass is rejected.
+
+        The name is resolved against the in-memory subclass tree, so an 
attacker-controlled
+        ``subprocess.check_output`` is never imported.
+        """
+        from airflow.exceptions import DeserializationError
+        from airflow.serialization.enums import DagAttributeTypes

Review Comment:
   Top level import pls.



##########
airflow-core/tests/unit/serialization/test_dag_serialization.py:
##########
@@ -2806,6 +2806,64 @@ 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"
 
+    def test_airflow_exc_deserialization_rejects_unknown_class(self):
+        """An AIRFLOW_EXC_SER name that is not a loaded AirflowException 
subclass is rejected.
+
+        The name is resolved against the in-memory subclass tree, so an 
attacker-controlled
+        ``subprocess.check_output`` is never imported.
+        """
+        from airflow.exceptions import DeserializationError
+        from airflow.serialization.enums import DagAttributeTypes
+
+        encoded = BaseSerialization._encode(
+            BaseSerialization.serialize(
+                {"exc_cls_name": "subprocess.check_output", "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_roundtrips_airflow_exception(self):
+        """A genuine AirflowException subclass round-trips via the registry, 
without importing."""
+        from airflow.exceptions import AirflowException
+
+        result = 
BaseSerialization.deserialize(BaseSerialization.serialize(AirflowException("boom")))
+        assert isinstance(result, AirflowException)
+        assert result.args == ("boom",)
+
+    def test_base_exc_deserialization_rejects_non_allowlisted_builtin(self):
+        """A BASE_EXC_SER name outside the {KeyError, AttributeError} the 
encoder emits is rejected."""
+        from airflow.exceptions import DeserializationError
+        from airflow.serialization.enums import DagAttributeTypes
+
+        # ``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 
disallowed builtin exception"
+            ):
+                BaseSerialization.deserialize(encoded)
+
+    def test_base_exc_deserialization_roundtrips_builtin_exception(self):
+        """The builtin exceptions the encoder emits (KeyError / 
AttributeError) still deserialize."""
+        from airflow.serialization.enums import DagAttributeTypes

Review Comment:
   Top level import pls.



##########
airflow-core/tests/unit/serialization/test_dag_serialization.py:
##########
@@ -2806,6 +2806,64 @@ 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"
 
+    def test_airflow_exc_deserialization_rejects_unknown_class(self):
+        """An AIRFLOW_EXC_SER name that is not a loaded AirflowException 
subclass is rejected.
+
+        The name is resolved against the in-memory subclass tree, so an 
attacker-controlled
+        ``subprocess.check_output`` is never imported.
+        """
+        from airflow.exceptions import DeserializationError
+        from airflow.serialization.enums import DagAttributeTypes
+
+        encoded = BaseSerialization._encode(
+            BaseSerialization.serialize(
+                {"exc_cls_name": "subprocess.check_output", "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_roundtrips_airflow_exception(self):
+        """A genuine AirflowException subclass round-trips via the registry, 
without importing."""
+        from airflow.exceptions import AirflowException
+
+        result = 
BaseSerialization.deserialize(BaseSerialization.serialize(AirflowException("boom")))
+        assert isinstance(result, AirflowException)
+        assert result.args == ("boom",)
+
+    def test_base_exc_deserialization_rejects_non_allowlisted_builtin(self):
+        """A BASE_EXC_SER name outside the {KeyError, AttributeError} the 
encoder emits is rejected."""
+        from airflow.exceptions import DeserializationError
+        from airflow.serialization.enums import DagAttributeTypes

Review Comment:
   Top level import pls.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to