kaxil commented on code in PR #74772:
URL: https://github.com/apache/airflow/pull/74772#discussion_r4282576227
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py:
##########
@@ -522,6 +528,121 @@ def
test_loop_task_reads_from_inside_the_loop_stay_on_one_pass(client, outside_l
assert client.get(f"{fanout_url}/item/-1").json() == "1-1"
[email protected]
+def deferred_loop_pass(session, outside_loop_run):
+ dr, _, _, steps, _ = outside_loop_run
+ trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={})
+ trigger.triggerer_id = TRIGGERER_ID
+ session.add(trigger)
+ session.flush()
+ steps[1].state = TaskInstanceState.DEFERRED
+ steps[1].trigger_id = trigger.id
+ session.commit()
+ yield dr, steps, trigger
+ clear_db_triggers()
+
+
+def test_trigger_in_a_loop_reads_and_writes_as_its_deferred_pass(session,
deferred_loop_pass):
+ dr, steps, _ = deferred_loop_pass
+ headers = {TRIGGER_TI_ID_HEADER: str(steps[1].id)}
+ step_url = f"/xcoms/{dr.dag_id}/{dr.run_id}/body.step/key"
+
+ sessions_begun = []
+
+ def record(api_session, *_):
+ sessions_begun.append(api_session)
+
+ with TestClient(InProcessExecutionAPI(triggerer_id=TRIGGERER_ID).app) as
api:
+ event.listen(Session, "after_begin", record)
+ try:
+ with capture_orm_selects("task_instance") as task_instance_selects:
+ assert api.get(step_url, headers=headers).json() == {"key":
"key", "value": "step-1"}
+ assert api.head(step_url,
headers=headers).headers["Content-Range"] == "map_indexes 1"
+ assert api.get(f"{step_url}/slice", headers=headers).json() ==
["step-1"]
+ assert api.get(f"{step_url}/item/0", headers=headers).json() ==
"step-1"
+ assert api.post(f"{step_url}-written", json="from-trigger",
headers=headers).status_code == 201
+ assert api.delete(step_url, headers=headers).status_code == 200
+ finally:
+ event.remove(Session, "after_begin", record)
+
+ liveness_checks = [
+ statement
+ for statement in task_instance_selects
+ if statement.startswith("SELECT task_instance.working_set FROM")
+ ]
+ assert len(liveness_checks) == 1
+ assert len(set(map(id, sessions_begun))) == 6
+ session.expire_all()
+ assert XComModelV2.get_for_attempt(steps[1].id, "key-written",
session=session).value == "from-trigger"
+ assert XComModelV2.get_for_attempt(steps[1].id, "key", session=session) is
None
+ assert XComModelV2.get_for_attempt(steps[0].id, "key",
session=session).value == "step-0"
+
+
[email protected](
+ ("api_triggerer_id", "header", "expected_status"),
+ [
+ pytest.param(None, None, 403, id="api-without-triggerer-id"),
+ pytest.param(TRIGGERER_ID, "not-a-uuid", 400, id="malformed-header"),
+ ],
+)
+def test_trigger_ti_id_header_rejected(deferred_loop_pass, api_triggerer_id,
header, expected_status):
+ dr, steps, _ = deferred_loop_pass
+
+ with TestClient(InProcessExecutionAPI(triggerer_id=api_triggerer_id).app)
as api:
+ response = api.get(
+ f"/xcoms/{dr.dag_id}/{dr.run_id}/body.step/key",
+ headers={TRIGGER_TI_ID_HEADER: header or str(steps[1].id)},
+ )
+
+ assert response.status_code == expected_status
+
+
[email protected](
+ ("method", "change", "expected_status", "expected_value"),
+ [
+ pytest.param(
+ "post",
+ {"triggerer_id": TRIGGERER_ID + 1},
+ 403,
+ None,
+ id="write-trigger-assigned-to-another-triggerer",
+ ),
+ pytest.param("post", {"state": TaskInstanceState.RUNNING}, 403, None,
id="write-task-not-deferred"),
+ pytest.param(
+ "get",
+ {"state": TaskInstanceState.FAILED, "trigger_removed": True},
+ 200,
+ "step-1",
+ id="read-live-task-no-longer-deferred",
+ ),
+ pytest.param("get", {"archived": True}, 403, None,
id="read-archived-try"),
Review Comment:
These cases only use `post` and `get`, so nothing pins which rule HEAD and
DELETE get. HEAD and DELETE only run in the happy-path test, where the task is
deferred on this triggerer and passes both checks. So changing `request.method
in ("GET", "HEAD")` in `_resolve_trigger_token` to `== "GET"`, or to `!=
"POST"` (which puts DELETE on the looser read rule), still passes everything. A
`head` case with the trigger removed (expect 200) and a `delete` case with
`RUNNING` (expect 403) would cover it. Asserting `detail["message"]` on the
403s would also show which check denied.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -252,29 +256,86 @@ def _build_xcom_read(
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error
-def xcom_query(
- dag_id: str,
- run_id: str,
- task_id: str,
- key: str,
+def _resolve_trigger_token(
+ request: Request,
session: SessionDep,
- dag_bag: DagBagDep,
- map_index: Annotated[int | None, Query()] = None,
- previous_iteration: bool = False,
token: ExecutionToken = CurrentExecutionToken,
-) -> Select:
- return _build_xcom_read(
- dag_id=dag_id,
- run_id=run_id,
- task_id=task_id,
- key=key,
- session=session,
- dag_bag=dag_bag,
- token=token,
- map_index=map_index,
- previous_iteration=previous_iteration,
- sequence=True,
- ).statement
+ dependency_solver=Depends(current_dependency_solver),
+) -> ExecutionToken:
+ """
+ Return the identity an XCom request acts as.
+
+ A request made by a trigger's code carries its task instance's id in
``TRIGGER_TI_ID_HEADER``. The
+ triggerer subprocess can forge the header, so it is only honoured on the
triggerer's in-process API.
+
+ A write acts as the claimed task instance only while it is deferred on a
trigger assigned to the calling
+ triggerer. A read acts as it while that try is live, in any state, because
a trigger's ``on_kill`` runs
+ after the task instance leaves ``deferred`` and the trigger row is gone;
reads are not a security
+ boundary on this API, which already serves every XCom. Anything else gets
a 403.
+ """
+ header = request.headers.get(TRIGGER_TI_ID_HEADER)
+ # The endpoint runs with Cadwyn's re-solved dependencies, so checking
there reuses the endpoint's session.
+ if header is None or dependency_solver != "cadwyn":
+ return token
+
+ triggerer_id = request.scope.get(_IN_PROCESS_TRIGGERER_ID)
+ if triggerer_id is None:
+ raise HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail={
+ "reason": "access_denied",
+ "message": f"{TRIGGER_TI_ID_HEADER} is only accepted from the
triggerer",
+ },
+ )
+
+ try:
+ claimed = UUID(header)
+ except ValueError:
+ raise HTTPException(
+ status_code=status.HTTP_400_BAD_REQUEST,
+ detail={"reason": "invalid_request", "message":
f"{TRIGGER_TI_ID_HEADER} must be a UUID"},
+ )
+
+ if request.method in ("GET", "HEAD"):
+ permitted = (
+ session.scalar(
+ select(TaskInstance.working_set)
+ .where(TaskInstance.id == claimed)
+ .execution_options(include_all_attempts=True)
+ )
+ is True
+ )
+ denial = "Task instance is not a live attempt"
+ else:
+ permitted = (
+ session.scalar(
+ select(TaskInstance.id)
+ .join(Trigger, Trigger.id == TaskInstance.trigger_id)
+ .where(
+ TaskInstance.id == claimed,
+ Trigger.triggerer_id == triggerer_id,
Review Comment:
What happens to the task when a trigger gets reassigned mid-run? If
triggerer A misses its heartbeat and the trigger moves to B, A's copy keeps
running until A cancels it. An XCom write from A in that window now gets a 403
here. The trigger raises before yielding an event, `cleanup_finished_triggers`
queues it as failed, and `Trigger.submit_failure` (which selects on
`trigger_id` and `DEFERRED` only, not on the owning triggerer) sends the task
to `__fail__` while B is still running it. Before this, a write from a trigger
outside a loop resolved by coordinates and went through. It's narrow, since no
in-tree trigger writes XComs from `run()`, but should `handle_failed_triggers`
skip triggers this triggerer no longer owns?
--
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]
