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:
@@ -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