Restrict asset state store access to registered tasks

airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py

@@ -22,8 +22,9 @@
 whichever identifier their inlet type carries: ``Asset``/``AssetNameRef`` use
 the name routes, ``AssetUriRef`` uses the URI routes.
 
-Per-task asset registration checks are intentionally not implemented here
-(deferred to AIP-93 — see TODO comment below).
+Task tokens must be registered with an asset as an inlet or outlet before
+accessing that asset's state store. Watcher tokens are not backed by task
+instances and are allowed through this route separately.
 """
 
 from __future__ import annotations
@@ -34,7 +35,7 @@
 
 from cadwyn import VersionedAPIRouter
 from fastapi import HTTPException, Query, status
-from sqlalchemy import select
+from sqlalchemy import exists, or_, select
 
 from airflow._shared.state import AssetScope, AssetStateStoreWriterKind
 from airflow.api_fastapi.common.db.common import SessionDep
@@ -44,7 +45,7 @@
 )
 from airflow.api_fastapi.execution_api.datamodels.token import TIToken
 from airflow.api_fastapi.execution_api.security import CurrentTIToken, ExecutionAPIRoute
-from airflow.models.asset import AssetModel
+from airflow.models.asset import AssetModel, TaskInletAssetReference, TaskOutletAssetReference
 from airflow.models.taskinstance import TaskInstance
 from airflow.state import get_state_backend
 from airflow.state.metastore import MetastoreBackend
@@ -71,12 +72,6 @@ def _fetch_ti_writer_fields(token: TIToken, session: SessionDep) -> _TIWriterFie
     return row.dag_id, row.run_id, row.task_id, row.map_index
 
 
-# TODO(AIP-103): enforce that the requesting task is registered with the asset
-# (via task_inlet_asset_reference or task_outlet_asset_reference) before
-# allowing reads/writes. Currently any task with a valid execution token can
-# access any asset's state store — the same gap exists in /assets and /asset-events.
-# Proper fix is a unified asset-registration check across all asset routes,
-# not just here.
 router = VersionedAPIRouter(
     route_class=ExecutionAPIRoute,
     responses={
@@ -106,14 +101,51 @@ def _resolve_asset_id_by_uri(uri: str, session: SessionDep) -> int:
     return asset_id
 
 
+def _authorize_asset_state_store_access(
+    token: TIToken, asset_id: int, session: SessionDep
+) -> _TIWriterFields | None:
+    if token.id == NULL_UUID:
+        return None
+
+    ti_fields = _fetch_ti_writer_fields(token, session)
+    dag_id, _, task_id, _ = ti_fields
+    is_registered = session.scalar(
+        select(
+            or_(
+                exists().where(
+                    TaskInletAssetReference.asset_id == asset_id,
+                    TaskInletAssetReference.dag_id == dag_id,
+                    TaskInletAssetReference.task_id == task_id,
+                ),
+                exists().where(
+                    TaskOutletAssetReference.asset_id == asset_id,
+                    TaskOutletAssetReference.dag_id == dag_id,
+                    TaskOutletAssetReference.task_id == task_id,
+                ),
+            )
+        )
+    )
+    if not is_registered:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail={
+                "reason": "forbidden",
+                "message": "Task is not registered with this asset as an inlet or outlet",
+            },
+        )
+    return ti_fields
+
+
 @router.get("/by-name/value")
 def get_asset_state_store_by_name(
     name: Annotated[str, Query(min_length=1)],
     key: Annotated[str, Query(min_length=1)],
     session: SessionDep,
+    token: TIToken = CurrentTIToken,
 ) -> AssetStateStoreResponse:
     """Get an asset state store value by asset name."""
     asset_id = _resolve_asset_id_by_name(name, session)
+    _authorize_asset_state_store_access(token, asset_id, session)
     value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session)
     if value is None:
         raise HTTPException(
@@ -129,6 +161,7 @@ def _put_asset_state_store(
     body: AssetStateStorePutBody,
     token: TIToken,
     session: SessionDep,
+    ti_fields: _TIWriterFields | None,
 ) -> None:
     backend = get_state_backend()
     if isinstance(backend, MetastoreBackend):
@@ -142,7 +175,8 @@ def _put_asset_state_store(
                 session=session,
             )
         else:
-            ti_fields = _fetch_ti_writer_fields(token, session)
+            if ti_fields is None:
+                ti_fields = _fetch_ti_writer_fields(token, session)
             dag_id, run_id, task_id, map_index = ti_fields
 
             backend.set_asset_state_store(
@@ -169,29 +203,33 @@ def set_asset_state_store_by_name(
     token: TIToken = CurrentTIToken,
 ) -> None:
     """Set an asset state store value by asset name."""
-    _put_asset_state_store(
-        AssetScope(asset_id=_resolve_asset_id_by_name(name, session)), key, body, token, session
-    )
+    asset_id = _resolve_asset_id_by_name(name, session)
+    ti_fields = _authorize_asset_state_store_access(token, asset_id, session)
+    _put_asset_state_store(AssetScope(asset_id=asset_id), key, body, token, session, ti_fields)
 
 
 @router.delete("/by-name/value", status_code=status.HTTP_204_NO_CONTENT)
 def delete_asset_state_store_by_name(
     name: Annotated[str, Query(min_length=1)],
     key: Annotated[str, Query(min_length=1)],
     session: SessionDep,
+    token: TIToken = CurrentTIToken,
 ) -> None:
     """Delete a single asset state store key by asset name."""
     asset_id = _resolve_asset_id_by_name(name, session)
+    _authorize_asset_state_store_access(token, asset_id, session)
     get_state_backend().delete(AssetScope(asset_id=asset_id), key, session=session)
 
 
 @router.delete("/by-name/clear", status_code=status.HTTP_204_NO_CONTENT)
 def clear_asset_state_store_by_name(
     name: Annotated[str, Query(min_length=1)],
     session: SessionDep,
+    token: TIToken = CurrentTIToken,
 ) -> None:
     """Delete all state store keys for an asset by asset name."""
     asset_id = _resolve_asset_id_by_name(name, session)
+    _authorize_asset_state_store_access(token, asset_id, session)
     get_state_backend().clear(AssetScope(asset_id=asset_id), session=session)
 
 
@@ -200,9 +238,11 @@ def get_asset_state_store_by_uri(
     uri: Annotated[str, Query(min_length=1)],
     key: Annotated[str, Query(min_length=1)],
     session: SessionDep,
+    token: TIToken = CurrentTIToken,
 ) -> AssetStateStoreResponse:
     """Get an asset state store value by asset URI."""
     asset_id = _resolve_asset_id_by_uri(uri, session)
+    _authorize_asset_state_store_access(token, asset_id, session)
     value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session)
     if value is None:
         raise HTTPException(
@@ -221,27 +261,31 @@ def set_asset_state_store_by_uri(
     token: TIToken = CurrentTIToken,
 ) -> None:
     """Set an asset state store value by asset URI."""
-    _put_asset_state_store(
-        AssetScope(asset_id=_resolve_asset_id_by_uri(uri, session)), key, body, token, session
-    )
+    asset_id = _resolve_asset_id_by_uri(uri, session)
+    ti_fields = _authorize_asset_state_store_access(token, asset_id, session)
+    _put_asset_state_store(AssetScope(asset_id=asset_id), key, body, token, session, ti_fields)
 
 
 @router.delete("/by-uri/value", status_code=status.HTTP_204_NO_CONTENT)
 def delete_asset_state_store_by_uri(
     uri: Annotated[str, Query(min_length=1)],
     key: Annotated[str, Query(min_length=1)],
     session: SessionDep,
+    token: TIToken = CurrentTIToken,
 ) -> None:
     """Delete a single asset state store key by asset URI."""
     asset_id = _resolve_asset_id_by_uri(uri, session)
+    _authorize_asset_state_store_access(token, asset_id, session)
     get_state_backend().delete(AssetScope(asset_id=asset_id), key, session=session)
 
 
 @router.delete("/by-uri/clear", status_code=status.HTTP_204_NO_CONTENT)
 def clear_asset_state_store_by_uri(
     uri: Annotated[str, Query(min_length=1)],
     session: SessionDep,
+    token: TIToken = CurrentTIToken,
 ) -> None:
     """Delete all state store keys for an asset by asset URI."""
     asset_id = _resolve_asset_id_by_uri(uri, session)
+    _authorize_asset_state_store_access(token, asset_id, session)
     get_state_backend().clear(AssetScope(asset_id=asset_id), session=session)

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_state_store.py

@@ -25,7 +25,7 @@
 
 from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
 from airflow.api_fastapi.execution_api.security import require_auth
-from airflow.models.asset import AssetActive, AssetModel
+from airflow.models.asset import AssetActive, AssetModel, TaskInletAssetReference, TaskOutletAssetReference
 from airflow.models.asset_state_store import AssetStateStoreModel
 from airflow.utils.session import create_session
 
@@ -49,6 +49,25 @@ def _create_asset_store_row(asset_id: int, key: str, value: str) -> None:
         s.add(AssetStateStoreModel(asset_id=asset_id, key=key, value=json.dumps(value)))
 
 
+def _create_task_asset_reference(
+    session: Session,
+    task_instance: TaskInstance,
+    asset: AssetModel,
+    reference_model=TaskOutletAssetReference,
+) -> None:
+    session.add(
+        reference_model(asset_id=asset.id, dag_id=task_instance.dag_id, task_id=task_instance.task_id)
+    )
+    session.commit()
+
+
+def _request_asset_state_store(client: TestClient, method: str, path: str, *, params, json_body=None):
+    kwargs = {"params": params}
+    if json_body is not None:
+        kwargs["json"] = json_body
+    return getattr(client, method)(path, **kwargs)
+
+
 _BY_NAME_CLEAR = "/execution/store/asset/by-name/clear"
 _BY_URI_VALUE = "/execution/store/asset/by-uri/value"
 _BY_URI_CLEAR = "/execution/store/asset/by-uri/clear"
@@ -58,6 +77,8 @@ def _create_asset_store_row(asset_id: int, key: str, value: str) -> None:
 def reset_state_tables():
     with create_session() as session:
         session.execute(delete(AssetStateStoreModel))
+        session.execute(delete(TaskInletAssetReference))
+        session.execute(delete(TaskOutletAssetReference))
         session.execute(delete(AssetModel))
 
 
@@ -129,6 +150,10 @@ async def _auth(request: Request) -> TIToken:
 
         exec_app.dependency_overrides[require_auth] = _auth
 
+    @pytest.fixture(autouse=True)
+    def setup_asset_reference(self, setup_ti_auth, asset: AssetModel, session: Session):
+        _create_task_asset_reference(session, self._ti, asset)
+
     def test_put_creates_row(self, client: TestClient, asset: AssetModel, session: Session):
         response = client.put(
             _BY_NAME_VALUE, params={"name": asset.name, "key": "watermark"}, json={"value": "2026-04-29"}
@@ -302,6 +327,10 @@ async def _auth(request: Request) -> TIToken:
 
         exec_app.dependency_overrides[require_auth] = _auth
 
+    @pytest.fixture(autouse=True)
+    def setup_asset_reference(self, setup_ti_auth, asset: AssetModel, session: Session):
+        _create_task_asset_reference(session, self._ti, asset)
+
     def test_put_creates_row(self, client: TestClient, asset: AssetModel, session: Session):
         response = client.put(
             _BY_URI_VALUE, params={"uri": asset.uri, "key": "watermark"}, json={"value": "2026-04-29"}
@@ -370,6 +399,90 @@ def test_clear_removes_all_keys(self, client: TestClient, asset: AssetModel):
             assert row is None
 
 
+class TestTaskAssetRegistration:
+    @pytest.fixture(autouse=True)
+    def setup_ti_auth(
+        self,
+        exec_app: FastAPI,
+        create_task_instance: CreateTaskInstance,
+    ):
+        self._ti: TaskInstance = create_task_instance()
+
+        async def _auth(request: Request) -> TIToken:
+            return TIToken(id=self._ti.id, claims=TIClaims(scope="execution"))
+
+        exec_app.dependency_overrides[require_auth] = _auth
+
+    @pytest.mark.parametrize(
+        ("method", "path", "params", "json_body"),
+        [
+            ("get", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, None),
+            ("put", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, {"value": "x"}),
+            ("delete", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, None),
+            ("delete", _BY_NAME_CLEAR, {"name": "test_asset"}, None),
+            ("get", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, None),
+            ("put", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, {"value": "x"}),
+            ("delete", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, None),
+            ("delete", _BY_URI_CLEAR, {"uri": "s3://bucket/test"}, None),
+        ],
+    )
+    def test_unregistered_task_is_forbidden(
+        self,
+        client: TestClient,
+        asset: AssetModel,
+        method: str,
+        path: str,
+        params: dict,
+        json_body: dict | None,
+    ):
+        response = _request_asset_state_store(client, method, path, params=params, json_body=json_body)
+
+        assert response.status_code == 403
+        assert response.json()["detail"]["reason"] == "forbidden"
+
+    @pytest.mark.parametrize(
+        ("method", "path", "params", "json_body", "expected_status"),
+        [
+            ("get", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, None, 200),
+            ("put", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, {"value": "x"}, 204),
+            ("delete", _BY_NAME_VALUE, {"name": "test_asset", "key": "watermark"}, None, 204),
+            ("delete", _BY_NAME_CLEAR, {"name": "test_asset"}, None, 204),
+            ("get", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, None, 200),
+            ("put", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, {"value": "x"}, 204),
+            ("delete", _BY_URI_VALUE, {"uri": "s3://bucket/test", "key": "watermark"}, None, 204),
+            ("delete", _BY_URI_CLEAR, {"uri": "s3://bucket/test"}, None, 204),
+        ],
+    )
+    def test_outlet_registered_task_can_access_asset_state_store(
+        self,
+        client: TestClient,
+        asset: AssetModel,
+        session: Session,
+        method: str,
+        path: str,
+        params: dict,
+        json_body: dict | None,
+        expected_status: int,
+    ):
+        _create_task_asset_reference(session, self._ti, asset)
+        _create_asset_store_row(asset.id, "watermark", "2026-04-29")
+
+        response = _request_asset_state_store(client, method, path, params=params, json_body=json_body)
+
+        assert response.status_code == expected_status
+
+    def test_inlet_registered_task_can_access_asset_state_store(
+        self, client: TestClient, asset: AssetModel, session: Session
+    ):
+        _create_task_asset_reference(session, self._ti, asset, TaskInletAssetReference)
+        _create_asset_store_row(asset.id, "watermark", "2026-04-29")
+
+        response = client.get(_BY_NAME_VALUE, params={"name": asset.name, "key": "watermark"})
+
+        assert response.status_code == 200
+        assert response.json() == {"value": "2026-04-29"}
+
+
 class TestInactiveAssetRejected:
     """An asset row without a corresponding asset_active entry is treated as not found."""