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