diff --git a/backend/app/controllers/api/chat.py b/backend/app/controllers/api/chat.py index 02a2349..af32475 100644 --- a/backend/app/controllers/api/chat.py +++ b/backend/app/controllers/api/chat.py @@ -102,7 +102,7 @@ def uow_factory() -> ChatUnitOfWork: return ChatOrchestrationService( uow_factory=uow_factory, - usage_limiter=ChatUsageLimiter(uow_factory, daily_query_cap=settings.chat_daily_query_cap), + usage_limiter=ChatUsageLimiter(daily_query_cap=settings.chat_daily_query_cap), audit_service=ChatAuditService(uow_factory), agent_port=_agent_port(), sql_executor=ReadOnlySqlExecutor(), diff --git a/backend/app/services/chat/chat_orchestration_service.py b/backend/app/services/chat/chat_orchestration_service.py index a358a95..6757e9c 100644 --- a/backend/app/services/chat/chat_orchestration_service.py +++ b/backend/app/services/chat/chat_orchestration_service.py @@ -10,6 +10,8 @@ from datetime import datetime from typing import TYPE_CHECKING, Protocol +from sqlalchemy import text + from app.core.clock import MARKET_TIMEZONE from app.integrations.openai.llm_agent_port import ( ChatCompletedEvent, @@ -24,6 +26,8 @@ from collections.abc import Callable, Iterator from types import TracebackType + from sqlalchemy.orm import Session + from app.integrations.openai.llm_agent_port import ( ChatStreamEvent, LlmAgentPort, @@ -72,6 +76,12 @@ class ChatOrchestrationUnitOfWork(Protocol): def chat_sessions(self) -> ChatSessionRepository: ... @property def chat_messages(self) -> ChatMessageRepository: ... + @property + def session(self) -> Session: + """The raw SQLAlchemy session -- `begin_turn()` uses it directly for the transaction-scoped + advisory lock (`pg_advisory_xact_lock`) that closes the daily-cap race (S11 §5.3).""" + ... + def commit(self) -> None: ... def __enter__(self) -> ChatOrchestrationUnitOfWork: ... def __exit__( @@ -103,9 +113,19 @@ def __init__( def begin_turn( self, *, session_id: uuid.UUID, customer_id: uuid.UUID, message_text: str ) -> BegunTurn: - self._usage_limiter.check(customer_id) - with self._uow_factory() as uow: + # Transaction-scoped advisory lock, released automatically on commit or rollback -- + # closes the daily-cap check-then-insert race (chat security audit, 2026-09-07) by + # serializing this customer's concurrent begin_turn() calls, across sessions too. The + # lock covers only this bookkeeping transaction, never the slow, network-bound LLM + # call in stream_turn() -- a customer's second tab still works while the first + # streams, only the "am I allowed to start a turn" decision is serialized. + uow.session.execute( + text("SELECT pg_advisory_xact_lock(hashtext(:customer_id)::bigint)"), + {"customer_id": str(customer_id)}, + ) + self._usage_limiter.check(uow, customer_id) + history = tuple( ConversationTurn(role=row.role.value, content=row.content) for row in uow.chat_messages.list_for_session(session_id) diff --git a/backend/app/services/chat/chat_usage_limiter.py b/backend/app/services/chat/chat_usage_limiter.py index db32575..7586b6d 100644 --- a/backend/app/services/chat/chat_usage_limiter.py +++ b/backend/app/services/chat/chat_usage_limiter.py @@ -1,4 +1,13 @@ -"""Per-customer daily chat query cap, checked before any model call (S11 §5.2 step 1, NFR-16).""" +"""Per-customer daily chat query cap, checked before any model call (S11 §5.2 step 1, NFR-16). + +`check()` takes an already-open unit of work rather than opening its own -- a security-review +finding (chat audit, 2026-09-07): the previous version ran its `SELECT count(...)` in its own, +separate transaction that committed and closed *before* `begin_turn()`'s own transaction inserted +the new message rows. Under READ COMMITTED (Postgres's default), two concurrent requests could +each read a count just under the cap and both proceed -- no reordering of the count-then-insert +sequence closes that race on its own; it needs the two to share a transaction that a lock +serializes. See `ChatOrchestrationService.begin_turn()`, which now acquires a +`pg_advisory_xact_lock` keyed on the customer before calling this.""" from __future__ import annotations @@ -7,8 +16,6 @@ if TYPE_CHECKING: import uuid - from collections.abc import Callable - from types import TracebackType class _ChatMessageRepoProtocol(Protocol): @@ -16,17 +23,11 @@ def count_for_customer_since(self, customer_id: uuid.UUID, since: datetime) -> i class ChatUsageLimiterUnitOfWork(Protocol): - """Satisfied by `ChatUnitOfWork`.""" + """Satisfied by `ChatUnitOfWork`. No `__enter__`/`__exit__` here -- the caller owns the + transaction (and the advisory lock inside it); this type only reads through it.""" @property def chat_messages(self) -> _ChatMessageRepoProtocol: ... - def __enter__(self) -> ChatUsageLimiterUnitOfWork: ... - def __exit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: ... class DailyQueryCapExceededError(Exception): @@ -34,17 +35,21 @@ class DailyQueryCapExceededError(Exception): class ChatUsageLimiter: - def __init__( - self, uow_factory: Callable[[], ChatUsageLimiterUnitOfWork], *, daily_query_cap: int - ) -> None: - self._uow_factory = uow_factory + def __init__(self, *, daily_query_cap: int) -> None: self._daily_query_cap = daily_query_cap - def check(self, customer_id: uuid.UUID, *, now: datetime | None = None) -> None: - """Raises `DailyQueryCapExceededError` if the cap was hit in the trailing 24 hours.""" + def check( + self, + uow: ChatUsageLimiterUnitOfWork, + customer_id: uuid.UUID, + *, + now: datetime | None = None, + ) -> None: + """Raises `DailyQueryCapExceededError` if the cap was hit in the trailing 24 hours. Must + run inside a transaction that already holds this customer's advisory lock -- otherwise + this reintroduces the exact race the caller-owned-transaction signature exists to close.""" since = (now or datetime.now(UTC)) - timedelta(hours=24) - with self._uow_factory() as uow: - count = uow.chat_messages.count_for_customer_since(customer_id, since) + count = uow.chat_messages.count_for_customer_since(customer_id, since) if count >= self._daily_query_cap: raise DailyQueryCapExceededError( f"daily chat query cap ({self._daily_query_cap}) reached for this customer" diff --git a/backend/tests/integration/test_chat_usage_cap_race.py b/backend/tests/integration/test_chat_usage_cap_race.py new file mode 100644 index 0000000..cf2b8eb --- /dev/null +++ b/backend/tests/integration/test_chat_usage_cap_race.py @@ -0,0 +1,171 @@ +"""S11 chat security audit (2026-09-07), medium finding -- the daily-cap check-then-insert race. + +Before the fix, `ChatUsageLimiter.check()` ran its own `SELECT count(...)` in a transaction that +committed and closed before `ChatOrchestrationService.begin_turn()`'s transaction inserted the new +message row. Under READ COMMITTED (Postgres's default), two concurrent `begin_turn()` calls for +the same customer could each read a count just under the cap and both proceed. The fix acquires a +`pg_advisory_xact_lock` keyed on the customer as the first statement of `begin_turn()`'s +transaction, serializing the cap check against any other in-flight `begin_turn()` for that +customer -- across sessions, not just within one. + +This test proves that serialization directly: one thread holds the same lock key while inserting +the message that fills the cap, a second thread's `begin_turn()` call is shown to genuinely block +on that lock (not just get lucky with scheduling) and, once unblocked, correctly sees the +just-committed row and rejects -- the race is closed, not merely less likely. +""" + +from __future__ import annotations + +import threading +import time +import uuid +from typing import TYPE_CHECKING + +import pytest +from sqlalchemy import text + +from app.config import Settings +from app.core.uow import SessionRole +from app.extensions import dispose_engines, init_engines +from app.integrations.openai.fake_agent_adapter import FakeAgentAdapter +from app.models.chat.chat_message import ChatMessage, ChatMessageRole +from app.models.chat.chat_session import ChatSession +from app.models.chat.chat_tool_call import ChatToolCall +from app.models.identity.customer import AccountApprovalStatus, Customer, KycStatus +from app.services.chat.chat_audit_service import ChatAuditService +from app.services.chat.chat_orchestration_service import ChatOrchestrationService +from app.services.chat.chat_usage_limiter import ChatUsageLimiter, DailyQueryCapExceededError +from app.services.chat.read_only_sql_executor import ReadOnlySqlExecutor +from app.services.chat.uow import ChatUnitOfWork + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Engine + from sqlalchemy.orm import Session + +_TABLES = [Customer.__table__, ChatSession.__table__, ChatMessage.__table__, ChatToolCall.__table__] + +pytestmark = pytest.mark.usefixtures("_tables") + + +@pytest.fixture(autouse=True) +def _engines(test_settings: Settings) -> Iterator[None]: + init_engines(test_settings) + yield None + dispose_engines() + + +@pytest.fixture +def _tables(owner_engine: Engine) -> Iterator[None]: + for table in _TABLES: + table.create(bind=owner_engine, checkfirst=True) + yield None + with owner_engine.begin() as connection: + for table in reversed(_TABLES): + connection.execute(text(f'DROP TABLE IF EXISTS "{table.name}" CASCADE')) + + +def _customer_with_two_sessions(db_committing: Session) -> tuple[uuid.UUID, uuid.UUID, uuid.UUID]: + customer = Customer( + email=f"{uuid.uuid4()}@trueup.test", + password_hash="hash", + kyc_status=KycStatus.approved, + account_approval_status=AccountApprovalStatus.approved, + ) + db_committing.add(customer) + db_committing.flush() + session_a = ChatSession(customer_id=customer.id) + session_b = ChatSession(customer_id=customer.id) + db_committing.add(session_a) + db_committing.add(session_b) + db_committing.commit() + return customer.id, session_a.id, session_b.id + + +def _build_orchestration( + customer_id: uuid.UUID, *, daily_query_cap: int +) -> ChatOrchestrationService: + def uow_factory() -> ChatUnitOfWork: + return ChatUnitOfWork(customer_id=customer_id, role=SessionRole.CUSTOMER) + + return ChatOrchestrationService( + uow_factory=uow_factory, + usage_limiter=ChatUsageLimiter(daily_query_cap=daily_query_cap), + audit_service=ChatAuditService(uow_factory), + agent_port=FakeAgentAdapter(), + sql_executor=ReadOnlySqlExecutor(), + max_tool_iterations=6, + ) + + +def test_concurrent_begin_turn_calls_at_the_cap_boundary_cannot_both_succeed( + db_committing: Session, +) -> None: + customer_id, session_a, session_b = _customer_with_two_sessions(db_committing) + + hold_seconds = 0.4 + lock_acquired = threading.Event() + second_call_blocked_for: list[float] = [] + second_call_outcome: list[str] = [] + + def hold_the_customer_lock_and_fill_the_cap() -> None: + with ChatUnitOfWork(customer_id=customer_id, role=SessionRole.CUSTOMER) as uow: + # Same lock key `ChatOrchestrationService.begin_turn()` acquires. + uow.session.execute( + text("SELECT pg_advisory_xact_lock(hashtext(:customer_id)::bigint)"), + {"customer_id": str(customer_id)}, + ) + lock_acquired.set() + time.sleep(hold_seconds) + uow.chat_messages.add( + ChatMessage( + session_id=session_a, role=ChatMessageRole.USER, content="first question" + ) + ) + uow.commit() + + def attempt_begin_turn_at_the_cap() -> None: + orchestration = _build_orchestration(customer_id, daily_query_cap=1) + started_at = time.monotonic() + try: + orchestration.begin_turn( + session_id=session_b, customer_id=customer_id, message_text="second question" + ) + second_call_outcome.append("began") + except DailyQueryCapExceededError: + second_call_outcome.append("rejected") + second_call_blocked_for.append(time.monotonic() - started_at) + + holder_thread = threading.Thread(target=hold_the_customer_lock_and_fill_the_cap) + holder_thread.start() + assert lock_acquired.wait(timeout=5), "holder thread never acquired the advisory lock" + + challenger_thread = threading.Thread(target=attempt_begin_turn_at_the_cap) + challenger_thread.start() + holder_thread.join(timeout=5) + challenger_thread.join(timeout=5) + + assert len(second_call_outcome) == 1, "challenger thread never completed" + # Assert on the challenger's own block duration -- a genuine cross-connection lock wait, not a + # Python thread-scheduling coincidence. Tolerance below the full hold; near-instant would mean + # the two transactions interleaved instead of serializing. + assert second_call_blocked_for[0] >= hold_seconds * 0.75 + # Unblocked only after the holder committed the capping row -- so the challenger's own count + # check sees it and correctly rejects, closing the race rather than narrowing its window. + assert second_call_outcome == ["rejected"] + + +def test_second_session_can_begin_a_turn_once_the_first_is_under_the_cap( + db_committing: Session, +) -> None: + """Sanity check alongside the race test: the lock serializes, it does not wrongly block a + second session's turn when the customer is still under the cap.""" + customer_id, _session_a, session_b = _customer_with_two_sessions(db_committing) + + orchestration = _build_orchestration(customer_id, daily_query_cap=5) + begun = orchestration.begin_turn( + session_id=session_b, customer_id=customer_id, message_text="only question" + ) + + assert begun.session_id == session_b