Fix team name authorization bypass in edgeworker

providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py

@@ -19,7 +19,7 @@
 
 from typing import TYPE_CHECKING, Annotated
 
-from fastapi import Body, Depends, status
+from fastapi import Body, Depends, HTTPException, status
 from sqlalchemy import select, update
 
 from airflow.api_fastapi.common.db.common import SessionDep  # noqa: TC001
@@ -28,6 +28,7 @@
 from airflow.executors.workloads import ExecuteTask
 from airflow.providers.common.compat.sdk import Stats, timezone
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
+from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
 from airflow.providers.edge3.version_compat import AIRFLOW_V_3_3_PLUS
 from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest
 from airflow.providers.edge3.worker_api.datamodels import (
@@ -61,6 +62,7 @@ def parse_command(command: str, dag_id: str, run_id: str) -> ExecuteTypeBody:
         [
             status.HTTP_400_BAD_REQUEST,
             status.HTTP_403_FORBIDDEN,
+            status.HTTP_404_NOT_FOUND,
         ]
     ),
 )
@@ -76,6 +78,10 @@ def fetch(
     session: SessionDep,
 ) -> EdgeJobFetched | None:
     """Fetch a job to execute on the edge worker."""
+    worker = session.scalar(select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name))
+    if not worker:
+        raise HTTPException(status.HTTP_404_NOT_FOUND, "Worker not found")
+
     query = (
         select(EdgeJobModel)
         .where(
@@ -86,7 +92,8 @@ def fetch(
     )
     if body.queues:
         query = query.where(EdgeJobModel.queue.in_(body.queues))
-    query = query.where(EdgeJobModel.team_name == body.team_name)
+    if worker.team_name is not None:
+        query = query.where(EdgeJobModel.team_name == worker.team_name)
     query = query.limit(1)
     query = query.with_for_update(skip_locked=True)
     job: EdgeJobModel | None = session.scalar(query)

providers/edge3/src/airflow/providers/edge3/worker_api/v2-edge-generated.yaml

@@ -65,6 +65,12 @@ paths:
               schema:
                 $ref: '#/components/schemas/HTTPExceptionResponse'
           description: Forbidden
+        '404':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Not Found
         '422':
           description: Validation Error
           content:

providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py

@@ -23,11 +23,13 @@
 from uuid import uuid4
 
 import pytest
+from fastapi import HTTPException, status
 from sqlalchemy import delete, select
 
 from airflow.executors.workloads import BundleInfo, ExecuteTask
 from airflow.providers.common.compat.sdk import Stats
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
+from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState
 from airflow.providers.edge3.models.types import EXECUTE_CALLBACK_TAG
 from airflow.providers.edge3.worker_api.datamodels import WorkerQueuesBody
 from airflow.providers.edge3.worker_api.routes.jobs import fetch, parse_command, state
@@ -80,6 +82,7 @@ class TestJobsApiRoutes:
     @pytest.fixture(autouse=True)
     def setup_test_cases(self, dag_maker, session: Session):
         session.execute(delete(EdgeJobModel))
+        session.execute(delete(EdgeWorkerModel))
         session.commit()
 
     @patch(f"{Stats.__module__}.Stats.incr")
@@ -135,8 +138,13 @@ def test_state(self, mock_stats_incr, session: Session):
             assert db_job is not None
             assert db_job.state == TaskInstanceState.SUCCESS
 
-    def test_fetch_filters_by_team_name(self, session: Session):
+    def test_fetch_filters_by_worker_team_name(self, session: Session):
         with create_session() as session:
+            session.add(
+                EdgeWorkerModel(
+                    worker_name="worker1", state=EdgeWorkerState.IDLE, queues=[QUEUE], team_name="team_a"
+                )
+            )
             job_team_a = EdgeJobModel(
                 dag_id="dag_a",
                 task_id="task_a",
@@ -164,15 +172,29 @@ def test_fetch_filters_by_team_name(self, session: Session):
             session.add_all([job_team_a, job_team_b])
             session.commit()
 
-            body = WorkerQueuesBody(free_concurrency=1, queues=[QUEUE], team_name="team_a")
+            body = WorkerQueuesBody(free_concurrency=1, queues=[QUEUE], team_name="team_b")
             result = fetch("worker1", body, session)
             assert result is not None
             assert result.dag_id == "dag_a"
             assert result.task_id == "task_a"
 
+    def test_fetch_unknown_worker_raises_404(self, session: Session):
+        body = WorkerQueuesBody(free_concurrency=1, queues=[QUEUE], team_name="team_a")
+
+        with pytest.raises(HTTPException) as exc_info:
+            fetch("unknown-worker", body, session)
+
+        assert exc_info.value.status_code == status.HTTP_404_NOT_FOUND
+        assert exc_info.value.detail == "Worker not found"
+
     def test_fetch_without_team_name_returns_any_team(self, session: Session):
-        """When team_name is None, no team filter is applied so any queued job can be returned."""
+        """When a worker has no team_name, no team filter is applied so any queued job can be returned."""
         with create_session() as session:
+            session.add(
+                EdgeWorkerModel(
+                    worker_name="worker1", state=EdgeWorkerState.IDLE, queues=[QUEUE], team_name=None
+                )
+            )
             job_team_a = EdgeJobModel(
                 dag_id="dag_a",
                 task_id="task_a",