From 1dcd1e2ec45f7f5f4eeb81f5e2f4b2bb8480e117 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Mon, 4 May 2026 13:09:39 -0700 Subject: [PATCH] fix(oidc): serialise role-mapping concurrency + skip no-op write lock (cumulative bug-2, perf-1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit bug-2 (Postgres) — replace_oidc_roles read existing rows under default READ COMMITTED with no row lock. Two concurrent OIDC callbacks for the same user_id (racing token refreshes with differing claim sets) could both observe the same baseline and produce a final role state matching neither caller's intent. Adds .with_for_update() to the SELECT so the existing rows for this user are locked for the duration of the transaction. The lock is per-user_id, not table-wide; unrelated user writes are unaffected. Empty result sets acquire no locks, so a brand-new user with no rows yet still allows two callers to proceed and merge via ON CONFLICT DO NOTHING — that's a permissive race that self-heals on the next reconciliation cycle, documented in code. perf-1 (SQLite) — replace_oidc_roles took the SQLite global write lock unconditionally via BEGIN IMMEDIATE before reading. Steady-state re-logins (claims unchanged, no INSERT/DELETE needed) paid the lock cost for nothing and serialised against unrelated writers. Replaces with a double-check pattern: phase 1 reads under the default deferred transaction (no write lock), computes the diff, and returns (set(), set()) on no-op. Phase 2, only when mutation is needed, commits the read txn, escalates to BEGIN IMMEDIATE, RE-READS, and re-computes the diff under the lock before writing. The returned (added, removed) reflects what was actually written, so caller logging in apply_role_mapping stays truthful even when concurrent writers shifted state between the two reads. The OR IGNORE on insert is now defense-in-depth (the lock makes it unnecessary) but kept as a safety net. (cherry picked from commit d5087ef3b94ccc6dc67453370b2cdc8ecd8b188f) --- tests/test_oidc_storage.py | 47 +++++++++++++++++++++++++++ turnstone/core/storage/_postgresql.py | 21 ++++++++++-- turnstone/core/storage/_sqlite.py | 38 +++++++++++++++++++++- 3 files changed, 102 insertions(+), 4 deletions(-) diff --git a/tests/test_oidc_storage.py b/tests/test_oidc_storage.py index 6f92e3ae..f15b9d82 100644 --- a/tests/test_oidc_storage.py +++ b/tests/test_oidc_storage.py @@ -550,3 +550,50 @@ class TestReplaceOIDCRoles: assert removed == {"role-oidc-old"} roles = {r["role_id"]: r["assigned_by"] for r in db.list_user_roles("u1")} assert roles == {"role-manual": "admin-ui", "role-default": "oidc-default"} + + def test_replace_oidc_roles_no_op_steady_state(self, db): + """Steady-state re-login: claims unchanged, function must short-circuit. + + This pins the contract that drives the SQLite optimistic-read fast + path — the common case (token refresh with identical role claims) + must not acquire a write lock. + """ + db.create_user("u1", "alice", "Alice", "h") + self._seed_role(db, "role-a") + self._seed_role(db, "role-b") + db.assign_role("u1", "role-a", "oidc") + db.assign_role("u1", "role-b", "oidc") + + added, removed = db.replace_oidc_roles("u1", {"role-a", "role-b"}) + + assert added == set() + assert removed == set() + # All rows still oidc-assigned with identical membership. + roles = {r["role_id"]: r["assigned_by"] for r in db.list_user_roles("u1")} + assert roles == {"role-a": "oidc", "role-b": "oidc"} + + def test_replace_oidc_roles_returns_post_lock_diff(self, db): + """Returned (added, removed) reflects the post-lock state, not the optimistic read. + + The SQLite implementation re-reads under the write lock to defend + against races; the values returned must come from that re-read so + callers (apply_role_mapping audit logs) see the actual transition + that hit the table. Steady-state input must collapse to empty + sets and leave row timestamps unchanged. + """ + db.create_user("u1", "alice", "Alice", "h") + self._seed_role(db, "role-a") + db.assign_role("u1", "role-a", "oidc") + + before = db.list_user_roles("u1") + assert len(before) == 1 + original_created = before[0]["assignment_created"] + + added, removed = db.replace_oidc_roles("u1", {"role-a"}) + + assert added == set() + assert removed == set() + # No write occurred — the assignment row's timestamp is untouched. + after = db.list_user_roles("u1") + assert len(after) == 1 + assert after[0]["assignment_created"] == original_created diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py index 8ae4f559..68fb9a96 100644 --- a/turnstone/core/storage/_postgresql.py +++ b/turnstone/core/storage/_postgresql.py @@ -2105,10 +2105,25 @@ class PostgreSQLBackend: now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") with self._conn() as conn: + # `with_for_update()` takes per-row locks on this user's + # `user_roles` rows for the duration of the transaction. + # Two concurrent OIDC callbacks for the same `user_id` (e.g. + # racing token refreshes with differing claim sets) would + # otherwise both read the same baseline under READ COMMITTED + # and produce a final role state matching neither caller's + # intent. The lock is per-`user_id`, so unrelated user writes + # are unaffected. + # + # Note: FOR UPDATE on an empty result set acquires no locks, + # so on a brand-new user with no rows yet, two concurrent + # callers can proceed in parallel; their inserts merge via + # ON CONFLICT DO NOTHING (final state is the union of the + # two desired sets). The next single-caller reconciliation + # cycle self-heals. existing_rows = conn.execute( - sa.select(user_roles.c.role_id, user_roles.c.assigned_by).where( - user_roles.c.user_id == user_id - ) + sa.select(user_roles.c.role_id, user_roles.c.assigned_by) + .where(user_roles.c.user_id == user_id) + .with_for_update() ).fetchall() current_oidc: set[str] = {r[0] for r in existing_rows if r[1] == "oidc"} # Roles assigned by any other source (admin-ui, oidc-default, etc.) diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py index e95dc93a..ce266372 100644 --- a/turnstone/core/storage/_sqlite.py +++ b/turnstone/core/storage/_sqlite.py @@ -2238,9 +2238,24 @@ class SQLiteBackend: def replace_oidc_roles( self, user_id: str, desired_role_ids: set[str] ) -> tuple[set[str], set[str]]: + # Double-check pattern: the steady-state OIDC re-login (claims + # unchanged from the last login) is overwhelmingly the common + # case, and SQLite's `BEGIN IMMEDIATE` takes the database-wide + # write lock — serialising every unrelated writer in the + # process. Acquiring it for a no-op diff is pure waste. + # + # Phase 1 reads under the default deferred transaction (no + # write lock) and bails out cheaply when the diff is empty. + # Phase 2 only fires when work is needed: it commits the read + # txn, escalates to `BEGIN IMMEDIATE`, and re-reads + re-diffs + # under the lock. The re-read is required — between the two + # reads any concurrent writer (admin-ui assignment, another + # racing OIDC callback) could have changed the row set, and + # acting on the optimistic snapshot would clobber that change. + # The post-lock diff is what we return so callers see the + # actual transition that hit the table. now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") with self._conn() as conn: - conn.execute(sa.text("BEGIN IMMEDIATE")) existing_rows = conn.execute( sa.select(user_roles.c.role_id, user_roles.c.assigned_by).where( user_roles.c.user_id == user_id @@ -2251,6 +2266,27 @@ class SQLiteBackend: # are off-limits to OIDC reconciliation per apply_role_mapping's contract. blocked: set[str] = {r[0] for r in existing_rows if r[1] != "oidc"} + effective_desired = desired_role_ids - blocked + added = effective_desired - current_oidc + removed = current_oidc - effective_desired + if not added and not removed: + # Steady state: claims unchanged from prior login. Skip + # the write lock entirely — this is the perf win. + return set(), set() + + # Mutation needed. Release the implicit read txn, acquire + # the SQLite write lock, and re-read so the diff reflects + # any state change that landed between the two reads. + conn.commit() + conn.execute(sa.text("BEGIN IMMEDIATE")) + existing_rows = conn.execute( + sa.select(user_roles.c.role_id, user_roles.c.assigned_by).where( + user_roles.c.user_id == user_id + ) + ).fetchall() + current_oidc = {r[0] for r in existing_rows if r[1] == "oidc"} + blocked = {r[0] for r in existing_rows if r[1] != "oidc"} + effective_desired = desired_role_ids - blocked added = effective_desired - current_oidc removed = current_oidc - effective_desired