From b2d36a8e0a60072dbfb1fe42e0068b85a2983835 Mon Sep 17 00:00:00 2001 From: Shivam <6463385+shivaam@users.noreply.github.com> Date: Wed, 29 Jul 2026 23:55:31 -0700 Subject: [PATCH] Prevent tasks from writing XComs as other task instances Execution API task tokens identify a single task instance, but XCom writes previously trusted the URL coordinates independently. A valid task token could therefore write under another task instance identity within its allowed team boundary. --- .../airflow/api_fastapi/execution_api/app.py | 34 +++++- .../api_fastapi/execution_api/routes/xcoms.py | 48 +++++++-- .../api_fastapi/execution_api/test_app.py | 46 ++++++++ .../execution_api/versions/head/test_xcoms.py | 100 ++++++++++++++++-- 4 files changed, 206 insertions(+), 22 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/app.py b/airflow-core/src/airflow/api_fastapi/execution_api/app.py index f9e7cb1725000..9b16641a559cb 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/app.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/app.py @@ -32,10 +32,11 @@ Cadwyn, current_dependency_solver, ) -from fastapi import Depends, FastAPI, Request, Response +from fastapi import Depends, FastAPI, HTTPException, Request, Response, status from fastapi.responses import JSONResponse from fastapi.routing import APIRoute from opentelemetry import context as otel_context, propagate as otel_propagate +from sqlalchemy import select from starlette.middleware.base import BaseHTTPMiddleware from airflow.api_fastapi.auth.tokens import ( @@ -44,6 +45,7 @@ get_sig_validation_args, get_signing_args, ) +from airflow.api_fastapi.common.db.common import SessionDep if TYPE_CHECKING: import httpx @@ -390,8 +392,9 @@ def app(self): from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken from airflow.api_fastapi.execution_api.routes.connections import has_connection_access from airflow.api_fastapi.execution_api.routes.variables import has_variable_access - from airflow.api_fastapi.execution_api.routes.xcoms import has_xcom_access + from airflow.api_fastapi.execution_api.routes.xcoms import get_xcom_write_ti, has_xcom_access from airflow.api_fastapi.execution_api.security import _jwt_bearer + from airflow.models.taskinstance import TaskInstance # Give this app its own lifespan + services registry so that stubbing services # (e.g. JWTValidator) doesn't affect the module-level ``lifespan.registry``. @@ -415,10 +418,37 @@ async def always_allow(request: Request): claims = TIClaims(scope="execution") return TIToken(id=ti_id, claims=claims) + def resolve_xcom_write_ti( + dag_id: str, + run_id: str, + task_id: str, + map_index: int = -1, + *, + session: SessionDep, + ) -> TaskInstance: + ti = session.scalar( + select(TaskInstance).where( + TaskInstance.dag_id == dag_id, + TaskInstance.run_id == run_id, + TaskInstance.task_id == task_id, + TaskInstance.map_index == map_index, + ) + ) + if ti is None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "reason": "access_denied", + "message": "Task may only set XComs for its own task instance", + }, + ) + return ti + self._app.dependency_overrides[_jwt_bearer] = always_allow self._app.dependency_overrides[has_connection_access] = always_allow self._app.dependency_overrides[has_variable_access] = always_allow self._app.dependency_overrides[has_xcom_access] = always_allow + self._app.dependency_overrides[get_xcom_write_ti] = resolve_xcom_write_ti return self._app diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py index 0cb6ccb23ce54..8f4c3142fff7a 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py @@ -27,12 +27,14 @@ from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.core_api.base import BaseModel +from airflow.api_fastapi.execution_api.datamodels.token import TIToken from airflow.api_fastapi.execution_api.datamodels.xcom import ( XComResponse, XComSequenceIndexResponse, XComSequenceSliceResponse, ) from airflow.api_fastapi.execution_api.security import CurrentTIToken +from airflow.models.taskinstance import TaskInstance from airflow.models.taskmap import TaskMap from airflow.models.xcom import XComModel from airflow.utils.db import get_query_count @@ -116,6 +118,33 @@ def has_xcom_access( log = logging.getLogger(__name__) +def get_xcom_write_ti( + dag_id: str, + run_id: str, + task_id: str, + map_index: int = -1, + token: TIToken = CurrentTIToken, + *, + session: SessionDep, +) -> TaskInstance: + """Resolve and authorize the task instance that owns an XCom write.""" + ti = session.get(TaskInstance, token.id) + if ti is None or (dag_id, run_id, task_id, map_index) != ( + ti.dag_id, + ti.run_id, + ti.task_id, + ti.map_index, + ): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "reason": "access_denied", + "message": "Task may only set XComs for its own task instance", + }, + ) + return ti + + async def xcom_query( dag_id: str, run_id: str, @@ -356,8 +385,6 @@ def get_xcom( return XComResponse(key=key, value=(result[0] if isinstance(result, tuple) else result).value) -# TODO: once we have JWT tokens, then remove dag_id/run_id/task_id from the URL and just use the info in -# the token @router.post( "/{dag_id}/{run_id}/{task_id}/{key:path}", status_code=status.HTTP_201_CREATED, @@ -368,6 +395,7 @@ def set_xcom( task_id: str, key: Annotated[str, Path(min_length=1)], session: SessionDep, + ti: Annotated[TaskInstance, Depends(get_xcom_write_ti)], value: Annotated[ JsonValue, Body( @@ -410,10 +438,10 @@ def set_xcom( if mapped_length is not None: task_map = TaskMap( - dag_id=dag_id, - task_id=task_id, - run_id=run_id, - map_index=map_index, + dag_id=ti.dag_id, + task_id=ti.task_id, + run_id=ti.run_id, + map_index=ti.map_index, length=mapped_length, keys=None, ) @@ -437,10 +465,10 @@ def set_xcom( XComModel.set( key=key, value=value, - run_id=run_id, - task_id=task_id, - dag_id=dag_id, - map_index=map_index, + run_id=ti.run_id, + task_id=ti.task_id, + dag_id=ti.dag_id, + map_index=ti.map_index, serialize=False, dag_result=dag_result, session=session, diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py index bb2d2d557dc36..d0bbdde79f141 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_app.py @@ -29,6 +29,7 @@ from fastapi.routing import APIRoute from fastapi.testclient import TestClient from opentelemetry import context as otel_context, propagate as otel_propagate +from sqlalchemy import select from sqlalchemy.exc import SQLAlchemyError from airflow.api_fastapi.execution_api.app import ( @@ -40,6 +41,8 @@ from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken from airflow.api_fastapi.execution_api.security import require_auth from airflow.api_fastapi.execution_api.versions import bundle +from airflow.models.xcom import XComModel +from airflow.sdk.serde import serialize from tests_common.test_utils.config import conf_vars @@ -163,6 +166,49 @@ def test_in_process_execution_api_runs_without_jwt_secret(): assert response.status_code == 200 +@pytest.mark.parametrize("map_index", [-1, 2]) +def test_in_process_execution_api_sets_xcom_for_route_task_instance(create_task_instance, session, map_index): + ti = create_task_instance(map_index=map_index) + session.commit() + params = {"map_index": map_index} if map_index >= 0 else None + + api = InProcessExecutionAPI() + with httpx.Client(transport=api.transport, base_url="http://localhost") as client: + response = client.post( + f"/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/in_process", + params=params, + json=serialize('"value"'), + ) + + assert response.status_code == status.HTTP_201_CREATED, response.json() + xcom = session.scalar( + select(XComModel).where( + XComModel.dag_id == ti.dag_id, + XComModel.run_id == ti.run_id, + XComModel.task_id == ti.task_id, + XComModel.map_index == map_index, + XComModel.key == "in_process", + ) + ) + assert xcom is not None + assert xcom.value == '"value"' + + +def test_in_process_execution_api_rejects_xcom_for_missing_route_task_instance(session): + api = InProcessExecutionAPI() + with httpx.Client(transport=api.transport, base_url="http://localhost") as client: + response = client.post("/xcoms/missing/run/task/key", json=serialize('"value"')) + + assert response.status_code == status.HTTP_403_FORBIDDEN, response.json() + assert response.json() == { + "detail": { + "reason": "access_denied", + "message": "Task may only set XComs for its own task instance", + } + } + assert session.scalar(select(XComModel).where(XComModel.key == "key")) is None + + def test_in_process_execution_api_transport_lifecycle(): """The background loop + thread lifecycle is tied to the ``.transport``, not the factory instance. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py index ae476d0344ca0..e2f40a83e4fb6 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py @@ -83,6 +83,17 @@ def _( exec_app.dependency_overrides = {} +@pytest.fixture +def authenticate_as(exec_app): + def _authenticate_as(ti_id): + async def _auth(request: Request) -> TIToken: + return TIToken(id=ti_id, claims=TIClaims(scope="execution")) + + exec_app.dependency_overrides[require_auth] = _auth + + return _authenticate_as + + class TestXComsGetEndpoint: @pytest.mark.parametrize( ("db_value"), @@ -340,7 +351,7 @@ class TestXComsSetEndpoint: (None, None), ], ) - def test_xcom_set(self, client, create_task_instance, session, value, expected_value): + def test_xcom_set(self, client, create_task_instance, session, authenticate_as, value, expected_value): """ Test that XCom value is set correctly. The request body can be either: - a JSON string (e.g. '"value"', '{"k":"v"}', '[1]'), which is stored as-is (a string) in the DB @@ -351,6 +362,7 @@ def test_xcom_set(self, client, create_task_instance, session, value, expected_v """ ti = create_task_instance() session.commit() + authenticate_as(ti.id) value = serialize(value) response = client.post( f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/xcom_1", @@ -373,6 +385,57 @@ def test_xcom_set(self, client, create_task_instance, session, value, expected_v ).one_or_none() assert task_map is None, "Should not be mapped" + @pytest.mark.parametrize("mismatch", ["dag_id", "run_id", "task_id", "map_index"]) + def test_xcom_set_rejects_identity_mismatch( + self, client, create_task_instance, session, authenticate_as, mismatch + ): + ti = create_task_instance(map_index=3) + session.commit() + authenticate_as(ti.id) + identity = { + "dag_id": ti.dag_id, + "run_id": ti.run_id, + "task_id": ti.task_id, + "map_index": ti.map_index, + } + identity[mismatch] = 4 if mismatch == "map_index" else f"other_{mismatch}" + + response = client.post( + f"/execution/xcoms/{identity['dag_id']}/{identity['run_id']}/{identity['task_id']}/xcom_1", + params={"map_index": identity["map_index"]}, + json='"value1"', + ) + + assert response.status_code == 403 + assert response.json() == { + "detail": { + "reason": "access_denied", + "message": "Task may only set XComs for its own task instance", + } + } + assert session.scalar(select(XComModel).where(XComModel.key == "xcom_1")) is None + + def test_xcom_set_rejects_missing_token_task_instance( + self, client, create_task_instance, session, authenticate_as + ): + ti = create_task_instance() + session.commit() + authenticate_as(uuid4()) + + response = client.post( + f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/xcom_1", + json='"value1"', + ) + + assert response.status_code == 403 + assert response.json() == { + "detail": { + "reason": "access_denied", + "message": "Task may only set XComs for its own task instance", + } + } + assert session.scalar(select(XComModel).where(XComModel.key == "xcom_1")) is None + @pytest.mark.parametrize( ("orig_value", "ser_value", "deser_value"), [ @@ -394,7 +457,16 @@ def test_xcom_set(self, client, create_task_instance, session, value, expected_v ), ], ) - def test_xcom_round_trip(self, client, create_task_instance, session, orig_value, ser_value, deser_value): + def test_xcom_round_trip( + self, + client, + create_task_instance, + session, + authenticate_as, + orig_value, + ser_value, + deser_value, + ): """ Test that deserialization works when XCom values are stored directly in the DB with API Server. @@ -408,6 +480,7 @@ def test_xcom_round_trip(self, client, create_task_instance, session, orig_value ti = create_task_instance() session.commit() + authenticate_as(ti.id) # Serialize the value to simulate the client SDK value = serialize(orig_value) @@ -437,15 +510,16 @@ def test_xcom_round_trip(self, client, create_task_instance, session, orig_value # Ensure that the deserialized value on the client side is the same as the original value assert deserialize(deserialized_value) == orig_value - def test_xcom_set_mapped(self, client, create_task_instance, session): - ti = create_task_instance() + def test_xcom_set_mapped(self, client, create_task_instance, session, authenticate_as): + ti = create_task_instance(map_index=2) session.commit() + authenticate_as(ti.id) value = serialize("value1") response = client.post( f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/xcom_1", - params={"map_index": -1, "mapped_length": 3}, + params={"map_index": 2, "mapped_length": 3}, json=value, ) @@ -457,7 +531,7 @@ def test_xcom_set_mapped(self, client, create_task_instance, session): XComModel.task_id == ti.task_id, XComModel.dag_id == ti.dag_id, XComModel.key == "xcom_1", - XComModel.map_index == -1, + XComModel.map_index == 2, ) ).first() assert xcom.value == "value1" @@ -468,7 +542,7 @@ def test_xcom_set_mapped(self, client, create_task_instance, session): assert task_map.dag_id == "dag" assert task_map.run_id == "test" assert task_map.task_id == "op1" - assert task_map.map_index == -1 + assert task_map.map_index == 2 assert task_map.length == 3 @pytest.mark.parametrize( @@ -479,7 +553,7 @@ def test_xcom_set_mapped(self, client, create_task_instance, session): ], ) def test_xcom_set_downstream_of_mapped( - self, client, create_task_instance, session, length, expected_status + self, client, create_task_instance, session, authenticate_as, length, expected_status ): """ Test that XCom value is set correctly. The value is passed as a JSON string in the request body. @@ -488,6 +562,7 @@ def test_xcom_set_downstream_of_mapped( """ ti = create_task_instance() session.commit() + authenticate_as(ti.id) response = client.post( f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/xcom_1", @@ -527,7 +602,9 @@ def test_xcom_access_denied(self, client, caplog): ('["value1"]', '["value1"]'), ], ) - def test_xcom_roundtrip(self, client, create_task_instance, session, value, expected_value): + def test_xcom_roundtrip( + self, client, create_task_instance, session, authenticate_as, value, expected_value + ): """ Test that XCom value is set and retrieved correctly using API. @@ -540,6 +617,7 @@ def test_xcom_roundtrip(self, client, create_task_instance, session, value, expe value = serialize(value) session.commit() + authenticate_as(ti.id) client.post( f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/test_xcom_roundtrip", json=value, @@ -559,11 +637,13 @@ def test_xcom_roundtrip(self, client, create_task_instance, session, value, expe assert response.status_code == 200 assert XComResponse.model_validate_json(response.read()).value == expected_value - def test_xcom_dag_result(self, client, create_task_instance, session): + def test_xcom_dag_result(self, client, create_task_instance, session, authenticate_as): """ Test that the dag_result flag propagates to XComModel. """ ti = create_task_instance() + session.commit() + authenticate_as(ti.id) client.post( f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/return_value", params={"dag_result": True},