134 lines
4.9 KiB
Python
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"]
|