Files
aios/active-runtime/app/aios_v2r1/router.py
T

336 lines
13 KiB
Python

"""Authoritative tenant registry and fail-closed message router."""
from __future__ import annotations
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timedelta
import sqlite3
from pathlib import Path
from typing import Callable, Iterator
from .tenant_context import (
CanonicalMessage,
ConflictError,
ScopeError,
TenantContext,
ValidationError,
isoformat,
parse_time,
require_ident,
utc_now,
)
@dataclass(frozen=True)
class Route:
tenant_id: str
cell_id: str
channel: str
channel_account_id: str
principal_id: str
session_id: str | None
class TenantRouter:
"""SQLite-backed routing registry.
Channel accounts are authoritative tenant ownership boundaries. Principal and
session bindings may narrow a route, but can never move traffic to another
tenant or Cell.
"""
def __init__(
self, database: str | Path, *, clock: Callable[[], datetime] = utc_now
) -> None:
self.database = str(database)
self._clock = clock
self._initialize()
def _connect(self) -> sqlite3.Connection:
connection = sqlite3.connect(self.database, timeout=30, isolation_level=None)
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA foreign_keys=ON")
connection.execute("PRAGMA journal_mode=WAL")
connection.execute("PRAGMA synchronous=FULL")
connection.execute("PRAGMA busy_timeout=30000")
return connection
@contextmanager
def _transaction(self) -> Iterator[sqlite3.Connection]:
connection = self._connect()
try:
connection.execute("BEGIN IMMEDIATE")
yield connection
connection.execute("COMMIT")
except Exception:
connection.execute("ROLLBACK")
raise
finally:
connection.close()
def _initialize(self) -> None:
with self._connect() as db:
db.executescript(
"""
CREATE TABLE IF NOT EXISTS tenants (
tenant_id TEXT PRIMARY KEY,
cell_id TEXT NOT NULL UNIQUE,
status TEXT NOT NULL CHECK(status IN ('active','suspended','retired')),
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS channel_routes (
channel TEXT NOT NULL,
channel_account_id TEXT NOT NULL,
tenant_id TEXT NOT NULL,
cell_id TEXT NOT NULL,
created_at TEXT NOT NULL,
PRIMARY KEY(channel, channel_account_id),
FOREIGN KEY(tenant_id) REFERENCES tenants(tenant_id)
);
CREATE TABLE IF NOT EXISTS principal_routes (
channel TEXT NOT NULL,
channel_account_id TEXT NOT NULL,
principal_id TEXT NOT NULL,
tenant_id TEXT NOT NULL,
cell_id TEXT NOT NULL,
created_at TEXT NOT NULL,
PRIMARY KEY(channel, channel_account_id, principal_id),
FOREIGN KEY(channel, channel_account_id)
REFERENCES channel_routes(channel, channel_account_id)
);
CREATE TABLE IF NOT EXISTS session_bindings (
tenant_id TEXT NOT NULL,
session_id TEXT NOT NULL,
cell_id TEXT NOT NULL,
principal_id TEXT NOT NULL,
expires_at TEXT NOT NULL,
created_at TEXT NOT NULL,
PRIMARY KEY(tenant_id, session_id),
FOREIGN KEY(tenant_id) REFERENCES tenants(tenant_id)
);
CREATE INDEX IF NOT EXISTS idx_session_expiry
ON session_bindings(expires_at);
"""
)
def register_tenant(self, tenant_id: str, cell_id: str) -> None:
tenant_id = require_ident("tenant_id", tenant_id, tenant=True)
cell_id = require_ident("cell_id", cell_id)
now = isoformat(self._clock())
with self._transaction() as db:
row = db.execute(
"SELECT cell_id, status FROM tenants WHERE tenant_id=?", (tenant_id,)
).fetchone()
if row:
if row["cell_id"] != cell_id:
raise ConflictError("tenant already bound to another cell")
db.execute(
"UPDATE tenants SET status='active', updated_at=? WHERE tenant_id=?",
(now, tenant_id),
)
return
conflict = db.execute(
"SELECT tenant_id FROM tenants WHERE cell_id=?", (cell_id,)
).fetchone()
if conflict:
raise ConflictError("cell already belongs to another tenant")
db.execute(
"INSERT INTO tenants VALUES(?,?,?,?,?)",
(tenant_id, cell_id, "active", now, now),
)
def set_tenant_status(self, tenant_id: str, status: str) -> None:
tenant_id = require_ident("tenant_id", tenant_id, tenant=True)
if status not in {"active", "suspended", "retired"}:
raise ValidationError("invalid tenant status")
with self._transaction() as db:
cursor = db.execute(
"UPDATE tenants SET status=?, updated_at=? WHERE tenant_id=?",
(status, isoformat(self._clock()), tenant_id),
)
if cursor.rowcount != 1:
raise ValidationError("unknown tenant")
def bind_channel(
self,
*,
tenant_id: str,
channel: str,
channel_account_id: str,
) -> None:
tenant_id = require_ident("tenant_id", tenant_id, tenant=True)
channel = require_ident("channel", channel)
channel_account_id = require_ident("channel_account_id", channel_account_id)
with self._transaction() as db:
tenant = db.execute(
"SELECT cell_id, status FROM tenants WHERE tenant_id=?", (tenant_id,)
).fetchone()
if not tenant or tenant["status"] != "active":
raise ValidationError("tenant is not active")
existing = db.execute(
"SELECT tenant_id, cell_id FROM channel_routes "
"WHERE channel=? AND channel_account_id=?",
(channel, channel_account_id),
).fetchone()
if existing:
if existing["tenant_id"] != tenant_id or existing["cell_id"] != tenant["cell_id"]:
raise ConflictError("channel account already belongs to another tenant")
return
db.execute(
"INSERT INTO channel_routes VALUES(?,?,?,?,?)",
(channel, channel_account_id, tenant_id, tenant["cell_id"], isoformat(self._clock())),
)
def bind_principal(
self,
*,
tenant_id: str,
channel: str,
channel_account_id: str,
principal_id: str,
) -> None:
principal_id = require_ident("principal_id", principal_id)
route = self.resolve_channel(channel, channel_account_id)
if route.tenant_id != tenant_id:
raise ScopeError("principal cannot cross tenant boundary")
with self._transaction() as db:
existing = db.execute(
"SELECT tenant_id, cell_id FROM principal_routes WHERE "
"channel=? AND channel_account_id=? AND principal_id=?",
(channel, channel_account_id, principal_id),
).fetchone()
if existing:
if existing["tenant_id"] != tenant_id or existing["cell_id"] != route.cell_id:
raise ConflictError("principal route conflict")
return
db.execute(
"INSERT INTO principal_routes VALUES(?,?,?,?,?,?)",
(
channel,
channel_account_id,
principal_id,
tenant_id,
route.cell_id,
isoformat(self._clock()),
),
)
def bind_session(
self,
*,
tenant_id: str,
session_id: str,
cell_id: str,
principal_id: str,
ttl_seconds: int = 86400,
) -> None:
tenant_id = require_ident("tenant_id", tenant_id, tenant=True)
session_id = require_ident("session_id", session_id)
cell_id = require_ident("cell_id", cell_id)
principal_id = require_ident("principal_id", principal_id)
if not 60 <= ttl_seconds <= 31 * 86400:
raise ValidationError("invalid session TTL")
now = self._clock()
expires = isoformat(now + timedelta(seconds=ttl_seconds))
with self._transaction() as db:
tenant = db.execute(
"SELECT cell_id, status FROM tenants WHERE tenant_id=?", (tenant_id,)
).fetchone()
if not tenant or tenant["status"] != "active" or tenant["cell_id"] != cell_id:
raise ScopeError("session cell does not match tenant")
existing = db.execute(
"SELECT cell_id, principal_id FROM session_bindings "
"WHERE tenant_id=? AND session_id=?",
(tenant_id, session_id),
).fetchone()
if existing and (
existing["cell_id"] != cell_id or existing["principal_id"] != principal_id
):
raise ConflictError("session is already bound to another principal or cell")
db.execute(
"INSERT INTO session_bindings VALUES(?,?,?,?,?,?) "
"ON CONFLICT(tenant_id,session_id) DO UPDATE SET expires_at=excluded.expires_at",
(tenant_id, session_id, cell_id, principal_id, expires, isoformat(now)),
)
def resolve_channel(self, channel: str, channel_account_id: str) -> Route:
channel = require_ident("channel", channel)
channel_account_id = require_ident("channel_account_id", channel_account_id)
with self._connect() as db:
row = db.execute(
"SELECT r.tenant_id, r.cell_id, t.status FROM channel_routes r "
"JOIN tenants t ON t.tenant_id=r.tenant_id "
"WHERE r.channel=? AND r.channel_account_id=?",
(channel, channel_account_id),
).fetchone()
if not row:
raise ScopeError("unknown channel account")
if row["status"] != "active":
raise ScopeError("tenant is not active")
return Route(row["tenant_id"], row["cell_id"], channel, channel_account_id, "", None)
def resolve(
self,
*,
channel: str,
channel_account_id: str,
principal_id: str,
session_id: str | None = None,
) -> Route:
principal_id = require_ident("principal_id", principal_id)
base = self.resolve_channel(channel, channel_account_id)
with self._connect() as db:
principal = db.execute(
"SELECT tenant_id, cell_id FROM principal_routes WHERE "
"channel=? AND channel_account_id=? AND principal_id=?",
(channel, channel_account_id, principal_id),
).fetchone()
if principal and (
principal["tenant_id"] != base.tenant_id or principal["cell_id"] != base.cell_id
):
raise ScopeError("principal route violates channel ownership")
if session_id is not None:
session_id = require_ident("session_id", session_id)
session = db.execute(
"SELECT cell_id, principal_id, expires_at FROM session_bindings "
"WHERE tenant_id=? AND session_id=?",
(base.tenant_id, session_id),
).fetchone()
if not session:
raise ScopeError("unknown session")
if parse_time(session["expires_at"]) <= self._clock():
raise ScopeError("session expired")
if session["cell_id"] != base.cell_id or session["principal_id"] != principal_id:
raise ScopeError("session scope mismatch")
return Route(
base.tenant_id,
base.cell_id,
channel,
channel_account_id,
principal_id,
session_id,
)
def assert_message_route(self, message: CanonicalMessage) -> Route:
route = self.resolve(
channel=message.channel,
channel_account_id=message.channel_account_id,
principal_id=message.principal_id,
session_id=message.session_id,
)
if message.tenant_id != route.tenant_id or message.cell_id != route.cell_id:
raise ScopeError("message claims do not match authoritative route")
return route
def assert_context_route(self, context: TenantContext) -> Route:
route = self.resolve(
channel=context.channel,
channel_account_id=context.channel_account_id,
principal_id=context.principal_id,
session_id=context.session_id,
)
if context.tenant_id != route.tenant_id or context.cell_id != route.cell_id:
raise ScopeError("context claims do not match authoritative route")
return route