Add verified JWT edge assertions via JWKS

.env.example

@@ -19,6 +19,21 @@ DRUKS_LOG_LEVEL=INFO
 #   python3 -c 'import base64, os; print(base64.b64encode(os.urandom(32)).decode())'
 DRUKS_SECRETS_KEY=
 
+# Browser identity. `none` = no authentication (loopback dashboard, exactly one
+# operator account, created by the first harness connect). Behind a trusted
+# identity edge, switch to header mode:
+#   DRUKS_AUTH_MODE=header
+#   DRUKS_AUTH_HEADER=X-ExeDev-Email
+# Or, for an edge that signs its assertions (Teleport, Cloudflare Access),
+# jwt mode verifies the token in DRUKS_AUTH_HEADER against the edge's JWKS:
+#   DRUKS_AUTH_MODE=jwt
+#   DRUKS_AUTH_HEADER=Teleport-Jwt-Assertion
+#   DRUKS_AUTH_JWKS_URL=https://teleport.example.com/.well-known/jwks.json
+#   DRUKS_AUTH_JWT_ISSUER=https://teleport.example.com
+#   DRUKS_AUTH_JWT_AUDIENCE=https://druks.example.com
+#   DRUKS_AUTH_JWT_IDENTITY_CLAIM=email
+DRUKS_AUTH_MODE=none
+
 # Webhook authentication and public ingress. DRUKS_WEBHOOK_HOST is used only
 # for the doctor's external probe; leave it empty when no public ingress exists.
 DRUKS_WEBHOOK_SECRET=change-me

backend/druks/accounts/constants.py

@@ -5,9 +5,6 @@
 # start; no provider email ever collides with it.
 SYSTEM_ACCOUNT_ID = "system"
 
-# A signed-in session in Redis: token -> account_id.
-SESSION_PREFIX = "druks:session:"
-
 # A personal access token serializes as druks_pat_<prefix>_<secret>. The prefix
 # is the row's lookup key and the only part that may appear in errors, logs, or
 # lists; the secret exists only inside the copy-once plaintext.

backend/druks/accounts/context.py

@@ -0,0 +1,5 @@
+from contextvars import ContextVar
+
+# The request's authenticated account, stamped by the auth gate — Workflow.start
+# reads it so a browser-origin run attributes itself without route ceremony.
+current_account_id: ContextVar[str | None] = ContextVar("current_account_id", default=None)

backend/druks/accounts/dependencies.py

@@ -1,55 +1,159 @@
+from collections.abc import AsyncIterator
+
 from fastapi import HTTPException, Request
+from fastapi.concurrency import run_in_threadpool
 
-from druks.accounts import sessions
-from druks.accounts.exceptions import InvalidPatError
+from druks.accounts.context import current_account_id
+from druks.accounts.exceptions import (
+    AuthConfigurationError,
+    InvalidAssertionError,
+    InvalidPatError,
+)
+from druks.accounts.jwt import verify_assertion
 from druks.accounts.models import Account, PersonalAccessToken
 
 _BEARER_CHALLENGE = 'Bearer realm="druks"'
 
 
