From 85d32d06d9c804c0554d1525aa5fd5f504fd19c7 Mon Sep 17 00:00:00 2001 From: BADRI NARAYANAN Date: Sun, 6 Sep 2026 18:53:07 -0700 Subject: [PATCH] fix(chat): close the daily-cap check-then-insert race with an advisory lock Security audit (2026-09-07) found ChatUsageLimiter.check() ran its own SELECT count(...) in a transaction that committed before begin_turn()'s transaction inserted the new message rows. Under READ COMMITTED, two concurrent requests for the same customer could each read a count just under the daily cap and both proceed. begin_turn() now acquires a transaction-scoped pg_advisory_xact_lock keyed on the customer as the first statement of its transaction, so the cap check and the message insert are serialized against any other in-flight begin_turn() for that customer, across sessions. The lock is released automatically on commit/rollback and covers only the fast bookkeeping transaction, never the slow LLM stream, so a second tab still works while the first is streaming. ChatUsageLimiter.check() now takes the caller's already-open unit of work instead of opening its own, so the count read and the message insert share one transaction under one lock. Adds an integration test that proves the serialization directly: one thread holds the lock while inserting the capping message, a second thread's begin_turn() call is shown to genuinely block on that lock and, once unblocked, correctly rejects. --- backend/app/controllers/api/chat.py | 2 +- .../chat/chat_orchestration_service.py | 24 ++- .../app/services/chat/chat_usage_limiter.py | 43 +++-- .../integration/test_chat_usage_cap_race.py | 171 ++++++++++++++++++ 4 files changed, 218 insertions(+), 22 deletions(-) create mode 100644 backend/tests/integration/test_chat_usage_cap_race.py 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