feat(auth): AuthMiddleware injects workspace ContextVar from JWT wid/role
deps.get_current_user_from_request now stashes the decoded payload on request.state.auth_payload so AuthMiddleware can populate the workspace_context ContextVar without a second decode. Reset is paired in the same try/finally as user_context to keep teardown atomic. Also adds: - auth.models.ActiveWorkspace — minimal proxy that satisfies the CurrentWorkspace protocol (id + role only). - auth.models.User.default_workspace_id — surfaces the DB column added in T4.4 so the eventual /auth/me payload can reference it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -31,6 +31,12 @@ class User(BaseModel):
|
|||||||
needs_setup: bool = Field(default=False, description="True for auto-created admin until setup completes")
|
needs_setup: bool = Field(default=False, description="True for auto-created admin until setup completes")
|
||||||
token_version: int = Field(default=0, description="Incremented on password change to invalidate old JWTs")
|
token_version: int = Field(default=0, description="Incremented on password change to invalidate old JWTs")
|
||||||
|
|
||||||
|
# Workspace linkage (Stage 0 PR4)
|
||||||
|
default_workspace_id: str | None = Field(
|
||||||
|
default=None,
|
||||||
|
description="The workspace the user lands in by default after login. NULL → /select-workspace.",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class UserResponse(BaseModel):
|
class UserResponse(BaseModel):
|
||||||
"""Response model for user info endpoint."""
|
"""Response model for user info endpoint."""
|
||||||
@@ -39,3 +45,18 @@ class UserResponse(BaseModel):
|
|||||||
email: str
|
email: str
|
||||||
system_role: Literal["admin", "user"]
|
system_role: Literal["admin", "user"]
|
||||||
needs_setup: bool = False
|
needs_setup: bool = False
|
||||||
|
|
||||||
|
|
||||||
|
class ActiveWorkspace(BaseModel):
|
||||||
|
"""Lightweight workspace proxy injected into the request-scoped contextvar.
|
||||||
|
|
||||||
|
Implements the structural ``CurrentWorkspace`` protocol expected by
|
||||||
|
``deerflow.runtime.workspace_context``: only ``.id`` (str) and
|
||||||
|
``.role`` (str) are required. We intentionally do *not* embed the
|
||||||
|
full ``WorkspaceRow`` here — the middleware needs to set the
|
||||||
|
contextvar on every request and an extra DB lookup just to populate
|
||||||
|
a name/slug we don't use yet would be wasted work.
|
||||||
|
"""
|
||||||
|
|
||||||
|
id: str
|
||||||
|
role: str
|
||||||
|
|||||||
@@ -17,9 +17,11 @@ from starlette.responses import JSONResponse
|
|||||||
from starlette.types import ASGIApp
|
from starlette.types import ASGIApp
|
||||||
|
|
||||||
from app.gateway.auth.errors import AuthErrorCode, AuthErrorResponse
|
from app.gateway.auth.errors import AuthErrorCode, AuthErrorResponse
|
||||||
|
from app.gateway.auth.models import ActiveWorkspace
|
||||||
from app.gateway.authz import _ALL_PERMISSIONS, AuthContext
|
from app.gateway.authz import _ALL_PERMISSIONS, AuthContext
|
||||||
from app.gateway.internal_auth import INTERNAL_AUTH_HEADER_NAME, get_internal_user, is_valid_internal_auth_token
|
from app.gateway.internal_auth import INTERNAL_AUTH_HEADER_NAME, get_internal_user, is_valid_internal_auth_token
|
||||||
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
||||||
|
from deerflow.runtime.workspace_context import reset_current_workspace, set_current_workspace
|
||||||
|
|
||||||
# Paths that never require authentication.
|
# Paths that never require authentication.
|
||||||
_PUBLIC_PATH_PREFIXES: tuple[str, ...] = (
|
_PUBLIC_PATH_PREFIXES: tuple[str, ...] = (
|
||||||
@@ -119,8 +121,22 @@ class AuthMiddleware(BaseHTTPMiddleware):
|
|||||||
# JWT-decode + DB-lookup pipeline a second time per request).
|
# JWT-decode + DB-lookup pipeline a second time per request).
|
||||||
request.state.user = user
|
request.state.user = user
|
||||||
request.state.auth = AuthContext(user=user, permissions=_ALL_PERMISSIONS)
|
request.state.auth = AuthContext(user=user, permissions=_ALL_PERMISSIONS)
|
||||||
token = set_current_user(user)
|
user_token = set_current_user(user)
|
||||||
|
|
||||||
|
# Inject workspace contextvar from the JWT's wid/role claims.
|
||||||
|
# decode_token has already rejected legacy no-wid tokens upstream,
|
||||||
|
# so by the time we get here payload.wid is guaranteed non-None
|
||||||
|
# for cookie-authenticated requests. Internal-auth requests skip
|
||||||
|
# the workspace contextvar (they don't have a workspace scope —
|
||||||
|
# the internal user is a system actor).
|
||||||
|
ws_token = None
|
||||||
|
payload = getattr(request.state, "auth_payload", None)
|
||||||
|
if payload is not None and payload.wid is not None:
|
||||||
|
ws_token = set_current_workspace(ActiveWorkspace(id=payload.wid, role=payload.role or "owner"))
|
||||||
|
|
||||||
try:
|
try:
|
||||||
return await call_next(request)
|
return await call_next(request)
|
||||||
finally:
|
finally:
|
||||||
reset_current_user(token)
|
if ws_token is not None:
|
||||||
|
reset_current_workspace(ws_token)
|
||||||
|
reset_current_user(user_token)
|
||||||
|
|||||||
@@ -220,6 +220,10 @@ async def get_current_user_from_request(request: Request):
|
|||||||
detail=AuthErrorResponse(code=AuthErrorCode.TOKEN_INVALID, message="Token revoked (password changed)").model_dump(),
|
detail=AuthErrorResponse(code=AuthErrorCode.TOKEN_INVALID, message="Token revoked (password changed)").model_dump(),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Stash decoded payload on request.state so AuthMiddleware can read
|
||||||
|
# wid/role for the workspace contextvar without a second decode.
|
||||||
|
request.state.auth_payload = payload
|
||||||
|
|
||||||
return user
|
return user
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,130 @@
|
|||||||
|
"""AuthMiddleware injects the workspace ContextVar (Stage 0 PR4 T4.7).
|
||||||
|
|
||||||
|
After PR4 every authenticated request has a workspace bound on
|
||||||
|
``deerflow.runtime.workspace_context._current_workspace``. The middleware
|
||||||
|
populates it from the JWT's ``wid`` / ``role`` claims, mirrors what it
|
||||||
|
already does for ``user_context``, and tears both down in a single
|
||||||
|
``try/finally`` so leaks don't cross requests.
|
||||||
|
|
||||||
|
Legacy 4-field tokens (no ``wid``) are rejected upstream by
|
||||||
|
``decode_token`` (T4.6) — those should never reach the workspace
|
||||||
|
injection branch; this file pins that the 401 they trigger carries
|
||||||
|
``AuthErrorCode.WORKSPACE_REQUIRED``.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from datetime import UTC, datetime, timedelta
|
||||||
|
from unittest.mock import AsyncMock, patch
|
||||||
|
from uuid import uuid4
|
||||||
|
|
||||||
|
import jwt
|
||||||
|
import pytest
|
||||||
|
from fastapi import FastAPI
|
||||||
|
from starlette.testclient import TestClient
|
||||||
|
|
||||||
|
from app.gateway.auth import create_access_token
|
||||||
|
from app.gateway.auth.config import get_auth_config
|
||||||
|
from app.gateway.auth.models import User
|
||||||
|
from app.gateway.auth_middleware import AuthMiddleware
|
||||||
|
from deerflow.runtime.workspace_context import get_current_workspace
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture(autouse=True)
|
||||||
|
def _stable_jwt_secret(monkeypatch):
|
||||||
|
monkeypatch.setenv("AUTH_JWT_SECRET", "test-secret-key-for-jwt-testing-minimum-32-chars")
|
||||||
|
yield
|
||||||
|
|
||||||
|
|
||||||
|
def _make_app() -> FastAPI:
|
||||||
|
"""App with AuthMiddleware + an inspect route that surfaces the contextvar."""
|
||||||
|
app = FastAPI()
|
||||||
|
app.add_middleware(AuthMiddleware)
|
||||||
|
|
||||||
|
@app.get("/api/v1/auth/setup-status") # public — never gates on wid
|
||||||
|
async def setup_status():
|
||||||
|
return {"needs_setup": False}
|
||||||
|
|
||||||
|
@app.get("/api/models") # protected — exercises wid injection
|
||||||
|
async def inspect_workspace():
|
||||||
|
ws = get_current_workspace()
|
||||||
|
if ws is None:
|
||||||
|
return {"workspace": None}
|
||||||
|
return {"workspace": {"id": ws.id, "role": ws.role}}
|
||||||
|
|
||||||
|
return app
|
||||||
|
|
||||||
|
|
||||||
|
def _make_user(uid: str) -> User:
|
||||||
|
return User(id=uid, email="t@example.com", password_hash="hash", token_version=0)
|
||||||
|
|
||||||
|
|
||||||
|
def _make_legacy_token() -> str:
|
||||||
|
"""Encode a pre-PR4 JWT (no wid/role) directly."""
|
||||||
|
now = datetime.now(UTC)
|
||||||
|
payload = {
|
||||||
|
"sub": str(uuid4()),
|
||||||
|
"exp": now + timedelta(hours=1),
|
||||||
|
"iat": now,
|
||||||
|
"ver": 0,
|
||||||
|
}
|
||||||
|
return jwt.encode(payload, get_auth_config().jwt_secret, algorithm="HS256")
|
||||||
|
|
||||||
|
|
||||||
|
def test_new_jwt_injects_workspace_into_contextvar() -> None:
|
||||||
|
"""Cookie with wid+role → route observes the workspace via the contextvar."""
|
||||||
|
uid = str(uuid4())
|
||||||
|
token = create_access_token(uid, workspace_id="ws-abc", role="owner")
|
||||||
|
|
||||||
|
with patch("app.gateway.deps.get_local_provider") as fn:
|
||||||
|
fn.return_value.get_user = AsyncMock(return_value=_make_user(uid))
|
||||||
|
client = TestClient(_make_app())
|
||||||
|
res = client.get("/api/models", cookies={"access_token": token})
|
||||||
|
|
||||||
|
assert res.status_code == 200, res.text
|
||||||
|
assert res.json() == {"workspace": {"id": "ws-abc", "role": "owner"}}
|
||||||
|
|
||||||
|
|
||||||
|
def test_legacy_jwt_rejected_with_workspace_required() -> None:
|
||||||
|
"""No-wid tokens get 401 with AuthErrorCode.WORKSPACE_REQUIRED, not generic token_invalid."""
|
||||||
|
client = TestClient(_make_app())
|
||||||
|
res = client.get("/api/models", cookies={"access_token": _make_legacy_token()})
|
||||||
|
|
||||||
|
assert res.status_code == 401
|
||||||
|
assert res.json()["detail"]["code"] == "workspace_required"
|
||||||
|
|
||||||
|
|
||||||
|
def test_public_path_skips_workspace_check() -> None:
|
||||||
|
"""Public whitelist (e.g. /api/v1/auth/setup-status) does not require wid."""
|
||||||
|
client = TestClient(_make_app())
|
||||||
|
res = client.get("/api/v1/auth/setup-status") # no cookie at all
|
||||||
|
assert res.status_code == 200
|
||||||
|
|
||||||
|
|
||||||
|
def test_workspace_contextvar_resets_between_requests() -> None:
|
||||||
|
"""After dispatch returns the contextvar must be clear (no leak across requests).
|
||||||
|
|
||||||
|
Why we test this: if the try/finally is wired only for user_context but
|
||||||
|
not workspace_context, two back-to-back requests can see each other's
|
||||||
|
workspace under asyncio task switching.
|
||||||
|
"""
|
||||||
|
uid = str(uuid4())
|
||||||
|
token = create_access_token(uid, workspace_id="ws-first", role="owner")
|
||||||
|
|
||||||
|
# First request resolves to ws-first
|
||||||
|
with patch("app.gateway.deps.get_local_provider") as fn:
|
||||||
|
fn.return_value.get_user = AsyncMock(return_value=_make_user(uid))
|
||||||
|
client = TestClient(_make_app())
|
||||||
|
res1 = client.get("/api/models", cookies={"access_token": token})
|
||||||
|
assert res1.json() == {"workspace": {"id": "ws-first", "role": "owner"}}
|
||||||
|
|
||||||
|
# Outside the request scope the contextvar must be empty again.
|
||||||
|
assert get_current_workspace() is None
|
||||||
|
|
||||||
|
# Second request with a different workspace must not see ws-first.
|
||||||
|
token2 = create_access_token(uid, workspace_id="ws-second", role="owner")
|
||||||
|
with patch("app.gateway.deps.get_local_provider") as fn:
|
||||||
|
fn.return_value.get_user = AsyncMock(return_value=_make_user(uid))
|
||||||
|
client = TestClient(_make_app())
|
||||||
|
res2 = client.get("/api/models", cookies={"access_token": token2})
|
||||||
|
assert res2.json() == {"workspace": {"id": "ws-second", "role": "owner"}}
|
||||||
Reference in New Issue
Block a user