-async def resolve_session_account(request: Request) -> Account | None:
-    token = request.cookies.get(sessions.SESSION_COOKIE)
-    if not token:
-        return
-    account_id = await sessions.resolve_session(token)
-    if not account_id:
-        return
-    return Account.get(account_id)
+def resolve_pat_account(header: str) -> Account:
+    """Parse an Authorization header value as exactly ``Bearer <PAT>``.
+    Present means it must authenticate — a malformed or dead credential is a
+    401, never a fall-through to edge identity."""
+    scheme, _, credential = header.partition(" ")
+    if scheme == "Bearer" and credential and " " not in credential:
+        try:
+            pat = PersonalAccessToken.authenticate(credential)
+        except InvalidPatError as error:
+            raise HTTPException(
+                status_code=401,
+                detail=str(error),
+                headers={"WWW-Authenticate": f'{_BEARER_CHALLENGE}, error="invalid_token"'},
+            ) from error
+        return pat.account
+    raise HTTPException(
+        status_code=401,
+        detail="Authorization must be exactly: Bearer <token>.",
+        headers={"WWW-Authenticate": _BEARER_CHALLENGE},
+    )
 
 
-async def current_session_account(request: Request) -> Account:
-    """The signed-in session's account, else 401 — the only door for PAT
-    management, so a token can never manage tokens."""
-    account = await resolve_session_account(request)
-    if account:
-        sessions.current_account_id.set(account.id)
-        return account
-    raise HTTPException(status_code=401, detail="Sign in to use this API.")
+def resolve_none_operator() -> Account | None:
+    """``none`` mode's operator: the only non-system account. None while zero
+    exist (setup state); more than one is configuration drift — refuse loudly
+    rather than guess. The startup validator runs this same check."""
+    operators = Account.list_non_system()
+    if len(operators) > 1:
+        raise AuthConfigurationError(
+            f"auth mode 'none' expects exactly one operator account, found "
+            f"{len(operators)} — remove the extras or switch to header mode"
+        )
+    return operators[0] if operators else None
 
 
-async def current_account(request: Request) -> Account:
-    """The calling account: the Bearer personal access token when an
-    Authorization header is present — never falling back to the cookie —
-    else the signed-in session."""
-    header = request.headers.get("Authorization")
-    # Present-but-empty is still present: it must be challenged, never slide
-    # to the cookie.
-    if header is not None:
-        scheme, _, credential = header.partition(" ")
-        if scheme == "Bearer" and credential and " " not in credential:
-            try:
-                pat = PersonalAccessToken.authenticate(credential)
-            except InvalidPatError as error:
-                raise HTTPException(
-                    status_code=401,
-                    detail=str(error),
-                    headers={"WWW-Authenticate": f'{_BEARER_CHALLENGE}, error="invalid_token"'},
-                ) from error
-            sessions.current_account_id.set(pat.account_id)
-            return pat.account
+async def _resolve_operator(request: Request) -> Account | None:
+    """Edge/none operator identity for this request; None only in none/zero
+    (the setup state). Both edge modes are open enrollment on exactly one
+    nonblank asserted value, trimmed of outer whitespace (the edge gates who
+    reaches druks at all; CITEXT handles case): header mode maps the value
+    itself, jwt mode maps the verified identity claim of the asserted token.
+    ``none`` mode ignores any asserted header — there is no edge to trust."""
+    settings = request.app.state.settings
+    if settings.auth_mode == "none":
+        return resolve_none_operator()
+    values = request.headers.getlist(settings.auth_header)
+    if len(values) != 1:
+        raise HTTPException(
+            status_code=401,
+            detail=f"The edge must assert exactly one {settings.auth_header} identity.",
+        )
+    asserted = values[0].strip()
+    if not asserted:
         raise HTTPException(
             status_code=401,
-            detail="Authorization must be exactly: Bearer <token>.",
-            headers={"WWW-Authenticate": _BEARER_CHALLENGE},
+            detail=f"The edge asserted a blank {settings.auth_header} identity.",
         )
