This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 3bd5f761f5c Stop Task SDK client identifiers from reaching other API
endpoints (#68617)
3bd5f761f5c is described below
commit 3bd5f761f5c0e0b38cc44ff4642ea9970041dd10
Author: Nayyar <[email protected]>
AuthorDate: Sat Nov 7 15:29:10 2026 +0530
Stop Task SDK client identifiers from reaching other API endpoints (#68617)
An identifier with "/", "." or ".." could send the request to a different
Execution API endpoint.
The server decodes "%2F" before it routes the request, so the client
rejects "/" in a Dag, run or task ID.
Co-authored-by: Shahar Epstein <[email protected]>
---
task-sdk/src/airflow/sdk/api/client.py | 67 ++++++++++-------
task-sdk/tests/task_sdk/api/test_client.py | 114 +++++++++++++++++++++++++++++
2 files changed, 156 insertions(+), 25 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/api/client.py
b/task-sdk/src/airflow/sdk/api/client.py
index 2e8a2753424..0eb04302736 100644
--- a/task-sdk/src/airflow/sdk/api/client.py
+++ b/task-sdk/src/airflow/sdk/api/client.py
@@ -250,6 +250,22 @@ def _log_and_trace_retry(retry_state) -> None:
)
+def _path(prefix: str, *identifiers: str | uuid.UUID, tail: str | None = None)
-> str:
+ """
+ Build a request path, quoting each identifier so it stays a single segment.
+
+ ``prefix`` is the fixed, trusted portion of the path. The server decodes
``%2F`` before
+ routing, so a ``/`` in a positional identifier would shift it onto another
route and is
+ rejected. ``tail`` is the last segment (a ``:path`` key or a final id) and
may contain ``/``.
+ """
+ if any("/" in str(identifier) for identifier in identifiers):
+ raise ValueError(f"Path identifier must not contain '/':
{identifiers!r}")
+ segments = [*identifiers] if tail is None else [*identifiers, tail]
+ quoted = (quote(str(segment), safe="") for segment in segments)
+ # quote() leaves a bare "." or ".." as is, and httpx would still collapse
it as a dot segment.
+ return "/".join([prefix, *("%2E" * len(q) if q in (".", "..") else q for q
in quoted)])
+
+
class TaskInstanceOperations:
__slots__ = ("client",)
@@ -460,7 +476,7 @@ class TaskInstanceOperations:
if state:
params["state"] = state.value if isinstance(state,
TaskInstanceState) else state
- resp = self.client.get(f"task-instances/previous/{dag_id}/{task_id}",
params=params)
+ resp = self.client.get(_path("task-instances/previous", dag_id,
task_id), params=params)
return PreviousTIResult(task_instance=resp.json())
def get_task_states(
@@ -515,7 +531,7 @@ class ConnectionOperations:
def get(self, conn_id: str) -> ConnectionResponse | ErrorResponse:
"""Get a connection from the API server."""
try:
- resp = self.client.get(f"connections/{quote(conn_id, safe='')}")
+ resp = self.client.get(_path("connections", tail=conn_id))
except ServerResponseError as e:
if e.response.status_code == HTTPStatus.NOT_FOUND:
log.debug(
@@ -553,7 +569,7 @@ class VariableOperations:
def get(self, key: str) -> VariableResponse | ErrorResponse:
"""Get a variable from the API server."""
try:
- resp = self.client.get(f"variables/{key}")
+ resp = self.client.get(_path("variables", tail=key))
except ServerResponseError as e:
if e.response.status_code == HTTPStatus.NOT_FOUND:
log.debug(
@@ -582,7 +598,7 @@ class VariableOperations:
def set(self, key: str, value: str, description: str | None = None) ->
OKResponse:
"""Set an Airflow Variable via the API server."""
body = VariablePostBody(val=value, description=description)
- self.client.put(f"variables/{key}", content=body.model_dump_json())
+ self.client.put(_path("variables", tail=key),
content=body.model_dump_json())
# Any error from the server will anyway be propagated down to the
supervisor,
# so we choose to send a generic response to the supervisor over the
server response to
# decouple from the server response string
@@ -593,7 +609,7 @@ class VariableOperations:
key: str,
) -> OKResponse:
"""Delete a variable with given key via the API server."""
- self.client.delete(f"variables/{key}")
+ self.client.delete(_path("variables", tail=key))
# Any error from the server will anyway be propagated down to the
supervisor,
# so we choose to send a generic response to the supervisor over the
server response to
# decouple from the server response string
@@ -627,7 +643,7 @@ class XComOperations:
params: dict[str, str | int] = {}
if previous_iteration:
params["previous_iteration"] = True
- resp = self.client.head(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}",
params=params)
+ resp = self.client.head(_path("xcoms", dag_id, run_id, task_id,
tail=key), params=params)
# content_range: str | None
if not (content_range := resp.headers["Content-Range"]) or not
content_range.startswith(
@@ -656,7 +672,7 @@ class XComOperations:
if include_prior_dates:
params.update({"include_prior_dates": include_prior_dates})
try:
- resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}",
params=params)
+ resp = self.client.get(_path("xcoms", dag_id, run_id, task_id,
tail=key), params=params)
except ServerResponseError as e:
if e.response.status_code == HTTPStatus.NOT_FOUND:
log.debug(
@@ -696,7 +712,7 @@ class XComOperations:
params["map_index"] = map_index
if mapped_length is not None and mapped_length >= 0:
params["mapped_length"] = mapped_length
- self.client.post(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}",
params=params, json=value)
+ self.client.post(_path("xcoms", dag_id, run_id, task_id, tail=key),
params=params, json=value)
# Any error from the server will anyway be propagated down to the
supervisor,
# so we choose to send a generic response to the supervisor over the
server response to
# decouple from the server response string
@@ -714,7 +730,7 @@ class XComOperations:
params: dict[str, str | int] = {}
if map_index is not None and map_index >= 0:
params["map_index"] = map_index
- self.client.delete(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}",
params=params)
+ self.client.delete(_path("xcoms", dag_id, run_id, task_id, tail=key),
params=params)
# Any error from the server will anyway be propagated down to the
supervisor,
# so we choose to send a generic response to the supervisor over the
server response to
# decouple from the server response string
@@ -735,7 +751,7 @@ class XComOperations:
params["previous_iteration"] = True
try:
resp = self.client.get(
- f"xcoms/{dag_id}/{run_id}/{task_id}/{key}/item/{offset}",
+ f"{_path('xcoms', dag_id, run_id, task_id,
tail=key)}/item/{offset}",
params=params,
)
except ServerResponseError as e:
@@ -787,7 +803,7 @@ class XComOperations:
params["step"] = step
if include_prior_dates:
params["include_prior_dates"] = include_prior_dates
- resp =
self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}/slice", params=params)
+ resp = self.client.get(f"{_path('xcoms', dag_id, run_id, task_id,
tail=key)}/slice", params=params)
return XComSequenceSliceResponse.model_validate_json(resp.read())
@@ -800,7 +816,7 @@ class TaskStateStoreOperations:
def get(self, ti_id: uuid.UUID, key: str) -> TaskStateStoreResponse |
ErrorResponse:
"""Get a task store value from the API server."""
try:
- resp = self.client.get(f"store/ti/{ti_id}/{quote(key, safe='')}")
+ resp = self.client.get(_path("store/ti", ti_id, tail=key))
except ServerResponseError as e:
if e.response.status_code == HTTPStatus.NOT_FOUND:
log.debug("Task store key not found", ti_id=ti_id, key=key)
@@ -811,17 +827,17 @@ class TaskStateStoreOperations:
def set(self, ti_id: uuid.UUID, key: str, value: JsonValue, expires_at:
datetime | None) -> OKResponse:
"""Set a task store value via the API server."""
body = TaskStateStorePutBody(value=value, expires_at=expires_at)
- self.client.put(f"store/ti/{ti_id}/{quote(key, safe='')}",
content=body.model_dump_json())
+ self.client.put(_path("store/ti", ti_id, tail=key),
content=body.model_dump_json())
return OKResponse(ok=True)
def delete(self, ti_id: uuid.UUID, key: str) -> OKResponse:
"""Delete a single task store key via the API server."""
- self.client.delete(f"store/ti/{ti_id}/{quote(key, safe='')}")
+ self.client.delete(_path("store/ti", ti_id, tail=key))
return OKResponse(ok=True)
def clear(self, ti_id: uuid.UUID) -> OKResponse:
"""Clear all task store keys for a task instance via the API server."""
- self.client.delete(f"store/ti/{ti_id}")
+ self.client.delete(_path("store/ti", ti_id))
return OKResponse(ok=True)
@@ -1001,7 +1017,8 @@ class DagRunOperations:
try:
self.client.post(
- f"dag-runs/{dag_id}/{run_id}",
content=body.model_dump_json(exclude_defaults=True)
+ _path("dag-runs", dag_id, run_id),
+ content=body.model_dump_json(exclude_defaults=True),
)
except ServerResponseError as e:
if e.response.status_code == HTTPStatus.CONFLICT:
@@ -1017,18 +1034,18 @@ class DagRunOperations:
def clear(self, dag_id: str, run_id: str) -> OKResponse:
"""Clear a Dag run via the API server."""
- self.client.post(f"dag-runs/{dag_id}/{run_id}/clear")
+ self.client.post(f"{_path('dag-runs', dag_id, run_id)}/clear")
# TODO: Error handling
return OKResponse(ok=True)
def get_detail(self, dag_id: str, run_id: str) -> DagRun:
"""Get detail of a dag run."""
- resp = self.client.get(f"dag-runs/{dag_id}/{run_id}")
+ resp = self.client.get(_path("dag-runs", dag_id, run_id))
return DagRun.model_validate_json(resp.read())
def get_state(self, dag_id: str, run_id: str) -> DagRunStateResponse:
"""Get the state of a Dag run via the API server."""
- resp = self.client.get(f"dag-runs/{dag_id}/{run_id}/state")
+ resp = self.client.get(f"{_path('dag-runs', dag_id, run_id)}/state")
return DagRunStateResponse.model_validate_json(resp.read())
def get_count(
@@ -1077,7 +1094,7 @@ class DagsOperations:
def get(self, dag_id: str) -> DagResponse:
"""Get a DAG via the API server."""
- resp = self.client.get(f"dags/{quote(dag_id, safe='')}")
+ resp = self.client.get(_path("dags", tail=dag_id))
return DagResponse.model_validate_json(resp.read())
@@ -1117,7 +1134,7 @@ class HITLOperations:
assigned_users=assigned_users,
)
resp = self.client.post(
- f"/hitlDetails/{ti_id}",
+ _path("/hitlDetails", ti_id),
content=payload.model_dump_json(),
# One key per request, reused by the client's retries, so a
retried POST cannot replace
# the request (and discard an answer) after the first attempt was
stored.
@@ -1139,14 +1156,14 @@ class HITLOperations:
params_input=params_input,
)
resp = self.client.patch(
- f"/hitlDetails/{ti_id}",
+ _path("/hitlDetails", ti_id),
content=payload.model_dump_json(),
)
return HITLDetailResponse.model_validate_json(resp.read())
def get_detail_response(self, ti_id: uuid.UUID) -> HITLDetailResponse:
"""Get content part of a Human-in-the-loop response for a specific
Task Instance."""
- resp = self.client.get(f"/hitlDetails/{ti_id}")
+ resp = self.client.get(_path("/hitlDetails", ti_id))
return HITLDetailResponse.model_validate_json(resp.read())
@@ -1158,7 +1175,7 @@ class ConnectionTestOperations:
def get_connection(self, connection_test_id: uuid.UUID) ->
ConnectionTestConnectionResponse:
"""Fetch connection data for a test request from the API server."""
- resp =
self.client.get(f"connection-tests/{connection_test_id}/connection")
+ resp = self.client.get(f"{_path('connection-tests',
connection_test_id)}/connection")
return
ConnectionTestConnectionResponse.model_validate_json(resp.read())
def update_state(
@@ -1171,7 +1188,7 @@ class ConnectionTestOperations:
state=state,
result_message=ResultMessage(result_message) if result_message is
not None else None,
)
- self.client.patch(f"connection-tests/{id}",
content=body.model_dump_json())
+ self.client.patch(_path("connection-tests", id),
content=body.model_dump_json())
class CallbackOperations:
diff --git a/task-sdk/tests/task_sdk/api/test_client.py
b/task-sdk/tests/task_sdk/api/test_client.py
index c9b862fb336..ca0c2793a52 100644
--- a/task-sdk/tests/task_sdk/api/test_client.py
+++ b/task-sdk/tests/task_sdk/api/test_client.py
@@ -17,6 +17,7 @@
from __future__ import annotations
+import contextlib
import json
import pickle
from datetime import UTC, datetime
@@ -382,6 +383,119 @@ class TestClient:
with pytest.raises(ValueError, match="Both client_ssl_cert and
client_ssl_key must be set"):
make_client(httpx2.MockTransport(handle_request))
+ @pytest.mark.parametrize(
+ ("call", "expected_path"),
+ [
+ pytest.param(
+ lambda c: c.variables.get("../connections/aws_default"),
+ b"/variables/..%2Fconnections%2Faws_default",
+ id="variables-get",
+ ),
+ pytest.param(
+ lambda c: c.variables.set("../connections/aws_default", "v"),
+ b"/variables/..%2Fconnections%2Faws_default",
+ id="variables-set",
+ ),
+ pytest.param(
+ lambda c: c.variables.delete("../connections/aws_default"),
+ b"/variables/..%2Fconnections%2Faws_default",
+ id="variables-delete",
+ ),
+ pytest.param(
+ lambda c: c.xcoms.head("d", "r", "t", "x/../../variables/k"),
+ b"/xcoms/d/r/t/x%2F..%2F..%2Fvariables%2Fk",
+ id="xcoms-head",
+ ),
+ pytest.param(
+ lambda c: c.xcoms.get("d", "r", "t", "x/../../variables/k"),
+ b"/xcoms/d/r/t/x%2F..%2F..%2Fvariables%2Fk",
+ id="xcoms-get",
+ ),
+ pytest.param(
+ lambda c: c.xcoms.set("d", "r", "t", "x/../../variables/k",
"v"),
+ b"/xcoms/d/r/t/x%2F..%2F..%2Fvariables%2Fk",
+ id="xcoms-set",
+ ),
+ pytest.param(
+ lambda c: c.xcoms.delete("d", "r", "t", "x/../../variables/k"),
+ b"/xcoms/d/r/t/x%2F..%2F..%2Fvariables%2Fk",
+ id="xcoms-delete",
+ ),
+ pytest.param(
+ lambda c: c.xcoms.get_sequence_item("d", "r", "t",
"x/../../variables/k", 0),
+ b"/xcoms/d/r/t/x%2F..%2F..%2Fvariables%2Fk/item/0",
+ id="xcoms-item",
+ ),
+ pytest.param(
+ lambda c: c.xcoms.get_sequence_slice("d", "r", "t",
"x/../../variables/k", None, None, None),
+ b"/xcoms/d/r/t/x%2F..%2F..%2Fvariables%2Fk/slice",
+ id="xcoms-slice",
+ ),
+ pytest.param(
+ lambda c: c.dag_runs.trigger("..", "r"),
b"/dag-runs/%2E%2E/r", id="dag-runs-trigger"
+ ),
+ pytest.param(
+ lambda c: c.dag_runs.clear("d", ".."),
b"/dag-runs/d/%2E%2E/clear", id="dag-runs-clear"
+ ),
+ pytest.param(
+ lambda c: c.dag_runs.get_detail("..", "r"),
b"/dag-runs/%2E%2E/r", id="dag-runs-detail"
+ ),
+ pytest.param(
+ lambda c: c.dag_runs.get_state("d", ".."),
b"/dag-runs/d/%2E%2E/state", id="dag-runs-state"
+ ),
+ pytest.param(
+ lambda c: c.task_instances.get_previous("..", "t"),
+ b"/task-instances/previous/%2E%2E/t",
+ id="get-previous",
+ ),
+ pytest.param(lambda c: c.connections.get(".."),
b"/connections/%2E%2E", id="connection-dotdot"),
+ pytest.param(lambda c: c.connections.get("."),
b"/connections/%2E", id="connection-dot"),
+ pytest.param(
+ lambda c: c.xcoms.get("..", "..", "variables", "secret"),
+ b"/xcoms/%2E%2E/%2E%2E/variables/secret",
+ id="xcoms-dots",
+ ),
+ ],
+ )
+ def test_path_identifiers_stay_in_their_segment(self, call, expected_path):
+ """Caller-supplied identifiers cannot reach another endpoint, whatever
they contain."""
+ requests_seen = []
+
+ def handle_request(request: httpx2.Request) -> httpx2.Response:
+ requests_seen.append(request)
+ return httpx2.Response(
+ status_code=200, content=b"null", headers={"Content-Range":
"map_indexes 1"}
+ )
+
+ with contextlib.suppress(Exception): # the stub body need not parse;
only the URL matters
+ call(make_client(transport=httpx2.MockTransport(handle_request)))
+
+ assert requests_seen[0].url.raw_path.split(b"?")[0] == expected_path
+
+ @pytest.mark.parametrize(
+ "call",
+ [
+ pytest.param(lambda c: c.dag_runs.trigger("d", "x/clear"),
id="trigger-run-id"),
+ pytest.param(lambda c: c.dag_runs.trigger("a/b", "r"),
id="trigger-dag-id"),
+ pytest.param(lambda c: c.xcoms.get("a/b", "r", "t", "k"),
id="xcoms-dag-id"),
+ pytest.param(lambda c: c.xcoms.get("d", "a/b", "t", "k"),
id="xcoms-run-id"),
+ pytest.param(lambda c: c.xcoms.get("d", "r", "a/b", "k"),
id="xcoms-task-id"),
+ pytest.param(lambda c: c.task_instances.get_previous("a/b", "t"),
id="previous-dag-id"),
+ pytest.param(lambda c: c.task_instances.get_previous("d", "a/b"),
id="previous-task-id"),
+ ],
+ )
+ def test_path_rejects_slash_in_positional_identifier(self, call):
+ requests_seen = []
+
+ def handle_request(request: httpx2.Request) -> httpx2.Response:
+ requests_seen.append(request)
+ return httpx2.Response(status_code=200, content=b"null")
+
+ with pytest.raises(ValueError, match="must not contain '/'"):
+ call(make_client(transport=httpx2.MockTransport(handle_request)))
+
+ assert requests_seen == []
+
class TestTaskInstanceOperations:
"""
