Bind a worker's name and labels into its token; list and revoke tokens by jti

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013GAeg4Jo6xjZzjwy9QDWMQ

backend/fluksio/api/deps.py

@@ -20,7 +20,7 @@ from fluksio.flow.artifacts import is_reference
 from fluksio.flow.controller import FlowController
 from fluksio.flow.dashboards import DashboardDef, DashboardStore
 from fluksio.flow.workers import PythonWorkerPool
-from fluksio.models import OAuthClient, TokenPayload, User
+from fluksio.models import OAuthClient, TokenPayload, User, WorkerToken
 
 reusable_oauth2 = OAuth2PasswordBearer(
     tokenUrl=f"{settings.API_V1_STR}/login/access-token"
@@ -300,6 +300,25 @@ def oauth_client_lives(client_id: str) -> bool:
         return session.get(OAuthClient, key) is not None
 
 
+def worker_claims(token: str) -> dict[str, Any]:
+    """A worker credential's claims, if it is still good. Blocking.
+
+    Raises ``InvalidTokenError`` for one that does not verify, one minted
+    before tokens carried a `jti` and their labels, and one a superuser has
+    revoked. A token whose `jti` has no row is one the engine minted for
+    itself, a provisioned machine's, and is good until it expires.
+    """
+    claims = security.decode_worker_token(token)
+    jti = claims.get("jti")
+    if not jti:
+        raise InvalidTokenError("this worker token predates revocation; mint a new one")
+    with Session(engine) as session:
+        row = session.get(WorkerToken, str(jti))
+    if row is not None and row.revoked_at is not None:
+        raise InvalidTokenError("this worker token has been revoked")
+    return claims
+
+
 def _gate_agent(payload: dict[str, Any]) -> dict[str, Any]:
     """Refuse a token whose agent has been revoked."""
     if not oauth_client_lives(str(payload.get("client_id") or "")):

backend/fluksio/api/routes/artifacts.py

@@ -17,8 +17,7 @@ from pydantic import BaseModel
 from sqlmodel import Session
 from starlette.concurrency import run_in_threadpool
 
-from fluksio.api.deps import user_from_token
-from fluksio.core import security
+from fluksio.api.deps import user_from_token, worker_claims
 from fluksio.core.config import settings
 from fluksio.core.db import engine
 from fluksio.flow.artifacts import ArtifactStore
@@ -42,7 +41,7 @@ def artifact_caller(request: Request) -> str:
     if not token:
         raise HTTPException(status_code=401, detail="Not authenticated")
     try:
-        claims = security.decode_worker_token(token)
+        claims = worker_claims(token)
     except InvalidTokenError:
         pass
     else:

backend/fluksio/api/routes/workers.py

@@ -2,24 +2,39 @@
 
 A worker dials in rather than being dialled: the GPU box and the engine are
 usually on different networks, and only one of them can be reached. It presents
-a token minted here, says what it can do, and then answers calls on the socket
-it opened.
+a token minted here, which says who it is and which labels it carries, says
+what else it can do, and then answers calls on the socket it opened.
 """
 
 from __future__ import annotations
 
 import asyncio
 import logging
-from datetime import timedelta
+import uuid
+from datetime import UTC, datetime, timedelta
 from typing import Any
 
-from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket
+from fastapi import (
+    APIRouter,
+    Depends,
+    HTTPException,
+    Request,
+    WebSocket,
+    WebSocketDisconnect,
+)
 from fastapi.responses import PlainTextResponse
 from fluksio_worker import worker_main
 from jwt.exceptions import InvalidTokenError
 from pydantic import BaseModel, Field
+from sqlmodel import col, select
+from starlette.concurrency import run_in_threadpool
 
-from fluksio.api.deps import get_current_active_superuser, get_current_user
+from fluksio.api.deps import (
+    SessionDep,
+    get_current_active_superuser,
+    get_current_user,
+    worker_claims,
+)
 from fluksio.core import security
 from fluksio.flow.remote import (
     PROTOCOL,
@@ -27,6 +42,7 @@ from fluksio.flow.remote import (
     RemoteWorkerHub,
     WorkerInventory,
 )
+from fluksio.models import WorkerToken
 
 logger = logging.getLogger(__name__)
 
@@ -35,6 +51,9 @@ router = APIRouter(prefix="/workers", tags=["workers"])
 #: Long, because a worker is a machine somebody set up once and left running.
 TOKEN_DAYS = 365
 
+#: The sockets attached with each token, by `jti`, so revoking one drops them.
+_attached: dict[str, set[WebSocket]] = {}
+
 
 class WorkerInfo(BaseModel):
     name: str
@@ -51,12 +70,16 @@ class WorkerInfo(BaseModel):
 
 
 class TokenRequest(BaseModel):
-    name: str
+    name: str = Field(min_length=1, max_length=255)
+    #: What a node's `device` may match on this worker; it cannot add to them.
+    labels: list[str] = Field(default_factory=list)
 
 
 class TokenIssued(BaseModel):
     name: str
     token: str
+    jti: str
+    labels: list[str] = Field(default_factory=list)
     expires_days: int = TOKEN_DAYS
 
 
@@ -145,14 +168,58 @@ def read_resources(request: Request) -> Any:
     response_model=TokenIssued,
     dependencies=[Depends(get_current_active_superuser)],
 )
-def issue_token(body: TokenRequest) -> Any:
+def issue_token(session: SessionDep, body: TokenRequest) -> Any:
     """Mint the credential a worker presents when it dials in.
 
-    Shown once. It is signed with the same keypair the agent tokens use, so
-    rotating that key revokes every worker along with them.
+    Shown once. The worker's name and labels are signed into it, so a worker is
+    what it was minted as whatever it says about itself, and only its `jti` is
+    kept here, to revoke it by.
     """
-    token = security.create_worker_token(body.name, timedelta(days=TOKEN_DAYS))
-    return TokenIssued(name=body.name, token=token)
+    labels = sorted({label.strip() for label in body.labels if label.strip()})
+    row = WorkerToken(
+        jti=uuid.uuid4().hex,
+        name=body.name,
+        labels=labels,
+        expires_at=datetime.now(UTC) + timedelta(days=TOKEN_DAYS),
+    )
+    token = security.create_worker_token(
+        body.name, timedelta(days=TOKEN_DAYS), labels=labels, jti=row.jti
+    )
+    session.add(row)
+    session.commit()
+    return TokenIssued(name=body.name, token=token, jti=row.jti, labels=labels)
+
+
+@router.get(
+    "/tokens",
+    response_model=list[WorkerToken],
+    dependencies=[Depends(get_current_active_superuser)],
+)
+def read_tokens(session: SessionDep) -> Any:
+    """Every worker token minted here, newest first, without the tokens."""
+    return session.exec(
+        select(WorkerToken).order_by(col(WorkerToken.created_at).desc())
+    ).all()
+
+
+@router.post(
+    "/tokens/{jti}/revoke",
+    response_model=WorkerToken,
+    dependencies=[Depends(get_current_active_superuser)],
+)
+async def revoke_token(session: SessionDep, jti: str) -> Any:
+    """Refuse this token from now on, and drop a worker attached with it."""
+    row = session.get(WorkerToken, jti)
+    if row is None:
+        raise HTTPException(status_code=404, detail="No such worker token")
+    if row.revoked_at is None:
+        row.revoked_at = datetime.now(UTC)
+        session.add(row)
+        session.commit()
+        session.refresh(row)
+    for socket in list(_attached.get(jti, ())):
+        await socket.close(code=1008)
+    return row
 
 
 @router.get(
@@ -177,14 +244,29 @@ async def attach(websocket: WebSocket, token: str = "") -> None:
     does: a websocket handshake carries no headers of its own.
     """
     try:
-        claims = security.decode_worker_token(token)
+        claims = await run_in_threadpool(worker_claims, token)
     except InvalidTokenError:
         await websocket.close(code=1008)
         return
 
+    jti = str(claims["jti"])
     await websocket.accept()
+    _attached.setdefault(jti, set()).add(websocket)
+    try:
+        await _attached_session(websocket, claims)
+    finally:
+        sockets = _attached.get(jti, set())
+        sockets.discard(websocket)
+        if not sockets:
+            _attached.pop(jti, None)
+
+
+async def _attached_session(websocket: WebSocket, claims: dict[str, Any]) -> None:
+    """The hello, then calls, until the socket goes or the token is revoked."""
     try:
         hello = await asyncio.wait_for(websocket.receive_json(), timeout=30)
+    except WebSocketDisconnect:
+        return
     except (TimeoutError, ValueError):
         await websocket.close(code=1002)
         return
@@ -196,13 +278,27 @@ async def attach(websocket: WebSocket, token: str = "") -> None:
         await websocket.close(code=1002)
         return
 
-    # The token names the worker; what it calls itself is a suggestion, so two
-    # hosts cannot fight over one identity by claiming the same name.
+    # The token names the worker and its labels. What it calls itself is a
+    # suggestion, so two hosts cannot fight over one identity by claiming the
+    # same name; labels it may repeat but not change, since they decide which
+    # node code and secrets come its way.
+    labels = sorted({str(label) for label in claims.get("labels") or []})
+    claimed = sorted({str(label) for label in hello.get("labels") or []})
+    if claimed and claimed != labels:
+        await websocket.send_json(
+            {
+                "op": "refused",
+                "reason": f"this token carries the labels {labels}, not {claimed}",
+            }
+        )
+        await websocket.close(code=1008)
+        return
+
     name = str(claims.get("sub") or hello.get("name") or "worker")
     hub = _hub(websocket.app)
     worker = RemoteWorker(
         name=name,
-        labels=[str(label) for label in (hello.get("labels") or [])],
+        labels=labels,
         send=websocket.send_json,
         loop=asyncio.get_running_loop(),
         max_parallel=max(1, int(hello.get("max_parallel") or 1)),