diff --git a/src/orchestration/models/model_deployment/deployment_server.py b/src/orchestration/models/model_deployment/deployment_server.py index 279230de..7b50d76d 100644 --- a/src/orchestration/models/model_deployment/deployment_server.py +++ b/src/orchestration/models/model_deployment/deployment_server.py @@ -34,7 +34,7 @@ _deprovision_direct_node, ) from orchestration.provisioning.engine.base import AdapterType -from orchestration.ownership import caller_org, owns +from orchestration.ownership import caller_org, owns, owns_pool from orchestration.models.model_deployment.model_ref import ( resolve_artifact_uri, ) @@ -863,7 +863,7 @@ async def _lookup_org_id(resource_type: str, resource_id: str) -> str | None: "SELECT owner_type, owner_id FROM compute_pools WHERE id = $1", UUID(resource_id), ) - if row and row["owner_type"] == "org": + if row and row["owner_type"] in ("organization", "user"): return row["owner_id"] except Exception as e: logger.warning(f"Failed to look up org_id for {resource_type}/{resource_id}: {e}") @@ -1676,6 +1676,15 @@ async def deploy_model(req: DeployModelRequest, request: Request): worker-pools, return. Otherwise enqueue ProvisioningJob. 4. PoolAtCapacity is caught and surfaced as 503. """ + # The deployment's org is the caller's, never the body's. + org = caller_org(request) + if not org: + raise HTTPException( + status_code=403, + detail="An organization is required to deploy (CLI: --org-id)", + ) + req.org_id = org + from uuid import UUID, uuid4 from orchestration.models.model_deployment.pool_placer import ( PoolPlacer, @@ -1733,7 +1742,7 @@ async def deploy_model(req: DeployModelRequest, request: Request): except (TypeError, ValueError): raise HTTPException(status_code=422, detail="pool_id must be a UUID") pool_row = await pool_repo.get(pool_id_uuid) - if pool_row is None: + if pool_row is None or not owns_pool(request, pool_row): raise HTTPException(status_code=404, detail="pool not found") lifecycle = (pool_row.get("lifecycle_state") or "running").lower() if lifecycle in ("terminating", "terminated"): @@ -2969,11 +2978,16 @@ async def create_pool(req: CreatePoolRequest, request: Request): except InvalidRegionError as e: raise HTTPException(status_code=422, detail=str(e)) - # 3. Multi-org header override (api_gateway forwards resolved org) - hdr_org = request.headers.get("x-organization-id") - if hdr_org and req.owner_type == "user": - req.owner_id = hdr_org - req.owner_type = "organization" + # 3. The owner is the caller's org, never the body's: a body could + # otherwise record the pool under another org. + owner = caller_org(request) + if not owner: + raise HTTPException( + status_code=403, + detail="An organization is required to create a pool (CLI: --org-id)", + ) + req.owner_id = owner + req.owner_type = "organization" db_pool = request.app.state.pool pool_repo = ComputePoolRepository(db_pool) @@ -3075,7 +3089,8 @@ async def create_pool(req: CreatePoolRequest, request: Request): @router.get("/list/pool/{pool_id}/inventory") -async def list_pool_inventory(pool_id: str): +async def list_pool_inventory(pool_id: str, request: Request): + await _require_own_pool(request, pool_id) async with _auth_channel() as channel: stub = compute_pool_pb2_grpc.ComputePoolManagerStub(channel) @@ -3136,16 +3151,23 @@ def resolve_gpu_specs(allowed_gpu_types, gpu_resource_map) -> list[dict]: @router.get("/listPools/{owner_id}") async def list_pools( + request: Request, owner_id: str | None = None, limit: int = Query(100, ge=1, le=500), offset: int = Query(0, ge=0), ): + # The org comes from the gateway; the path segment is kept only so + # existing callers do not break. + owner_id = caller_org(request) + if not owner_id: + return {"pools": [], "total": 0} + async with _auth_channel() as channel: stub = compute_pool_pb2_grpc.ComputePoolManagerStub(channel) try: resp = await stub.ListPools( - compute_pool_pb2.ListPoolsRequest(owner_id=owner_id or "") + compute_pool_pb2.ListPoolsRequest(owner_id=owner_id) ) except grpc.RpcError as e: raise HTTPException(status_code=500, detail=e.details()) @@ -3238,7 +3260,8 @@ async def list_pools( @router.get("/pool/{pool_id}") -async def get_pool(pool_id: str): +async def get_pool(pool_id: str, request: Request): + await _require_own_pool(request, pool_id) async with _auth_channel() as channel: stub = compute_pool_pb2_grpc.ComputePoolManagerStub(channel) @@ -3300,6 +3323,7 @@ async def update_pool_metadata(pool_id: str, req: UpdatePoolMetadataRequest, req pool_uuid = UUID(pool_id) except ValueError: raise HTTPException(status_code=400, detail="Invalid pool_id") + await _require_own_pool(request, pool_id) conn = None try: @@ -3385,11 +3409,12 @@ async def update_pool_metadata(pool_id: str, req: UpdatePoolMetadataRequest, req @router.post("/stoppool/{pool_id}") -async def stop_pool(pool_id: str): +async def stop_pool(pool_id: str, request: Request): try: pool_uuid = UUID(pool_id) except ValueError: raise HTTPException(status_code=400, detail="Invalid pool_id") + await _require_own_pool(request, pool_id) conn = None try: @@ -3439,11 +3464,12 @@ async def stop_pool(pool_id: str): @router.post("/deletepool/{pool_id}") -async def delete_pool(pool_id: str): +async def delete_pool(pool_id: str, request: Request): try: pool_uuid = UUID(pool_id) except ValueError: raise HTTPException(status_code=400, detail="Invalid pool_id") + await _require_own_pool(request, pool_id) lifecycle_state = await _get_pool_lifecycle_state(pool_uuid) if lifecycle_state is None: @@ -3482,7 +3508,7 @@ async def delete_pool(pool_id: str): @router.delete("/pool/{pool_id}") -async def delete_pool_rest(pool_id: str): +async def delete_pool_rest(pool_id: str, request: Request): """Delete a pool from the dashboard in one call: tear down every node's EC2 (via the provisioning reconciler) and soft-delete the pool row. @@ -3521,6 +3547,7 @@ async def delete_pool_rest(pool_id: str): pool_uuid = UUID(pool_id) except ValueError: raise HTTPException(status_code=400, detail="Invalid pool_id") + await _require_own_pool(request, pool_id) conn = None try: @@ -3732,11 +3759,11 @@ async def delete_pool_rest(pool_id: str): @router.get("/listDeployments/{pool_id}") -async def list_deployments(pool_id: str | None = None): +async def list_deployments(request: Request, pool_id: str | None = None): """ - List all deployments. - Optionally filter by pool_id. + List the deployments in one of the caller's pools. """ + await _require_own_pool(request, pool_id or "") async with _auth_channel() as channel: stub = model_deployment_pb2_grpc.ModelDeploymentServiceStub(channel) @@ -3776,6 +3803,29 @@ async def _require_own_deployment(request, deployment_id: str) -> None: raise HTTPException(status_code=404, detail="Deployment not found") +async def _pool_owner(pool_id: str): + """The pool's (owner_type, owner_id) row, or None if absent or unreadable.""" + conn = None + try: + conn = await asyncpg.connect(POSTGRES_DSN) + return await conn.fetchrow( + "SELECT owner_type, owner_id FROM compute_pools WHERE id = $1", + UUID(pool_id), + ) + except Exception as e: + logger.warning("ownership lookup for pool %s failed: %s", pool_id, e) + return None + finally: + if conn: + await conn.close() + + +async def _require_own_pool(request, pool_id: str) -> None: + """404 unless the caller's org owns the pool, as for deployments.""" + if not owns_pool(request, await _pool_owner(pool_id)): + raise HTTPException(status_code=404, detail="Pool not found") + + @router.get("/logs/{deployment_id}") async def get_deployment_logs(deployment_id: str, request: Request): """ diff --git a/src/tests/orchestration/api/test_admin_workers_ws_proxy.py b/src/tests/orchestration/api/test_admin_workers_ws_proxy.py index 3e839596..6047f912 100644 --- a/src/tests/orchestration/api/test_admin_workers_ws_proxy.py +++ b/src/tests/orchestration/api/test_admin_workers_ws_proxy.py @@ -792,7 +792,7 @@ def test_an_unknown_node_is_refused(kind): @pytest.mark.parametrize("kind", ["shell", "logs"]) def test_an_older_pool_stored_as_user_typed_still_belongs_to_its_org(kind): - # createpool stores owner_type "user" with the org's id when no org is forwarded. + # createpool used to store owner_type "user" with the org's id; those rows remain. reg = FakeRegistry() app = _make_app(reg) admin_workers._deps.pool_repo = _FakePools( diff --git a/src/tests/orchestration/model_deployment/test_delete_pool_rest.py b/src/tests/orchestration/model_deployment/test_delete_pool_rest.py index 307c2bd9..8fec7b5d 100644 --- a/src/tests/orchestration/model_deployment/test_delete_pool_rest.py +++ b/src/tests/orchestration/model_deployment/test_delete_pool_rest.py @@ -64,6 +64,20 @@ async def close(self): pass +ORG = "org-1" + + +@pytest.fixture(autouse=True) +def _caller_owns_the_pool(): + # Every request here comes from the pool's own org; ownership is tested + # in test_pool_ownership.py. + with patch.object( + deployment_server, "_pool_owner", + AsyncMock(return_value={"owner_type": "organization", "owner_id": ORG}), + ): + yield + + def _app(): app = FastAPI() app.include_router(deployment_server.router) @@ -83,7 +97,7 @@ def _patches(conn): async def _delete(path): transport = ASGITransport(app=_app()) async with httpx.AsyncClient(transport=transport, base_url="http://t") as c: - return await c.delete(path) + return await c.delete(path, headers={"X-Organization-ID": ORG}) async def test_delete_pool_404_when_missing(): diff --git a/src/tests/orchestration/model_deployment/test_pool_ownership.py b/src/tests/orchestration/model_deployment/test_pool_ownership.py new file mode 100644 index 00000000..82bb6ece --- /dev/null +++ b/src/tests/orchestration/model_deployment/test_pool_ownership.py @@ -0,0 +1,279 @@ +"""Pool routes act only on the caller's own pools, and new pools and +deployments are recorded under the caller's org, never one named in the body. +""" +from __future__ import annotations + +from contextlib import asynccontextmanager +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from orchestration.models.model_deployment import deployment_server as ds +from orchestration.repositories import pool_repo as pool_repo_mod + +OWNER = "org-a" +POOL = str(uuid4()) + + +class _Reached(Exception): + """Raised by a stubbed next step: the request got past the ownership check.""" + + +ROUTES = [ + # method, path, json body, attribute that is the route's next step + ("GET", f"/deployment/list/pool/{POOL}/inventory", None, "_auth_channel"), + ("GET", f"/deployment/pool/{POOL}", None, "_auth_channel"), + ("PATCH", f"/deployment/updatepool/{POOL}", {}, "asyncpg.connect"), + ("POST", f"/deployment/stoppool/{POOL}", None, "asyncpg.connect"), + ("POST", f"/deployment/deletepool/{POOL}", None, "_get_pool_lifecycle_state"), + ("DELETE", f"/deployment/pool/{POOL}", None, "asyncpg.connect"), + ("GET", f"/deployment/listDeployments/{POOL}", None, "_auth_channel"), +] +IDS = ["inventory", "get", "update", "stop", "deletepool", "delete", "deployments"] + + +def _app(owner_row): + async def _owner(pool_id): + return owner_row if pool_id == POOL else None + + app = FastAPI() + app.include_router(ds.router) + app.state.pool = None + app.state.worker_controller = None + return app, _owner + + +def _stub_next_step(monkeypatch, attr): + reached = [] + + def _next(*_a, **_kw): + reached.append(True) + raise _Reached() + + owner, _, name = attr.rpartition(".") + monkeypatch.setattr(getattr(ds, owner) if owner else ds, name, _next) + return reached + + +def _call(app, method, path, body, org): + headers = {"X-Organization-ID": org} if org is not None else {} + client = TestClient(app, raise_server_exceptions=False) + return client.request(method, path, json=body, headers=headers) + + +def _run(monkeypatch, attr, method, path, body, org, owner_row): + app, owner = _app(owner_row) + monkeypatch.setattr(ds, "_pool_owner", owner) + reached = _stub_next_step(monkeypatch, attr) + return _call(app, method, path, body, org), reached + + +@pytest.mark.parametrize("method, path, body, attr", ROUTES, ids=IDS) +@pytest.mark.parametrize( + "org", ["org-b", "", "None", None], ids=["other-org", "empty", "None", "no-header"], +) +def test_another_orgs_pool_is_not_found(monkeypatch, method, path, body, attr, org): + r, reached = _run(monkeypatch, attr, method, path, body, org, + {"owner_type": "organization", "owner_id": OWNER}) + + # 404, not 403: the same answer as a missing pool reveals nothing. + assert r.status_code == 404, r.text + assert reached == [], "the route acted on a pool the caller does not own" + + +@pytest.mark.parametrize("method, path, body, attr", ROUTES, ids=IDS) +@pytest.mark.parametrize("owner_type", ["organization", "user"]) +def test_the_owning_org_gets_through(monkeypatch, method, path, body, attr, owner_type): + _, reached = _run(monkeypatch, attr, method, path, body, OWNER, + {"owner_type": owner_type, "owner_id": OWNER}) + + assert reached == [True] + + +@pytest.mark.parametrize("method, path, body, attr", ROUTES, ids=IDS) +def test_a_system_pool_belongs_to_no_org(monkeypatch, method, path, body, attr): + r, reached = _run(monkeypatch, attr, method, path, body, OWNER, + {"owner_type": "system", "owner_id": OWNER}) + + assert r.status_code == 404 + assert reached == [] + + +# ---- listPools ------------------------------------------------------------- + + +@pytest.fixture +def listing(monkeypatch): + asked = [] + + class _Stub: + def __init__(self, _channel): + pass + + async def ListPools(self, req): + asked.append(req.owner_id) + return SimpleNamespace(pools=[]) + + @asynccontextmanager + async def _channel(): + yield None + + monkeypatch.setattr(ds, "_auth_channel", _channel) + monkeypatch.setattr(ds.compute_pool_pb2_grpc, "ComputePoolManagerStub", _Stub) + app = FastAPI() + app.include_router(ds.router) + return app, asked + + +def test_list_pools_uses_the_callers_org_not_the_path(listing): + app, asked = listing + + _call(app, "GET", "/deployment/listPools/org-b", None, OWNER) + + assert asked == [OWNER] + + +@pytest.mark.parametrize("org", ["", "None", None], ids=["empty", "None", "no-header"]) +def test_a_caller_with_no_org_lists_no_pools(listing, org): + app, asked = listing + + r = _call(app, "GET", "/deployment/listPools/org-b", None, org) + + assert r.json() == {"pools": [], "total": 0} + assert asked == [] + + +# ---- createpool ------------------------------------------------------------ + + +def _pool_body(**over): + body = { + "pool_name": "p1", "owner_type": "organization", "owner_id": "org-b", + "provider": "on_prem", "allowed_gpu_types": ["any"], + "max_cost_per_hour": 0, "is_dedicated": False, + "provider_pool_id": "", "scheduling_policy_json": "{}", + } + body.update(over) + return body + + +@pytest.fixture +def created(monkeypatch): + stored = [] + + class _Repo: + def __init__(self, _db): + pass + + async def credential_exists(self, *_a, **_kw): + return True + + async def create_pool(self, data): + stored.append(data) + raise _Reached() + + monkeypatch.setattr(pool_repo_mod, "ComputePoolRepository", _Repo) + app = FastAPI() + app.include_router(ds.router) + app.state.pool = None + return app, stored + + +@pytest.mark.parametrize("owner_type", ["organization", "system", "user"]) +def test_a_new_pool_is_owned_by_the_caller_whatever_the_body_says(created, owner_type): + app, stored = created + + _call(app, "POST", "/deployment/createpool", _pool_body(owner_type=owner_type), OWNER) + + assert len(stored) == 1 + assert (stored[0]["owner_type"], stored[0]["owner_id"]) == ("organization", OWNER) + + +@pytest.mark.parametrize("org", ["", "None", None], ids=["empty", "None", "no-header"]) +def test_a_pool_cannot_be_created_without_an_org(created, org): + app, stored = created + + r = _call(app, "POST", "/deployment/createpool", _pool_body(), org) + + assert r.status_code == 403 + assert stored == [] + + +# ---- deploy ---------------------------------------------------------------- + + +def _deploy_body(**over): + body = {"model_name": "m", "model_version": "v1", "replicas": 1, + "gpu_per_replica": 1, "org_id": "org-b"} + body.update(over) + return body + + +@pytest.fixture +def deploying(monkeypatch): + recorded = [] + + async def _external(req, _deploys): + recorded.append(req.org_id) + return {"deployment_id": "d1", "state": "RUNNING"}, 201 + + async def _no_duplicate(*_a, **_kw): + return None + + async def _audit(*_a, **_kw): + return None + + monkeypatch.setattr(ds, "_create_external_deployment", _external) + monkeypatch.setattr(ds, "_reject_duplicate_deployment", _no_duplicate) + monkeypatch.setattr(ds, "log_audit_event", _audit) + app = FastAPI() + app.include_router(ds.router) + app.state.pool = None + app.state.worker_controller = None + return app, recorded + + +def test_a_deployment_is_recorded_under_the_callers_org(deploying): + app, recorded = deploying + + r = _call(app, "POST", "/deployment/deploy", + _deploy_body(workload_type="external"), OWNER) + + assert r.status_code == 201, r.text + assert recorded == [OWNER] + + +@pytest.mark.parametrize("org", ["", "None", None], ids=["empty", "None", "no-header"]) +def test_deploying_without_an_org_is_refused(deploying, org): + app, recorded = deploying + + r = _call(app, "POST", "/deployment/deploy", + _deploy_body(workload_type="external"), org) + + assert r.status_code == 403 + assert recorded == [] + + +@pytest.mark.parametrize("owner, expect", [ + ("org-b", 404), # another org's pool: same answer as a missing one + (OWNER, 409), # the caller's pool: past the check, refused as terminated +]) +def test_deploy_only_targets_the_callers_own_pools(monkeypatch, deploying, owner, expect): + app, _ = deploying + + class _Repo: + def __init__(self, _db): + pass + + async def get(self, _pool_id): + return {"owner_type": "organization", "owner_id": owner, + "lifecycle_state": "terminated"} + + monkeypatch.setattr(pool_repo_mod, "ComputePoolRepository", _Repo) + + r = _call(app, "POST", "/deployment/deploy", _deploy_body(pool_id=POOL), OWNER) + + assert r.status_code == expect, r.text