-    return await current_session_account(request)
+    if settings.auth_mode == "header":
+        return Account.get_or_create(asserted)
+    try:
+        # Verification can fetch JWKS behind a sync client — run it off the
+        # loop so a slow or dead key endpoint never stalls unrelated requests.
+        email = await run_in_threadpool(verify_assertion, asserted, settings)
+    except InvalidAssertionError as error:
+        raise HTTPException(status_code=401, detail=str(error)) from error
+    return Account.get_or_create(email)
+
+
+def _require_no_bearer(request: Request) -> None:
+    if request.headers.get("Authorization") is not None:
+        raise HTTPException(
+            status_code=401,
+            detail="This API accepts your edge or local operator identity only, "
+            "never a bearer token.",
+        )
+
+
+def _setup_required() -> HTTPException:
+    return HTTPException(
+        status_code=409,
+        detail="No operator account exists yet — connect a harness to finish setup.",
+    )
+
+
+async def current_account(request: Request) -> AsyncIterator[Account]:
+    """The calling account: the Bearer personal access token when an
+    Authorization header is present — present-but-empty is still present and
+    must be challenged — else the edge/none operator identity."""
+    header = request.headers.get("Authorization")
+    if header is not None:
+        account = resolve_pat_account(header)
+    else:
+        account = await _resolve_operator(request)
+        if not account:
+            raise _setup_required()
+    token = current_account_id.set(account.id)
+    try:
+        yield account
+    finally:
+        # The actor must not leak into whatever runs on this task next.
+        current_account_id.reset(token)
+
+
+async def current_operator_account(request: Request) -> AsyncIterator[Account]:
+    """The edge/none operator, and only that — the door for PAT management, so
+    a token can never manage tokens."""
+    _require_no_bearer(request)
+    account = await _resolve_operator(request)
+    if not account:
+        raise _setup_required()
+    token = current_account_id.set(account.id)
+    try:
+        yield account
+    finally:
+        current_account_id.reset(token)
+
+
+async def current_operator_or_setup(request: Request) -> AsyncIterator[Account | None]:
+    """The edge/none operator for the harness connection flow, where none/zero
+    is a legal state: the first completed connection creates the operator."""
+    _require_no_bearer(request)
+    account = await _resolve_operator(request)
+    token = current_account_id.set(account.id if account else None)
+    try:
+        yield account
+    finally:
+        current_account_id.reset(token)
+
+
+async def current_account_or_setup(request: Request) -> AsyncIterator[Account | None]:
+    """PAT-first identity that reads none/zero as the setup state instead of
+    refusing — only ``/api/auth/me``, which must answer during onboarding."""
+    header = request.headers.get("Authorization")
+    if header is not None:
+        account = resolve_pat_account(header)
+    else:
+        account = await _resolve_operator(request)
+    token = current_account_id.set(account.id if account else None)
+    try:
+        yield account
+    finally:
+        current_account_id.reset(token)

backend/druks/accounts/exceptions.py

@@ -1,3 +1,15 @@
 class InvalidPatError(Exception):
     """A presented bearer credential that resolves to no live personal access
     token — unknown, mismatched, revoked, or expired."""
+
+
+class AuthConfigurationError(Exception):
+    """The configured auth mode cannot resolve a single operator identity —
+    e.g. ``none`` mode with more than one non-system account. Refuses the
+    request (and startup) instead of guessing which account is the operator."""
+
+
+class InvalidAssertionError(Exception):
+    """An edge-minted JWT assertion that fails verification — bad signature,
+    wrong issuer or audience, expired, unknown signing key, or a missing
+    identity claim. The raw token never appears in the message."""

backend/druks/accounts/jwt.py

