Files
docqube_backend/app/middleware/tenant_context_middleware.py
T
2026-09-08 11:00:05 +05:30

134 lines
4.9 KiB
Python

"""
B1.4 — establish the tenant context for every request.
This is the wiring the rest of B1 was waiting on, and its absence was the single
thing standing between "isolation is built" and "isolation works". The ORM
listener (`app/core/tenant_filter.py`) and the RLS session variables
(`app/core/rls.py`) both read `current_tenant_id()`. Both were built, tested and
correct. Nothing in the application ever set it — the only callers of
`scoped_to()` in the repository were the tests, which supply their own context
and therefore proved the machinery while never exercising the wiring.
The consequence was quiet, because the flag defaults off: every request ran with
no context, both layers stayed inert, and the application behaved exactly as it
always had. Turning `TENANT_FILTER_ENABLED` on would have failed closed on every
query in the system.
**Why middleware and not a dependency.** `get_current_user` is a `def`, so
FastAPI runs it in a threadpool, and a `ContextVar` set inside a worker thread
does not propagate back to the request's task — the write is simply lost at the
thread boundary. Middleware runs in the request's own task, so what it sets is
visible to every dependency and handler that follows, including the sync ones
(the context is *copied into* the threadpool, which is all a reader needs).
**Why raw ASGI and not `BaseHTTPMiddleware`.** `BaseHTTPMiddleware` runs the
downstream app in a separate task and pumps the response through a queue, which
breaks streaming: this application serves file downloads through
`StreamingResponse` and has a Server-Sent Events endpoint on
`text/event-stream`, and SSE under that middleware stops arriving incrementally.
A plain ASGI callable adds no task, no queue and no buffering — it sets a
context variable and gets out of the way.
**Why the token and not the database.** Reading the user here would mean a query
per request before any route is chosen, and that query would itself need the
context it is trying to establish. The token is signature-verified before its
claims are read, and `get_current_user` independently checks the tenant claim
against the loaded row, rejecting any mismatch — so a forged claim fails there
rather than being trusted here.
"""
import logging
import uuid
from jose import JWTError, jwt
from starlette.datastructures import Headers
from starlette.types import ASGIApp, Message, Receive, Scope, Send
from app.core.request_cache import reset_request_cache
from app.core.scope import reset_scope_cache
from app.core.settings import settings
from app.core.tenant_context import reset_context, set_context
logger = logging.getLogger(__name__)
_COOKIE_NAMES = ("docqube_access_token", "access_token")
def _token_from(headers: Headers) -> str | None:
authorization = headers.get("authorization") or ""
if authorization.lower().startswith("bearer "):
candidate = authorization[7:].strip()
if candidate:
return candidate
cookie_header = headers.get("cookie")
if not cookie_header:
return None
for part in cookie_header.split(";"):
name, _, value = part.strip().partition("=")
if name in _COOKIE_NAMES and value:
return value
return None
def _context_from_token(token: str) -> tuple[uuid.UUID | None, bool] | None:
"""The tenant and privilege a verified token asserts, or None if unusable."""
try:
claims = jwt.decode(token, settings.APP_SECRET, algorithms=[settings.ALGORITHM])
except JWTError:
return None
if claims.get("type") in ("refresh", "device_management"):
return None
tenant_id = None
raw = claims.get("tenant_id")
if raw:
try:
tenant_id = uuid.UUID(str(raw))
except (ValueError, AttributeError, TypeError):
logger.warning("Token carried an unparseable tenant_id")
is_super = bool(claims.get("is_superadmin", False))
if tenant_id is None and not is_super:
return None
return tenant_id, is_super
class TenantContextMiddleware:
def __init__(self, app: ASGIApp) -> None:
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] not in ("http", "websocket"):
await self.app(scope, receive, send)
return
reset_scope_cache()
reset_request_cache()
tokens = None
raw_token = _token_from(Headers(scope=scope))
if raw_token:
resolved = _context_from_token(raw_token)
if resolved is not None:
tenant_id, is_super = resolved
tokens = set_context(tenant_id, is_super=is_super)
try:
await self.app(scope, receive, send)
finally:
if tokens is not None:
try:
reset_context(tokens)
except ValueError:
pass
def install(app: ASGIApp) -> None:
app.add_middleware(TenantContextMiddleware)
__all__ = ["TenantContextMiddleware", "install"]