150 lines
5.5 KiB
Python
150 lines
5.5 KiB
Python
"""Pushing the tenant context into PostgreSQL, so the database enforces isolation.
|
|
|
|
Until now isolation has been a convention: every query had to remember
|
|
`.filter(tenant_id == ...)`. Finding S-5 was one that forgot, and the reason it
|
|
was possible to forget is that nothing outside the developer's memory was
|
|
checking.
|
|
|
|
This makes it a property instead. Before any query runs, the session's tenant is
|
|
written into two PostgreSQL settings that the row-level security policies read:
|
|
|
|
app.tenant_id the workspace, or '' when none is set
|
|
app.bypass_rls 'on' for deliberate cross-workspace work
|
|
|
|
A policy that sees neither returns no rows. So a code path that forgets to
|
|
establish context fails loudly and safely, rather than quietly reading everyone's
|
|
data.
|
|
|
|
**This only bites if the application connects as a non-owner role.** A superuser
|
|
bypasses row-level security unconditionally, and the table owner does too unless
|
|
the table is FORCEd — which migration `c3d5e7f9a801` does. `scripts/create_app_role.py`
|
|
creates the role to connect as; `check_rls_enforced()` reports whether the
|
|
connection in use is actually subject to the policies, because "policies exist"
|
|
and "policies apply" look identical in the schema.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
from sqlalchemy import event, text
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.tenant_context import current_tenant_id, is_bypassed
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
TENANT_SETTING = "app.tenant_id"
|
|
BYPASS_SETTING = "app.bypass_rls"
|
|
|
|
_STATE_KEY = "_rls_state"
|
|
_APPLYING_KEY = "_rls_applying"
|
|
|
|
|
|
def _desired_state() -> tuple[str, str]:
|
|
if is_bypassed():
|
|
return "", "on"
|
|
tenant_id = current_tenant_id()
|
|
return (str(tenant_id) if tenant_id else ""), "off"
|
|
|
|
|
|
def apply_rls_context(session: Session) -> None:
|
|
"""Write the current context into the session, if it has changed.
|
|
|
|
Guarded against re-entry because setting the context is itself a query, and
|
|
the listener that calls this fires on every query.
|
|
|
|
`set_config(..., true)` makes the setting local to the transaction, so it is
|
|
discarded on commit or rollback and cannot leak into whatever uses this
|
|
connection next out of the pool.
|
|
"""
|
|
if session.info.get(_APPLYING_KEY):
|
|
return
|
|
|
|
desired = _desired_state()
|
|
if session.info.get(_STATE_KEY) == desired:
|
|
return
|
|
|
|
tenant_id, bypass = desired
|
|
session.info[_APPLYING_KEY] = True
|
|
try:
|
|
session.execute(
|
|
text("SELECT set_config(:k1, :v1, true), set_config(:k2, :v2, true)"),
|
|
{"k1": TENANT_SETTING, "v1": tenant_id, "k2": BYPASS_SETTING, "v2": bypass},
|
|
)
|
|
session.info[_STATE_KEY] = desired
|
|
finally:
|
|
session.info[_APPLYING_KEY] = False
|
|
|
|
|
|
def register_rls_listener() -> None:
|
|
@event.listens_for(Session, "do_orm_execute")
|
|
def _sync_on_query(execute_state: Any) -> None:
|
|
apply_rls_context(execute_state.session)
|
|
|
|
@event.listens_for(Session, "before_flush")
|
|
def _sync_on_flush(session: Session, flush_context: Any, instances: Any) -> None:
|
|
"""Writes need the context too, and do not come through `do_orm_execute`.
|
|
|
|
That event fires for ORM *queries*. An INSERT produced by `session.add()`
|
|
plus a flush does not pass through it, so without this hook every write
|
|
ran with no `app.tenant_id` set — and the policy's WITH CHECK refused it.
|
|
The symptom is a permission error on a perfectly ordinary insert, which
|
|
is a confusing way to discover a missing listener.
|
|
"""
|
|
apply_rls_context(session)
|
|
|
|
@event.listens_for(Session, "after_transaction_end")
|
|
def _forget(session: Session, transaction: Any) -> None:
|
|
session.info.pop(_STATE_KEY, None)
|
|
|
|
logger.debug("row-level security context listener registered")
|
|
|
|
|
|
def check_rls_enforced(session: Session) -> list[str]:
|
|
"""Whether isolation is actually in force on *this* connection.
|
|
|
|
Returns the problems found, empty when all is well. Worth calling from a
|
|
health check: a schema full of policies and a superuser connection look
|
|
exactly like a schema full of policies and a correct connection.
|
|
"""
|
|
problems: list[str] = []
|
|
|
|
rows = session.execute(
|
|
text(
|
|
"""
|
|
SELECT c.relname, c.relrowsecurity, c.relforcerowsecurity
|
|
FROM pg_class c
|
|
JOIN pg_namespace n ON n.oid = c.relnamespace
|
|
WHERE n.nspname = current_schema()
|
|
AND c.relkind = 'r'
|
|
AND EXISTS (SELECT 1 FROM pg_policy p WHERE p.polrelid = c.oid)
|
|
"""
|
|
)
|
|
).fetchall()
|
|
|
|
if not rows:
|
|
problems.append("no table carries a row-level security policy")
|
|
|
|
for name, enabled, forced in rows:
|
|
if not enabled:
|
|
problems.append(f"{name}: has policies but row security is disabled")
|
|
elif not forced:
|
|
problems.append(
|
|
f"{name}: row security is enabled but not FORCEd, so the owning role "
|
|
"bypasses it — and the owning role is what this application connects as"
|
|
)
|
|
|
|
is_super = session.execute(
|
|
text("SELECT usesuper FROM pg_user WHERE usename = current_user")
|
|
).scalar()
|
|
if is_super:
|
|
problems.append(
|
|
"connected as a PostgreSQL superuser, which bypasses row-level security "
|
|
"unconditionally — no policy applies to this connection. Point DATABASE_URL "
|
|
"at the role created by scripts/create_app_role.py"
|
|
)
|
|
|
|
return problems
|