diff --git a/apps/dashboard/src/pages/Compute/NodeDetail.test.tsx b/apps/dashboard/src/pages/Compute/NodeDetail.test.tsx index 2f0db2ce..139cdd8f 100644 --- a/apps/dashboard/src/pages/Compute/NodeDetail.test.tsx +++ b/apps/dashboard/src/pages/Compute/NodeDetail.test.tsx @@ -285,6 +285,25 @@ describe("NodeDetail", () => { }); }); + // A shell on a worker is a root shell, so it needs node:shell, not just + // the deployment permissions every member holds. + it("hides the Shell tab and redirects a deep link without node:shell", async () => { + const authCtx = await import("@/context/AuthContext"); + (authCtx.useAuth as ReturnType).mockReturnValue({ + hasPermission: (p: string) => p !== "node:shell", + user: { user_id: "u1", username: "test", permissions: [] }, + organizations: [], + }); + + renderNodeDetail("/dashboard/compute/pools/p1/nodes/n1/shell"); + + await waitFor(() => { + expect(screen.getAllByText("worker-1").length).toBeGreaterThan(0); + }); + expect(screen.queryByText("Shell")).not.toBeInTheDocument(); + expect(screen.queryByTestId("node-shell-stub")).not.toBeInTheDocument(); + }); + // ── Logs tab ────────────────────────────────────────────────────────────── it("shows NodeLogs stub on logs tab", async () => { const user = userEvent.setup(); diff --git a/apps/dashboard/src/pages/Compute/NodeDetail.tsx b/apps/dashboard/src/pages/Compute/NodeDetail.tsx index 9e1c8c7f..dc071487 100644 --- a/apps/dashboard/src/pages/Compute/NodeDetail.tsx +++ b/apps/dashboard/src/pages/Compute/NodeDetail.tsx @@ -37,6 +37,7 @@ export default function NodeDetail() { const { hasPermission } = useAuth(); const canRetry = hasPermission("deployment:create"); const canDelete = hasPermission("deployment:delete"); + const canShell = hasPermission("node:shell"); const [node, setNode] = useState(null); const [summary, setSummary] = useState(null); @@ -188,7 +189,7 @@ export default function NodeDetail() { { label: "Instance Details", value: "instance-details", hidden: !isNosana }, { label: "EC2 Details", value: "ec2", hidden: !isAws }, { label: "Metrics", value: "metrics" }, - { label: "Shell", value: "shell", hidden: isDepin }, + { label: "Shell", value: "shell", hidden: isDepin || !canShell }, { label: "Logs", value: "logs" }, ]; @@ -377,14 +378,18 @@ export default function NodeDetail() { - - + canShell ? ( + + + + ) : ( + + ) } /> diff --git a/apps/dashboard/src/pages/Compute/PoolDetail.test.tsx b/apps/dashboard/src/pages/Compute/PoolDetail.test.tsx index f6a60fdf..e567b6da 100644 --- a/apps/dashboard/src/pages/Compute/PoolDetail.test.tsx +++ b/apps/dashboard/src/pages/Compute/PoolDetail.test.tsx @@ -287,6 +287,34 @@ describe("PoolDetail", () => { expect(screen.queryByTestId("delete-node-node-abc")).not.toBeInTheDocument(); }); + // A shell on a worker is a root shell, so only node:shell holders see it. + it.each([ + [true, true], + [false, false], + ])("shows the per-row Shell link only with node:shell (has it: %s)", async (granted, shown) => { + const authCtx = await import("@/context/AuthContext"); + (authCtx.useAuth as ReturnType).mockReturnValue({ + hasPermission: (p: string) => (p === "node:shell" ? granted : true), + user: { org_id: "org-1", user_id: "u1", username: "test", permissions: [] }, + organizations: [], + }); + + const user = userEvent.setup(); + const nodeService = await import("@/services/nodeService"); + (nodeService.listNodes as ReturnType).mockResolvedValue([MOCK_NODE]); + + renderPoolDetail(); + await waitForPoolLoad(); + await user.click(screen.getByRole("button", { name: "Nodes" })); + + await waitFor(() => { + expect(screen.getByText("worker-1")).toBeInTheDocument(); + }); + const link = screen.queryByRole("link", { name: /shell/i }); + if (shown) expect(link).toBeInTheDocument(); + else expect(link).not.toBeInTheDocument(); + }); + // ── Per-row Delete: confirm true → deleteNode called + node list refetched ─ it("calls deleteNode then refetches the node list on a confirmed row delete", async () => { const user = userEvent.setup(); diff --git a/apps/dashboard/src/pages/Compute/PoolDetail.tsx b/apps/dashboard/src/pages/Compute/PoolDetail.tsx index ce9df8b8..2e7fbe0b 100644 --- a/apps/dashboard/src/pages/Compute/PoolDetail.tsx +++ b/apps/dashboard/src/pages/Compute/PoolDetail.tsx @@ -62,6 +62,7 @@ function PoolDetailContent() { const location = useLocation(); const { hasPermission } = useAuth(); const canDelete = hasPermission("deployment:delete"); + const canShell = hasPermission("node:shell"); // A pool is created and then linked straight to /pools/:id/workers, so the // tab has to survive a deep link rather than always opening on Overview. @@ -374,6 +375,7 @@ function PoolDetailContent() { nodes={nodes} poolId={id ?? ""} canDelete={canDelete} + canShell={canShell} onRefetch={fetchNodes} /> @@ -652,11 +654,13 @@ function NodeList({ nodes, poolId, canDelete, + canShell, onRefetch, }: { nodes: NodeView[]; poolId: string; canDelete: boolean; + canShell: boolean; onRefetch: () => void | Promise; }) { const [deletingId, setDeletingId] = useState(null); @@ -763,13 +767,15 @@ function NodeList({ > Status - e.stopPropagation()} - className="inline-flex items-center gap-1 text-xs text-muted-foreground hover:text-foreground" - > - Shell - + {canShell && ( + e.stopPropagation()} + className="inline-flex items-center gap-1 text-xs text-muted-foreground hover:text-foreground" + > + Shell + + )} e.stopPropagation()} diff --git a/src/api_gateway/gateway/proxy_routes.py b/src/api_gateway/gateway/proxy_routes.py index dcff634b..488a2515 100644 --- a/src/api_gateway/gateway/proxy_routes.py +++ b/src/api_gateway/gateway/proxy_routes.py @@ -574,11 +574,11 @@ async def proxy_admin_workers_logs(websocket: WebSocket, node_id: str): async def proxy_admin_workers_shell(websocket: WebSocket, node_id: str): """Proxy interactive shell WS to the dashboard. - Same auth pattern as /logs but requires DEPLOYMENT_UPDATE — a shell - can mutate the running container, so we gate it more tightly. + Same auth pattern as /logs but requires NODE_SHELL: a root shell on the + worker, which the default member role does not carry. """ await _proxy_admin_workers_ws(websocket, node_id, subpath="shell", - permission=PermissionEnum.DEPLOYMENT_UPDATE) + permission=PermissionEnum.NODE_SHELL) async def _proxy_admin_workers_ws(websocket: WebSocket, node_id: str, *, subpath: str, permission): diff --git a/src/api_gateway/rbac/initialization.py b/src/api_gateway/rbac/initialization.py index cf6b4212..e91d47c4 100644 --- a/src/api_gateway/rbac/initialization.py +++ b/src/api_gateway/rbac/initialization.py @@ -12,6 +12,22 @@ logger = logging.getLogger(__name__) +# Default dashboard operator permissions. Admins get every permission. +MEMBER_PERMISSIONS = [ + PermissionEnum.DEPLOYMENT_LIST.value, + PermissionEnum.DEPLOYMENT_CREATE.value, + PermissionEnum.DEPLOYMENT_UPDATE.value, + PermissionEnum.DEPLOYMENT_DELETE.value, + PermissionEnum.MODEL_ACCESS.value, + PermissionEnum.API_KEY_LIST.value, + PermissionEnum.API_KEY_CREATE.value, + PermissionEnum.API_KEY_REVOKE.value, + PermissionEnum.ORG_VIEW.value, + PermissionEnum.MODEL_LIST.value, + PermissionEnum.MODEL_ADD.value, + PermissionEnum.MODEL_DELETE.value, +] + async def initialize_default_org(db: AsyncSession): """ Initialize the default organization, roles, permissions, and superadmin if they don't exist. @@ -36,20 +52,7 @@ async def initialize_default_org(db: AsyncSession): db.add(admin_role) # Member Role: Default dashboard operator permissions - member_permissions = [ - PermissionEnum.DEPLOYMENT_LIST.value, - PermissionEnum.DEPLOYMENT_CREATE.value, - PermissionEnum.DEPLOYMENT_UPDATE.value, - PermissionEnum.DEPLOYMENT_DELETE.value, - PermissionEnum.MODEL_ACCESS.value, - PermissionEnum.API_KEY_LIST.value, - PermissionEnum.API_KEY_CREATE.value, - PermissionEnum.API_KEY_REVOKE.value, - PermissionEnum.ORG_VIEW.value, - PermissionEnum.MODEL_LIST.value, - PermissionEnum.MODEL_ADD.value, - PermissionEnum.MODEL_DELETE.value, - ] + member_permissions = list(MEMBER_PERMISSIONS) member_role_stmt = select(Role).where(Role.name == "member") member_role_res = await db.execute(member_role_stmt) diff --git a/src/api_gateway/rbac/middleware.py b/src/api_gateway/rbac/middleware.py index 92f32584..a9cc8d20 100644 --- a/src/api_gateway/rbac/middleware.py +++ b/src/api_gateway/rbac/middleware.py @@ -246,12 +246,8 @@ async def auth_middleware(request: Request, call_next): - oidc: try local (superadmin recovery), fall back to _resolve_oidc_token - local: _resolve_local_token only """ - # Skip auth for WebSocket connections - they handle auth differently (via query params) - # Check both 'upgrade' header and the connection type - upgrade_header = request.headers.get("upgrade", "").lower() - connection_header = request.headers.get("connection", "").lower() - if upgrade_header == "websocket" or "upgrade" in connection_header: - return await call_next(request) + # No Upgrade/Connection exemption: an HTTP middleware never sees a real + # handshake, so only a hand-set header could reach one. # Path relative to the ASGI root_path. Under the unified app the gateway is # mounted at /api, so Starlette sets scope['root_path']='/api' but does NOT diff --git a/src/api_gateway/rbac/permissions.py b/src/api_gateway/rbac/permissions.py index 6a541faa..c8452e96 100644 --- a/src/api_gateway/rbac/permissions.py +++ b/src/api_gateway/rbac/permissions.py @@ -60,7 +60,11 @@ def canonical_permissions() -> List[str]: PermissionEnum.ROLE_DELETE.value, ), "inferiallm:org:read": (PermissionEnum.ORG_VIEW.value,), - "inferiallm:org:write": (PermissionEnum.ORG_UPDATE.value,), + # node:shell sits here because only the admin role holds org:write. + "inferiallm:org:write": ( + PermissionEnum.ORG_UPDATE.value, + PermissionEnum.NODE_SHELL.value, + ), "inferiallm:audit:read": (PermissionEnum.AUDIT_LOG_LIST.value,), } diff --git a/src/api_gateway/schemas/auth.py b/src/api_gateway/schemas/auth.py index 29b4eb37..3169005b 100644 --- a/src/api_gateway/schemas/auth.py +++ b/src/api_gateway/schemas/auth.py @@ -42,6 +42,9 @@ class PermissionEnum(str, Enum): ORG_VIEW = "organization:view" ORG_UPDATE = "organization:update" + # Nodes. A root shell on a worker, so not in the default member role. + NODE_SHELL = "node:shell" + class OrganizationBasicInfo(BaseModel): id: str diff --git a/src/common/middleware.py b/src/common/middleware.py index 70464e0c..974bd9ce 100644 --- a/src/common/middleware.py +++ b/src/common/middleware.py @@ -66,12 +66,8 @@ async def internal_auth_middleware(request: Request, call_next): response.headers["X-Request-ID"] = request_id return response - # WebSocket handshake check - browsers can't set custom headers for WS - upgrade_header = request.headers.get("upgrade", "").lower() - if upgrade_header == "websocket": - logger.info(f"Allowing WebSocket handshake for {path}") - response = await call_next(request) - return response + # No Upgrade-header exemption: an HTTP middleware never sees a real + # handshake, so only a hand-set header could reach one. # 3. Only check if path prefix matches (if provided) if check_path_prefix and not path.startswith(check_path_prefix): diff --git a/src/orchestration/api/admin_workers.py b/src/orchestration/api/admin_workers.py index 33f4322b..7a7a6223 100644 --- a/src/orchestration/api/admin_workers.py +++ b/src/orchestration/api/admin_workers.py @@ -448,6 +448,28 @@ async def _run_until_either(*coros) -> None: logger.warning("_run_until_either crashed: %s", e) +async def _node_in_caller_org(ws: WebSocket, node_id: str) -> bool: + """Whether the node sits in a pool owned by the caller's organisation. + + The gateway checks the user's permission, not which nodes it covers, so a + node id from another organisation must be refused here. + """ + org_id = ws.headers.get("x-organization-id") or "" + if not org_id or _deps.inventory_repo is None or _deps.pool_repo is None: + return False + try: + node = await _deps.inventory_repo.get_node_by_id(node_id) + pool = await _deps.pool_repo.get(node["pool_id"]) if node else None + except Exception as e: + logger.warning("org check for node %s failed: %s", node_id, e) + return False + return ( + pool is not None + and pool["owner_type"] == "organization" + and str(pool["owner_id"]) == org_id + ) + + @router.websocket("/{node_id}/logs") async def proxy_worker_logs(ws: WebSocket, node_id: str): """Multiplex a live container logs stream over the worker's control WS. @@ -456,6 +478,9 @@ async def proxy_worker_logs(ws: WebSocket, node_id: str): bodies into the dashboard wire format, and tears the stream down on either side disconnect. """ + if not await _node_in_caller_org(ws, node_id): + await ws.close(code=1008) + return await ws.accept() stream_id = str(uuid.uuid4()) @@ -536,6 +561,9 @@ async def proxy_worker_shell(ws: WebSocket, node_id: str): and the worker's ShellOutput / ShellExit / ShellError frames back into the dashboard wire format. """ + if not await _node_in_caller_org(ws, node_id): + await ws.close(code=1008) + return await ws.accept() # DePIN nodes (Nosana/Akash) run the model in a confidential, externally diff --git a/src/orchestration/middleware.py b/src/orchestration/middleware.py index 4560ab0f..1faf7890 100644 --- a/src/orchestration/middleware.py +++ b/src/orchestration/middleware.py @@ -5,22 +5,42 @@ logger = logging.getLogger(__name__) +# Paths that bypass the internal key. Each one is health or discovery, or +# authenticates its caller itself (workers carry their own token). +INTERNAL_AUTH_SKIP_PATHS = [ + "/health", + "/v1/workers/register", + "/v1/workers/channel", + "/v3/discovery:clusters", + "/v3/discovery:endpoints", + "/v3/healthz", + "/v3/route-table", + "/v3/debug/resources", +] + class InternalAuthMiddleware: def __init__(self, app, internal_api_key: str, skip_paths: list = None): self.app = app self.internal_api_key = internal_api_key self.skip_paths = [p.rstrip("/") for p in (skip_paths or [])] - async def __call__(self, scope, receive, send): - if scope["type"] != "http" and scope["type"] != "websocket": - await self.app(scope, receive, send) + async def _reject(self, scope, receive, send, status_code: int, detail: str): + """Refuse a request in the shape its protocol expects.""" + if scope["type"] == "websocket": + # Sent before any accept, so the handshake itself fails. + await send({"type": "websocket.close", "code": 1008}) return + response = JSONResponse(status_code=status_code, content={"detail": detail}) + await response(scope, receive, send) - # 1. Skip validation for WebSockets (browsers can't send headers) - if scope["type"] == "websocket": + async def __call__(self, scope, receive, send): + if scope["type"] != "http" and scope["type"] != "websocket": await self.app(scope, receive, send) return + # WebSockets are checked like anything else. A browser cannot set the + # header, but a browser never reaches this port: the gateway mints its + # own upstream connection and does send it. path = scope.get("path", "") norm_path = path.rstrip("/") if path != "/" else path @@ -37,20 +57,18 @@ async def __call__(self, scope, receive, send): if not api_key: logger.warning(f"Unauthorized access attempt to {path}: Missing API Key") - response = JSONResponse( - status_code=status.HTTP_401_UNAUTHORIZED, - content={"detail": "Missing X-Internal-API-Key header"}, + await self._reject( + scope, receive, send, status.HTTP_401_UNAUTHORIZED, + "Missing X-Internal-API-Key header", ) - await response(scope, receive, send) return if api_key != self.internal_api_key: logger.warning(f"Unauthorized access attempt to {path}: Invalid API Key") - response = JSONResponse( - status_code=status.HTTP_403_FORBIDDEN, - content={"detail": "Invalid internal API key"}, + await self._reject( + scope, receive, send, status.HTTP_403_FORBIDDEN, + "Invalid internal API key", ) - await response(scope, receive, send) return await self.app(scope, receive, send) @@ -60,5 +78,5 @@ def internal_auth_middleware_factory(app): return InternalAuthMiddleware( app, internal_api_key=settings.internal_api_key, - skip_paths=["/health", "/deployment/ws"] + skip_paths=["/health"] ) diff --git a/src/orchestration/server.py b/src/orchestration/server.py index 75cd33c1..49ff3dd2 100644 --- a/src/orchestration/server.py +++ b/src/orchestration/server.py @@ -110,7 +110,7 @@ ) from orchestration.infra.redis_event_bus import RedisEventBus -from orchestration.middleware import InternalAuthMiddleware +from orchestration.middleware import INTERNAL_AUTH_SKIP_PATHS, InternalAuthMiddleware from orchestration.grpc_auth_interceptor import ( InternalAPIKeyInterceptor, ) @@ -442,17 +442,7 @@ async def _check(_authorization=None): app.add_middleware( InternalAuthMiddleware, internal_api_key=settings.internal_api_key, - skip_paths=[ - "/health", - "/deployment/ws", - "/v1/workers/register", - "/v1/workers/channel", - "/v3/discovery:clusters", - "/v3/discovery:endpoints", - "/v3/healthz", - "/v3/route-table", - "/v3/debug/resources", - ], + skip_paths=INTERNAL_AUTH_SKIP_PATHS, ) # Register standard exception handlers diff --git a/src/tests/api_gateway/test_auth_middleware_upgrade.py b/src/tests/api_gateway/test_auth_middleware_upgrade.py new file mode 100644 index 00000000..6009287a --- /dev/null +++ b/src/tests/api_gateway/test_auth_middleware_upgrade.py @@ -0,0 +1,34 @@ +"""The user-auth middleware must not trust Upgrade or Connection headers. + +It runs as an HTTP middleware, so a real websocket handshake never reaches it; +a request that carries these headers here is a plain one. With the exemption +in place, the request passed the middleware and reached body validation (422) +instead of being refused. +""" +import pytest +from fastapi.testclient import TestClient + +from api_gateway.app import app + +client = TestClient(app, raise_server_exceptions=False) + + +@pytest.mark.parametrize( + "headers", + [ + {"Upgrade": "websocket"}, + {"Upgrade": "WebSocket"}, + {"Connection": "upgrade"}, + {"Connection": "keep-alive, Upgrade"}, + ], +) +@pytest.mark.parametrize( + "method, path", + [ + ("POST", "/management/users"), + ("POST", "/management/api-keys"), + ("PUT", "/admin/users/x/role"), + ], +) +def test_upgrade_headers_do_not_skip_user_auth(method, path, headers): + assert client.request(method, path, headers=headers).status_code == 401 diff --git a/src/tests/api_gateway/test_internal_auth_bypass.py b/src/tests/api_gateway/test_internal_auth_bypass.py index df9bb7af..f81eb88e 100644 --- a/src/tests/api_gateway/test_internal_auth_bypass.py +++ b/src/tests/api_gateway/test_internal_auth_bypass.py @@ -138,6 +138,49 @@ async def _run(): asyncio.run(_run()) +class TestUpgradeHeaderCannotBypass: + """An Upgrade header is client-supplied and must not exempt anything. + + A real handshake arrives on the websocket scope and never reaches an HTTP + middleware, so a request that does carry the header here is a plain one. + """ + + GOOD = "valid-key-32chars-aaaaaaaaaaaaaaaa" + + def _get(self, headers): + app = _make_test_app(internal_api_key=self.GOOD) + + async def _run(): + async with AsyncClient( + transport=ASGITransport(app=app, raise_app_exceptions=False), + base_url="http://test", + ) as client: + return await client.get("/internal/secret", headers=headers) + + return asyncio.run(_run()) + + def test_upgrade_header_alone_does_not_grant_access(self): + assert self._get({"Upgrade": "websocket"}).status_code == 401 + + def test_upgrade_header_is_not_case_sensitive_either(self): + assert self._get({"Upgrade": "WebSocket"}).status_code == 401 + + def test_upgrade_header_does_not_rescue_a_wrong_key(self): + resp = self._get( + {"Upgrade": "websocket", "X-Internal-API-Key": "wrong-key"}, + ) + assert resp.status_code == 403 + + def test_a_valid_key_still_works_with_the_header_present(self): + resp = self._get( + { + "Upgrade": "websocket", + "X-Internal-API-Key": self.GOOD, + }, + ) + assert resp.status_code == 200 + + class TestInternalKeyValidation: """Verify correct API key validation behavior.""" diff --git a/src/tests/api_gateway/test_node_shell_permission.py b/src/tests/api_gateway/test_node_shell_permission.py new file mode 100644 index 00000000..5812e7da --- /dev/null +++ b/src/tests/api_gateway/test_node_shell_permission.py @@ -0,0 +1,85 @@ +"""A worker shell is a root shell, so it needs its own permission. + +Members get DEPLOYMENT_UPDATE by default, which the shell used to require. +NODE_SHELL is admin-only in every auth mode: local admins hold every +permission, and in the catalog modes it comes only from the admin-only +``inferiallm:org:write`` key. +""" +from types import SimpleNamespace + +import pytest +from fastapi.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +from api_gateway.app import app +from api_gateway.gateway import proxy_routes +from api_gateway.models import UserContext +from api_gateway.rbac.catalog import CATALOG +from api_gateway.rbac.initialization import MEMBER_PERMISSIONS +from api_gateway.rbac.permissions import ( + canonical_permissions, + expand_catalog_permissions, +) +from api_gateway.schemas.auth import PermissionEnum + +SHELL = PermissionEnum.NODE_SHELL.value + + +def _catalog_role(name): + return next(r for r in CATALOG.roles if r.name == name) + + +def test_local_admins_hold_it_and_members_do_not(): + assert SHELL in canonical_permissions() + assert SHELL not in MEMBER_PERMISSIONS + + +def test_catalog_admins_hold_it(): + assert SHELL in expand_catalog_permissions(_catalog_role("admin").permissions) + + +@pytest.mark.parametrize("role", ["member", "viewer"]) +def test_catalog_non_admins_do_not(role): + assert SHELL not in expand_catalog_permissions(_catalog_role(role).permissions) + + +def _user(*permissions): + return UserContext( + user_id="u1", username="a@b.com", email="a@b.com", roles=["member"], + permissions=list(permissions), org_id="org1", + quota_limit=10000, quota_used=0, + ) + + +def _no_upstream(*_a, **_kw): + raise OSError("no orchestration in this test") + + +@pytest.fixture +def connect_shell(monkeypatch): + def _connect(user): + async def _resolve(_token): + return user + + monkeypatch.setattr(proxy_routes, "_get_ws_user_context", _resolve) + # Replace the module reference, not its attribute: another test swaps + # the websockets package for a stub, which breaks its lazy imports. + monkeypatch.setattr( + proxy_routes, "websockets", SimpleNamespace(connect=_no_upstream), + ) + with TestClient(app).websocket_connect( + "/v1/admin/workers/n1/shell?access_token=t", + ): + pass + + return _connect + + +def test_deployment_update_alone_no_longer_opens_a_shell(connect_shell): + with pytest.raises(WebSocketDisconnect): + connect_shell(_user(PermissionEnum.DEPLOYMENT_UPDATE.value)) + + +def test_the_shell_permission_gets_past_the_gateway(connect_shell): + # Accepted; the relay then fails only because no orchestration runs here. + connect_shell(_user(SHELL)) 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 a8dda6a2..005c4212 100644 --- a/src/tests/orchestration/api/test_admin_workers_ws_proxy.py +++ b/src/tests/orchestration/api/test_admin_workers_ws_proxy.py @@ -53,6 +53,14 @@ NODE_ID = "11111111-2222-3333-4444-555555555555" +ORG_ID = "org-1" +POOL_ID = "pool-1" + + +def _client(app: FastAPI, org_id: str | None = ORG_ID) -> TestClient: + """A client carrying the org header the gateway forwards.""" + headers = {"X-Organization-ID": org_id} if org_id is not None else {} + return TestClient(app, headers=headers) # --------------------------------------------------------------------------- @@ -169,14 +177,18 @@ def wait_closed(self, deadline_s: float = 2.0) -> None: def _make_app(registry: FakeRegistry) -> FastAPI: """Build a fresh FastAPI app with admin_workers configured to use the - supplied registry. The proxy only touches ``worker_registry`` for - shell/logs — other deps are left as no-op stubs.""" + supplied registry. The node sits in a pool owned by ORG_ID, so the + org check passes for ``_client``.""" app = FastAPI() admin_workers._deps.worker_auth = None admin_workers._deps.worker_registry = registry - admin_workers._deps.inventory_repo = None - admin_workers._deps.pool_repo = None + admin_workers._deps.inventory_repo = _FakeInventory( + {"provider": "worker", "pool_id": POOL_ID}, + ) + admin_workers._deps.pool_repo = _FakePools( + {POOL_ID: {"owner_type": "organization", "owner_id": ORG_ID}}, + ) admin_workers._deps.control_plane_external_url = "" admin_workers._deps.require_permission = lambda perm: lambda *_a, **_kw: True admin_workers._deps.db_pool = None @@ -209,7 +221,7 @@ def reset_deps(): ) = saved -def _stream_id(registry: FakeRegistry, deadline_s: float = 2.0) -> str: +def _stream_id(registry: FakeRegistry, deadline_s: float = 10.0) -> str: """Return the stream_id the proxy minted for the open call. The proxy's call to ``open_*_stream`` happens on the TestClient's @@ -237,7 +249,7 @@ def test_shell_open_passes_query_params_into_body(): app = _make_app(reg) qs = "deployment=dep-1&container=ctr-2&shell=/bin/bash&user=root&cols=120&rows=40" - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell?{qs}", ) as ws: # Push an exit to terminate the worker_to_dashboard loop and let the @@ -267,7 +279,7 @@ def test_shell_stdin_translates_to_shell_input(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell?shell=/bin/sh", ) as ws: ws.send_json({"type": "stdin", "data": "ls -la\n"}) @@ -293,7 +305,7 @@ def test_shell_output_translates_to_dashboard_wire_format(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: sid = _stream_id(reg) @@ -314,7 +326,7 @@ def test_shell_resize_forwarded_as_shell_resize_envelope(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: ws.send_json({"type": "resize", "rows": 50, "cols": 132}) @@ -336,7 +348,7 @@ def test_shell_error_frame_closes_dashboard_with_error_envelope(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: sid = _stream_id(reg) @@ -359,7 +371,7 @@ def test_shell_worker_offline_reports_clean_error(): reg = FakeRegistry(offline=True) app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: msg = ws.receive_json() @@ -379,7 +391,7 @@ def test_shell_dashboard_disconnect_closes_stream(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: sid = _stream_id(reg) @@ -401,7 +413,7 @@ def test_shell_resize_missing_fields_safely_dropped(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: ws.send_json({"type": "resize"}) # nothing @@ -429,7 +441,7 @@ def test_shell_stdin_non_string_data_dropped(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: ws.send_json({"type": "stdin", "data": 12345}) # number, not str @@ -450,7 +462,7 @@ def test_shell_unknown_frame_type_ignored(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: ws.send_json({"type": "ping"}) # not in our schema @@ -472,7 +484,7 @@ def test_shell_default_cols_rows_when_missing(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: sid = _stream_id(reg) @@ -495,7 +507,7 @@ def test_logs_open_passes_deployment_and_container_into_body(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/logs?deployment=dep-7&container=ctr-3", ) as ws: sid = _stream_id(reg) @@ -515,7 +527,7 @@ def test_logs_multiple_lines_arrive_in_order(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/logs", ) as ws: sid = _stream_id(reg) @@ -540,7 +552,7 @@ def test_logs_worker_offline_reports_clean_error(): reg = FakeRegistry(offline=True) app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/logs", ) as ws: msg = ws.receive_json() @@ -560,7 +572,7 @@ def test_logs_worker_disconnect_surfaces_as_error(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/logs", ) as ws: sid = _stream_id(reg) @@ -580,7 +592,7 @@ def test_logs_dashboard_disconnect_closes_stream(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/logs", ) as ws: sid = _stream_id(reg) @@ -597,7 +609,7 @@ def test_logs_empty_query_params_yield_blank_targets(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/logs", ) as ws: sid = _stream_id(reg) @@ -622,7 +634,7 @@ def test_shell_logs_unexpected_frame_logged_and_ignored(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: sid = _stream_id(reg) @@ -641,7 +653,7 @@ def test_logs_unexpected_frame_logged_and_ignored(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/logs", ) as ws: sid = _stream_id(reg) @@ -666,12 +678,22 @@ async def get_node_by_id(self, node_id): return self._node +class _FakePools: + def __init__(self, pools): + self._pools = pools + + async def get(self, pool_id): + return self._pools.get(pool_id) + + def test_shell_depin_node_gated_with_clear_message(): reg = FakeRegistry() app = _make_app(reg) - admin_workers._deps.inventory_repo = _FakeInventory({"provider": "nosana"}) + admin_workers._deps.inventory_repo = _FakeInventory( + {"provider": "nosana", "pool_id": POOL_ID}, + ) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: msg = ws.receive_json() @@ -691,9 +713,11 @@ def test_shell_worker_node_not_gated(): DePIN gate must not fire for it.""" reg = FakeRegistry() app = _make_app(reg) - admin_workers._deps.inventory_repo = _FakeInventory({"provider": "worker"}) + admin_workers._deps.inventory_repo = _FakeInventory( + {"provider": "worker", "pool_id": POOL_ID}, + ) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell", ) as ws: sid = _stream_id(reg) @@ -709,7 +733,7 @@ def test_shell_malformed_cols_query_param_falls_back_to_default(): reg = FakeRegistry() app = _make_app(reg) - with TestClient(app).websocket_connect( + with _client(app).websocket_connect( f"/v1/admin/workers/{NODE_ID}/shell?cols=lol&rows=&shell=/bin/sh", ) as ws: sid = _stream_id(reg) @@ -719,3 +743,48 @@ def test_shell_malformed_cols_query_param_falls_back_to_default(): _, body = reg.opens[0] assert body.cols == 0 # fallback assert body.rows == 0 # empty string also falls back + + +# --------------------------------------------------------------------------- +# Organisation scoping. The gateway checks permission, not which nodes it +# covers, so another org's node must be refused here, before any stream opens. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("kind", ["shell", "logs"]) +@pytest.mark.parametrize( + "org_id, pool", + [ + ("org-2", {"owner_type": "organization", "owner_id": ORG_ID}), + (None, {"owner_type": "organization", "owner_id": ORG_ID}), + ("", {"owner_type": "organization", "owner_id": ORG_ID}), + (ORG_ID, {"owner_type": "system", "owner_id": ORG_ID}), + (ORG_ID, None), + ], + ids=["other-org", "no-header", "empty-header", "system-pool", "no-pool"], +) +def test_a_node_outside_the_callers_org_is_refused(kind, org_id, pool): + reg = FakeRegistry() + app = _make_app(reg) + admin_workers._deps.pool_repo = _FakePools({POOL_ID: pool} if pool else {}) + + with pytest.raises(WebSocketDisconnect): + with _client(app, org_id).websocket_connect( + f"/v1/admin/workers/{NODE_ID}/{kind}", + ): + pass # refused at the handshake, so entering raises + assert reg.opens == [] + + +@pytest.mark.parametrize("kind", ["shell", "logs"]) +def test_an_unknown_node_is_refused(kind): + reg = FakeRegistry() + app = _make_app(reg) + admin_workers._deps.inventory_repo = _FakeInventory(None) + + with pytest.raises(WebSocketDisconnect): + with _client(app).websocket_connect( + f"/v1/admin/workers/{NODE_ID}/{kind}", + ): + pass + assert reg.opens == [] diff --git a/src/tests/orchestration/test_internal_auth_websockets.py b/src/tests/orchestration/test_internal_auth_websockets.py new file mode 100644 index 00000000..97f7b9ff --- /dev/null +++ b/src/tests/orchestration/test_internal_auth_websockets.py @@ -0,0 +1,84 @@ +"""The internal-key check covers websockets, not just HTTP. + +Orchestration's admin websockets (worker logs and the interactive shell) carry +no auth of their own, so the middleware is the only thing between an open port +and a root shell on a worker. A browser cannot set the header, but a browser +never reaches this port: the gateway authenticates the user, then opens its own +upstream connection with the key attached. +""" +from fastapi import FastAPI, WebSocket +from fastapi.testclient import TestClient +from starlette.websockets import WebSocketDisconnect +import pytest + +from orchestration.middleware import InternalAuthMiddleware + +KEY = "internal-key-32chars-aaaaaaaaaaaa" + + +def _app(): + app = FastAPI() + app.add_middleware( + InternalAuthMiddleware, internal_api_key=KEY, + skip_paths=["/health", "/v1/workers/channel"], + ) + + @app.websocket("/v1/admin/workers/{node_id}/shell") + async def shell(ws: WebSocket, node_id: str): + await ws.accept() + await ws.send_text("root@worker") + + @app.websocket("/v1/workers/channel") + async def channel(ws: WebSocket): + await ws.accept() + await ws.send_text("worker channel") + + @app.get("/v1/admin/workers") + async def listing(): + return {"nodes": []} + + return TestClient(app) + + +def _shell(headers=None): + with _app().websocket_connect( + "/v1/admin/workers/n1/shell", headers=headers or {}, + ) as ws: + return ws.receive_text() + + +def test_a_shell_without_the_key_is_refused(): + with pytest.raises(WebSocketDisconnect): + _shell() + + +def test_a_shell_with_the_wrong_key_is_refused(): + with pytest.raises(WebSocketDisconnect): + _shell({"X-Internal-API-Key": "wrong-key"}) + + +def test_an_upgrade_header_does_not_substitute_for_the_key(): + # The HTTP path is what a hand-set Upgrade header actually reaches. + resp = _app().get("/v1/admin/workers", headers={"Upgrade": "websocket"}) + assert resp.status_code == 401 + + +def test_the_gateway_reaches_the_shell_with_the_key(): + assert _shell({"X-Internal-API-Key": KEY}) == "root@worker" + + +@pytest.mark.parametrize( + "path", + ["/deployment/ws", "/v1/admin/workers/x/shell", "/v1/admin/workers/x/logs"], +) +def test_no_unauthenticated_stream_is_on_the_real_skip_list(path): + from orchestration.middleware import INTERNAL_AUTH_SKIP_PATHS + + assert path not in INTERNAL_AUTH_SKIP_PATHS + + +def test_the_worker_channel_still_connects_without_the_key(): + # Workers hold a worker JWT, not the internal key, and authenticate + # themselves inside the route. + with _app().websocket_connect("/v1/workers/channel") as ws: + assert ws.receive_text() == "worker channel"