@@ -0,0 +1,94 @@
+import time
+from functools import lru_cache
+
+import jwt
+
+from druks.accounts.exceptions import InvalidAssertionError
+from druks.settings import Settings
+
+# The edge signs with an asymmetric key it publishes over JWKS; RS256 is the
+# pinned profile — an untrusted token never chooses its own algorithm.
+_ALGORITHMS = ["RS256"]
+_JWKS_TIMEOUT_SECONDS = 5
+# One trust window for both directions: the cached key set expires after it,
+# and an unknown kid can force at most one early refetch inside it — forged
+# kids must not turn key fetching into a request-rate amplifier.
+_JWKS_LIFESPAN_SECONDS = 300
+
+
+@lru_cache(maxsize=1)
+def _jwks_state(url: str) -> dict:
+    # Set-level caching only: a per-key cache would trust a rotated-away or
+    # compromised key until process restart.
+    client = jwt.PyJWKClient(
+        url, cache_keys=False, lifespan=_JWKS_LIFESPAN_SECONDS, timeout=_JWKS_TIMEOUT_SECONDS
+    )
+    # -inf: the first unknown kid is always eligible to force a refetch — a
+    # zero stamp would suppress it while the monotonic clock is still young.
+    return {"client": client, "forced_refresh_at": float("-inf")}
+
+
+def _usable_for_signing(key) -> bool:
+    # Both JWK restriction channels must permit verification — ``use`` (sig or
+    # unrestricted) and ``key_ops`` (verify or unrestricted) — so a
+    # mixed-purpose JWKS can never let an encryption key mint identities.
+    # key_ops has no public accessor on PyJWK.
+    if key.public_key_use not in ("sig", None):
+        return False
+    operations = key._jwk_data.get("key_ops")
+    # Absent means unrestricted; a present list — even an empty one — permits
+    # exactly what it names, and anything malformed permits nothing.
+    if operations is None:
+        return True
+    return isinstance(operations, list) and "verify" in operations
+
+
+def _find_key(jwk_set, kid: str):
+    return next(
+        (key for key in jwk_set.keys if key.key_id == kid and _usable_for_signing(key)), None
+    )
+
+
+def _signing_key(url: str, kid: str):
+    state = _jwks_state(url)
+    key = _find_key(state["client"].get_jwk_set(), kid)
+    if key:
+        return key.key
+    if time.monotonic() - state["forced_refresh_at"] >= _JWKS_LIFESPAN_SECONDS:
+        state["forced_refresh_at"] = time.monotonic()
+        key = _find_key(state["client"].get_jwk_set(refresh=True), kid)
+        if key:
+            return key.key
+    raise InvalidAssertionError("Assertion rejected: unknown signing key.")
+
+
+def verify_assertion(token: str, settings: Settings) -> str:
+    """The verified identity claim of an edge-minted assertion, else
+    InvalidAssertionError. Fails closed on every fetch, signature, claim, or
+    shape problem; only the failure class name reaches the message."""
+    try:
+        kid = jwt.get_unverified_header(token).get("kid")
+        if not kid:
+            raise InvalidAssertionError("Assertion rejected: no signing key named.")
+        claims = jwt.decode(
+            token,
+            _signing_key(settings.auth_jwks_url, kid),
+            algorithms=_ALGORITHMS,
+            issuer=settings.auth_jwt_issuer,
+            audience=settings.auth_jwt_audience,
+            options={"require": ["exp", "iss", "aud"]},
+        )
+    except InvalidAssertionError:
+        raise
+    except Exception as error:
+        # The whole path fails closed — a broken JWKS endpoint or a transport
+        # surprise is a rejected assertion, never a 500.
+        raise InvalidAssertionError(f"Assertion rejected: {type(error).__name__}.") from error
+    claim = claims.get(settings.auth_jwt_identity_claim)
+    # Claims are external data: the identity must be a nonblank string, not
+    # whatever shape the token happened to carry.
+    if isinstance(claim, str) and claim.strip():
+        return claim.strip()
+    raise InvalidAssertionError(
+        f"Assertion carries no usable {settings.auth_jwt_identity_claim} claim."
+    )

backend/druks/accounts/models.py

@@ -5,7 +5,7 @@
 from datetime import datetime
 
 from sqlalchemy import ForeignKey, Index, LargeBinary, String, select
-from sqlalchemy.dialects.postgresql import CITEXT
+from sqlalchemy.dialects.postgresql import CITEXT, insert
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from druks.accounts.constants import (
@@ -16,6 +16,7 @@
     PAT_PREFIX_LENGTH,
     PAT_SECRET_BYTES,
     PAT_TOKEN_TAG,
+    SYSTEM_ACCOUNT_ID,
 )
 from druks.accounts.exceptions import InvalidPatError
 from druks.core.models import Uuid7Pk
@@ -46,14 +47,24 @@ def get_for_username(cls, username: str) -> "Account | None":
 
     @classmethod
     def get_or_create(cls, username: str) -> "Account":
+        """Concurrency-safe lookup-or-create: racing requests both INSERT with
+        ON CONFLICT DO NOTHING, then converge on the one row through the
+        canonical CITEXT lookup."""
         account = cls.get_for_username(username)
         if account:
             return account
-        account = cls(username=username)
         session = db_session()
-        session.add(account)
-        session.flush()
-        return account
+        session.execute(
+            insert(cls)
+            .values(username=username)
+            .on_conflict_do_nothing(index_elements=["username"])
+        )
+        return session.scalars(select(cls).where(cls.username == username)).one()
+
+    @classmethod
+    def list_non_system(cls) -> list["Account"]:
+        stmt = select(cls).where(cls.username != SYSTEM_ACCOUNT_ID).order_by(cls.created_at)
+        return list(db_session().scalars(stmt))
 
 
 def _hash_token(token: str) -> bytes: