diff --git a/ui/backend/server/tests/test_db_methods/test_permissions.py b/ui/backend/server/tests/test_db_methods/test_permissions.py index 0d4dd4e45..40b813495 100644 --- a/ui/backend/server/tests/test_db_methods/test_permissions.py +++ b/ui/backend/server/tests/test_db_methods/test_permissions.py @@ -40,6 +40,7 @@ user_can_write_to_dag_template, user_can_write_to_project, user_project_visibility, + visible_project_ids_for_user, ) from trackingserver_projects.models import Project, ProjectTeamMembership, ProjectUserMembership from trackingserver_projects.schema import ProjectIn, ProjectUpdate, Visibility @@ -641,3 +642,41 @@ async def test_user_cannot_get_dag_runs_with_no_access(db, by): can_do, message = await user_can_get_latest_dag_runs(authenticated_request, **kwargs) assert can_do is False, message + + +@pytest.mark.asyncio +async def test_visible_project_ids_for_user_only_returns_member_projects(db): + """visible_project_ids_for_user must only return projects the user has a + membership to (directly or via a team). It must never return projects the + user has no link to -- that property backs the visibility scoping in the + "get latest dag runs" endpoint when neither project_id nor dag_template_id + is supplied. + """ + # Project A is accessible by user A only. + project_a, user_a = await _setup_project_accessible_only_by( + "user_individual@no_team.com", "write" + ) + # Project B is accessible by user B only. + project_b, user_b = await _setup_project_accessible_only_by( + "user_1_team_1@team1.com", "write" + ) + request_a = await _get_authenticated_request("user_individual@no_team.com") + request_b = await _get_authenticated_request("user_1_team_1@team1.com") + user_a_obj, teams_a = request_a.auth + user_b_obj, teams_b = request_b.auth + + ids_a = await visible_project_ids_for_user(user_a_obj, teams_a) + ids_b = await visible_project_ids_for_user(user_b_obj, teams_b) + assert project_a.id in ids_a + assert project_b.id not in ids_a + assert project_b.id in ids_b + assert project_a.id not in ids_b + + # A user with no memberships sees no projects. + request_stranger = await _get_authenticated_request( + "user_with_no_permissions@no_one_invited_me.com" + ) + stranger, stranger_teams = request_stranger.auth + ids_stranger = await visible_project_ids_for_user(stranger, stranger_teams) + assert project_a.id not in ids_stranger + assert project_b.id not in ids_stranger diff --git a/ui/backend/server/tests/test_lifecycle/test_run_tracking.py b/ui/backend/server/tests/test_lifecycle/test_run_tracking.py index 707eeedc3..6d76fc46a 100644 --- a/ui/backend/server/tests/test_lifecycle/test_run_tracking.py +++ b/ui/backend/server/tests/test_lifecycle/test_run_tracking.py @@ -307,3 +307,126 @@ async def test_node_run_lifecycle_multiple_updates(async_client: AsyncClient, db assert set(attr["name"] for attr in realized_attrs) == set( attr["name"] for attr in expected_attrs ) + + +async def _create_dag_run_via_api( + async_client: AsyncClient, username: str, dag_template_id: int, tags: dict +) -> int: + response = await async_client.post( + f"/api/v1/dag_runs?{urlencode({'dag_template_id': dag_template_id})}", + data=dict( + run_start_time=datetime.datetime.now(), + run_end_time=None, + run_status="RUNNING", + tags=tags, + inputs={}, + outputs=[], + ), + content_type="application/json", + headers={"test_username": username}, + ) + assert response.status_code == 200, response.content + return response.json()["id"] + + +@pytest.mark.asyncio +async def test_get_latest_dag_runs_without_params_excludes_other_users_runs( + async_client: AsyncClient, db +): + """A user calling ``GET /api/v1/dag_runs/latest/`` with no parameters must + only see runs from projects they have visibility into. Calls without any + filter previously returned every active run on the instance. + """ + user_a = "user_individual@no_team.com" + user_b = "user_1_team_1@team1.com" + dag_template_a, _ = await _setup_dag_template(async_client, user_a) + dag_template_b, _ = await _setup_dag_template(async_client, user_b) + + run_a_id = await _create_dag_run_via_api( + async_client, user_a, dag_template_a, {"owner": "a"} + ) + run_b_id = await _create_dag_run_via_api( + async_client, user_b, dag_template_b, {"owner": "b"} + ) + + # User A asks for latest runs with no filter -- must not see user B's run. + response = await async_client.get( + "/api/v1/dag_runs/latest/", headers={"test_username": user_a} + ) + assert response.status_code == 200, response.content + returned_ids = {item["id"] for item in response.json()} + assert run_a_id in returned_ids + assert run_b_id not in returned_ids + + # Symmetric: user B does not see user A's run either. + response_b = await async_client.get( + "/api/v1/dag_runs/latest/", headers={"test_username": user_b} + ) + assert response_b.status_code == 200, response_b.content + returned_ids_b = {item["id"] for item in response_b.json()} + assert run_b_id in returned_ids_b + assert run_a_id not in returned_ids_b + + +@pytest.mark.asyncio +async def test_get_latest_dag_runs_with_other_users_project_id_is_rejected( + async_client: AsyncClient, db +): + """Explicitly requesting another user's ``project_id`` must be rejected, + matching how every other "get" endpoint in the module behaves. + """ + user_a = "user_individual@no_team.com" + user_b = "user_1_team_1@team1.com" + # Setting up via the same helper returns dag_template ids, but we want a + # project id user A doesn't belong to, so we go via _setup_sample_project. + project_b_id, *_ = await _setup_sample_project(async_client, user_b) + + response = await async_client.get( + f"/api/v1/dag_runs/latest/?{urlencode({'project_id': project_b_id})}", + headers={"test_username": user_a}, + ) + # The permission decorator translates a failed check into HTTP 404. + assert response.status_code == 404, response.content + + +@pytest.mark.asyncio +async def test_get_latest_dag_runs_with_own_project_id_still_works( + async_client: AsyncClient, db +): + """The legitimate case -- user asking for runs in a project they own -- + must keep returning their runs after the visibility filter is applied. + """ + user_a = "user_individual@no_team.com" + dag_template_a, _ = await _setup_dag_template(async_client, user_a) + run_id = await _create_dag_run_via_api( + async_client, user_a, dag_template_a, {"owner": "a"} + ) + + # By dag_template_id. + response = await async_client.get( + f"/api/v1/dag_runs/latest/?{urlencode({'dag_template_id': dag_template_a})}", + headers={"test_username": user_a}, + ) + assert response.status_code == 200, response.content + returned_ids = {item["id"] for item in response.json()} + assert run_id in returned_ids + + # By project_id (we resolve the project from the dag template). + project_a_id, *_ = await _setup_sample_project(async_client, user_a) + # New template in this project, new run. + dag_template_a2, _ = await _setup_dag_template(async_client, user_a) + run_id_2 = await _create_dag_run_via_api( + async_client, user_a, dag_template_a2, {"owner": "a2"} + ) + response = await async_client.get( + f"/api/v1/dag_runs/latest/?{urlencode({'project_id': project_a_id})}", + headers={"test_username": user_a}, + ) + assert response.status_code == 200, response.content + # The freshly-created project_a_id should not contain run_id (which came + # from _setup_dag_template which makes its own project), but the call must + # succeed -- the assertion that matters here is the 200, plus that we + # never get back runs from elsewhere. + returned_ids = {item["id"] for item in response.json()} + assert run_id not in returned_ids + assert run_id_2 not in returned_ids diff --git a/ui/backend/server/trackingserver_base/permissions/permissions.py b/ui/backend/server/trackingserver_base/permissions/permissions.py index 806c8cfb1..7e74a1d8b 100644 --- a/ui/backend/server/trackingserver_base/permissions/permissions.py +++ b/ui/backend/server/trackingserver_base/permissions/permissions.py @@ -18,8 +18,9 @@ import typing from common.django_utils import amap +from django.db.models import Q from ninja.errors import HttpError -from trackingserver_auth.models import APIKey, Team +from trackingserver_auth.models import APIKey, Team, User from trackingserver_base.permissions.base import allowed from trackingserver_projects.models import Project, ProjectTeamMembership, ProjectUserMembership from trackingserver_projects.schema import ProjectIn, ProjectUpdate, Visibility @@ -121,6 +122,31 @@ async def user_project_visibility(request, project: Project) -> str | None: return None +async def visible_project_ids_for_user(user: User, organizations: list[Team]) -> list[int]: + """Returns the ids of all projects the user has visibility to. + + Mirrors the membership lookup used by the projects listing endpoint, but only + returns identifiers so it can be cheaply joined into other queries (e.g. the + "latest dag runs" filter). Includes both direct user memberships and team + memberships (including the "Public" team if applicable). + + @param user: The authenticated user + @param organizations: Teams the user belongs to (from request.auth[1]) + @return: A list of project IDs the user can read or write + """ + public_team = await Team.objects.aget(name="Public") + team_ids = [org.id for org in organizations] + [public_team.id] + return [ + project_id + async for project_id in Project.objects.filter( + Q(projectusermembership__user_id=user.id) + | Q(projectteammembership__team_id__in=team_ids) + ) + .distinct() + .values_list("id", flat=True) + ] + + async def user_can_get_project_by_id(request, project_id: int) -> tuple[bool, str]: """Checks if a user can get a project by ID, given a certain ID @@ -281,13 +307,16 @@ async def user_can_get_latest_dag_runs( project_id: int = None, dag_template_id: int = None, ): - """Tells whether or not the user can get a list of latest DAG runs. This is true - if the user can get the associated project, project version, or DAG template - (the one that is specified by the request). + """Tells whether or not the user can get a list of latest DAG runs. + + When a specific ``project_id`` or ``dag_template_id`` is supplied we verify + the caller can read it. When neither is supplied this returns ``(True, "")`` + and the endpoint MUST scope the query to the caller's visible projects (see + the visibility contract in the module docstring above). The endpoint body + relies on ``visible_project_ids_for_user`` for that scoping. @param request: Django request @param project_id: Project ID in question - @param project_version_id: Project version ID in question @param dag_template_id: DAG template ID in question @return: Tuple of (can_get, error_message) """ diff --git a/ui/backend/server/trackingserver_run_tracking/api.py b/ui/backend/server/trackingserver_run_tracking/api.py index 4c36448ff..ee30eb314 100644 --- a/ui/backend/server/trackingserver_run_tracking/api.py +++ b/ui/backend/server/trackingserver_run_tracking/api.py @@ -30,6 +30,7 @@ user_can_get_project_by_id, user_can_write_to_dag_run, user_can_write_to_dag_template, + visible_project_ids_for_user, ) from trackingserver_run_tracking.models import DAGRun, NodeRun, NodeRunAttribute from trackingserver_run_tracking.schema import ( @@ -79,9 +80,10 @@ async def get_latest_dag_runs( limit: int = 100, offset: int = 0, ) -> list[DAGRunOut]: - """Gets a list of DAG runs. This accepts one (and only one of) the following: - - project_id - - dag_template_id + """Gets a list of DAG runs. Accepts at most one of ``project_id`` or + ``dag_template_id``. When neither is supplied the result is scoped to the + set of projects the caller has visibility to, per the visibility contract + documented in ``trackingserver_base/permissions/permissions.py``. @param project_id: @param dag_template_id: @@ -106,6 +108,17 @@ async def get_latest_dag_runs( f"Can only specify one of project_id/dag_template_id, " f"got: project_id={project_id}, dag_template_id={dag_template_id}", ) + # When no specific identifier is provided we still need to honor the + # "get endpoints only return items the caller has visibility to" contract. + # Scope the query to projects the caller is a member of (directly or via a + # team, including the Public team). If the caller has no visible projects + # this short-circuits to an empty result. + if not key_kwargs: + visible_project_ids = await visible_project_ids_for_user(user, teams) + if not visible_project_ids: + logger.info(f"User {user.email} has no visible projects; returning no DAG runs") + return [] + key_kwargs["dag_template__project_id__in"] = visible_project_ids dag_runs = await django_utils.amap( DAGRunOut.create_with_username, DAGRun.objects.filter(**key_kwargs, **{"dag_template__is_active": True})