fix(sdk): coordinate concurrent token refresh

crates/services/src/auth/cache.rs

@@ -13,10 +13,11 @@ const TTL: Duration = Duration::from_secs(30);
 const MAX_ENTRIES: usize = 1_024;
 const MAX_BYTES: usize = 4 * 1024 * 1024;
 const MAX_ENTRY_BYTES: usize = 16 * 1024;
-// Remove/insert churn can retain 4096 buckets for 1024 live entries. As in
-// the model cache, reserve keys, values and structural overhead for that
-// allocation, including storage retained by invalidate().
-const CONTAINER_BYTES: usize = MAX_ENTRIES * 4 * (std::mem::size_of::<([u8; 32], Entry)>() + 64);
+// Reserve entry and session-fence buckets, including allocations retained
+// after invalidation. Churn can retain four buckets per live entry.
+const CONTAINER_BYTES: usize = MAX_ENTRIES
+    * 4
+    * (std::mem::size_of::<([u8; 32], Entry)>() + std::mem::size_of::<([u8; 32], u64)>() + 128);
 
 struct Entry {
     principal: AuthPrincipal,
@@ -34,6 +35,8 @@ impl Entry {
 #[derive(Default)]
 struct State {
     generation: u64,
+    global_invalidation: u64,
+    session_invalidations: HashMap<[u8; 32], u64>,
     bytes: usize,
     entries: HashMap<[u8; 32], Entry>,
 }
@@ -46,6 +49,14 @@ impl AuthenticationCache {
         Sha256::digest(token.as_bytes()).into()
     }
 
+    fn session_key(user_id: &str, session_id: &str) -> [u8; 32] {
+        let mut hash = Sha256::new();
+        hash.update((user_id.len() as u64).to_be_bytes());
+        hash.update(user_id.as_bytes());
+        hash.update(session_id.as_bytes());
+        hash.finalize().into()
+    }
+
     pub(super) fn lookup(&self, key: &[u8; 32]) -> (Option<AuthPrincipal>, u64) {
         let mut state = self
             .0
@@ -80,10 +91,18 @@ impl AuthenticationCache {
             .unwrap_or_else(|poisoned| poisoned.into_inner());
         let now = Instant::now();
         let now_unix = chrono::Utc::now().timestamp();
-        if state.generation != generation
-            || now_unix >= expires_at_unix
-            || now.duration_since(checked_at) >= TTL
-        {
+        if state.generation != generation {
+            let session_invalidated = principal.session_id.as_deref().is_some_and(|session| {
+                state
+                    .session_invalidations
+                    .get(&Self::session_key(&principal.user.user_id, session))
+                    .is_some_and(|last| *last > generation)
+            });
+            if generation < state.global_invalidation || session_invalidated {
+                return false;
+            }
+        }
+        if now_unix >= expires_at_unix || now.duration_since(checked_at) >= TTL {
             return false;
         }
         let bytes = principal_bytes(principal);
@@ -125,28 +144,71 @@ impl AuthenticationCache {
         true
     }
 
-    fn invalidate(&self) {
+    fn invalidate(&self, session: Option<(&str, &str)>) {
         let mut state = self
             .0
             .lock()
             .unwrap_or_else(|poisoned| poisoned.into_inner());
         state.generation = state.generation.wrapping_add(1);
-        state.entries.clear();
-        state.bytes = 0;
+        if let Some((user_id, session_id)) = session {
+            let key = Self::session_key(user_id, session_id);
+            // Bound the fence metadata. On overflow, fence all older fills
+            // before discarding their session-specific invalidations.
+            if !state.session_invalidations.contains_key(&key)
+                && state.session_invalidations.len() >= MAX_ENTRIES
+            {
+                state.global_invalidation = state.generation;
+                state.session_invalidations.clear();
+            }
+            let generation = state.generation;
+            state.session_invalidations.insert(key, generation);
+            state.entries.retain(|_, entry| {
+                entry.principal.user.user_id != user_id
+                    || entry.principal.session_id.as_deref() != Some(session_id)
+            });
+            state.bytes = state.entries.values().map(|entry| entry.bytes).sum();
+        } else {
+            state.global_invalidation = state.generation;
+            state.session_invalidations.clear();
+            state.entries.clear();
+            state.bytes = 0;
+        }
     }
 
     pub(super) fn mutation(self: &Arc<Self>) -> InvalidateOnDrop {
-        self.invalidate();
-        InvalidateOnDrop(Arc::clone(self))
+        self.invalidate(None);
+        InvalidateOnDrop {
+            cache: Arc::clone(self),
+            session: None,
+        }
+    }
+
+    pub(super) fn session_mutation(
+        self: &Arc<Self>,
+        user_id: &str,
+        session_id: &str,
+    ) -> InvalidateOnDrop {
+        self.invalidate(Some((user_id, session_id)));
+        InvalidateOnDrop {
+            cache: Arc::clone(self),
+            session: Some((user_id.to_owned(), session_id.to_owned())),
+        }
     }
 }
 
-pub(super) struct InvalidateOnDrop(Arc<AuthenticationCache>);
+pub(super) struct InvalidateOnDrop {
+    cache: Arc<AuthenticationCache>,
+    session: Option<(String, String)>,
+}
 
 impl Drop for InvalidateOnDrop {
     fn drop(&mut self) {
         // Includes errors, cancellation, and unknown commit acknowledgements.
-        self.0.invalidate();
+        self.cache.invalidate(
+            self.session
+                .as_ref()
+                .map(|(user, session)| (user.as_str(), session.as_str())),
+        );
     }
 }
 
@@ -216,6 +278,65 @@ mod tests {
         assert!(cache.lookup(&key).0.is_none());
     }
 
+    #[test]
+    fn session_invalidation_preserves_other_sessions_and_fences_inflight_fills() {
+        let cache = Arc::new(AuthenticationCache::default());
+        let first_key = AuthenticationCache::key("first");
+        let second_key = AuthenticationCache::key("second");
+        let other_key = AuthenticationCache::key("other");
+        let other_fill = AuthenticationCache::key("other-cold-fill");
+        let mut first = principal();
+        first.session_id = Some("first-session".into());
+        let mut second = first.clone();
+        second.session_id = Some("second-session".into());
+        let mut other = first.clone();
+        other.user.user_id = "other-user".into();
+        let now = Instant::now();
+        let expires = chrono::Utc::now().timestamp() + 60;
+        for (key, value) in [
+            (first_key, &first),
+            (second_key, &second),
+            (other_key, &other),
+        ] {
+            assert!(cache.publish(key, 0, now, expires, value));
+        }
+        let mutation = cache.session_mutation("owner", "first-session");
+        assert!(cache.lookup(&first_key).0.is_none());
+        assert_eq!(cache.lookup(&second_key).0, Some(second.clone()));
+        assert_eq!(cache.lookup(&other_key).0, Some(other.clone()));
+        assert!(!cache.publish(first_key, 0, now, expires, &first));
+        let (_, during) = cache.lookup(&first_key);
+        drop(mutation);
+        assert!(!cache.publish(first_key, during, now, expires, &first));
+        for _ in 0..100 {
+            drop(cache.session_mutation("owner", "first-session"));
+        }
+        assert!(
+            cache.publish(other_fill, 0, now, expires, &other),
+            "a different session's in-flight database observation stays cacheable"
+        );
+        assert_eq!(cache.lookup(&second_key).0, Some(second));
+        assert_eq!(cache.lookup(&other_key).0, Some(other));
+    }
+
+    #[test]
+    fn session_fence_overflow_is_bounded_and_fails_closed() {
+        let cache = Arc::new(AuthenticationCache::default());
+        let now = Instant::now();
+        let expires = chrono::Utc::now().timestamp() + 60;
+        for index in 0..=MAX_ENTRIES {
+            drop(cache.session_mutation("owner", &format!("session-{index}")));
+        }
+        let state = cache.0.lock().unwrap();
+        assert!(state.session_invalidations.len() <= MAX_ENTRIES);
+        assert!(state.global_invalidation > 0);
+        drop(state);
+        let key = AuthenticationCache::key("unrelated");
+        assert!(!cache.publish(key, 0, now, expires, &principal()));
+        let (_, current) = cache.lookup(&key);
+        assert!(cache.publish(key, current, now, expires, &principal()));
+    }
+
     #[test]
     fn cache_retention_is_bounded_and_oversized_principals_stay_local() {
         for content_bytes in [0, MAX_ENTRY_BYTES / 2] {

crates/services/src/auth/mod.rs

@@ -1799,33 +1799,26 @@ impl AuthService for DatabaseAuthService {
             .fetch_one(&mut *tx)
             .await
             .map_err(internal_error)?;
-        // A concurrent disconnect revokes this row under the same user lock;
-        // the compare-and-revoke prevents refresh from resurrecting it.
-        let revoked = query(
-            "UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE token_hash = ? AND is_revoked = 0",
+        // Replace the current token under the account lock. The old hash no
+        // longer matches, so refresh and disconnect still have one winner.
+        let rotated = query(
+            "UPDATE auth_refresh_tokens SET token_hash = ?, expires_at = ?, created_at = NOW(6), token_prefix = NULL \
+             WHERE token_hash = ? AND user_id = ? AND session_id = ? AND is_revoked = 0",
         )
+        .bind(&new_refresh_token_hash)
+        .bind(expires_at)
         .bind(&refresh_token_hash)
+        .bind(&user_id)
+        .bind(&session_id)
         .execute(&mut *tx)
         .await
-        .map_err(|e| map_auth_sqlx(e, "refresh.revoke_old_token", Some(&pool)))?;
-        if revoked.rows_affected() != 1 {
+        .map_err(|e| map_auth_sqlx(e, "refresh.rotate_token", Some(&pool)))?;
+        if rotated.rows_affected() != 1 {
             return Err(error_response(
                 StatusCode::UNAUTHORIZED,
                 "Token expired or revoked",
             ));
         }
-        query(
-            "INSERT INTO auth_refresh_tokens (token_id, user_id, session_id, token_hash, expires_at, is_revoked) \
-             VALUES (?, ?, ?, ?, ?, 0)",
-        )
-        .bind(Uuid::new_v4().to_string())
-        .bind(&user.user_id)
-        .bind(&session_id)
-        .bind(&new_refresh_token_hash)
-        .bind(expires_at)
-        .execute(&mut *tx)
-        .await
-        .map_err(|e| map_auth_sqlx(e, "refresh.insert_new_refresh_token", Some(&pool)))?;
         tx.commit()
             .await
             .map_err(|e| map_auth_sqlx(e, "refresh.commit_tx", Some(&pool)))?;
@@ -1861,42 +1854,50 @@ impl AuthService for DatabaseAuthService {
             .clone()
             .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Invalid token"))?;
         let (session_id, _, _) = self.parse_token_session(&claims)?;
+        // JWT validation allows clock-skew leeway; logout proofs do not.
+        if claims.exp <= Utc::now().timestamp() {
+            return Err(error_response(
+                StatusCode::UNAUTHORIZED,
+                "Token expired or revoked",
+            ));
+        }
 
         let pool = self
             .get_pool()
             .await
             .map_err(|e| map_auth_sqlx(e, "auth.get_pool", None))?;
-        let refresh_token_hash = sha256_hex(&request.refresh_token);
-        let stored = self
-            .fetch_refresh_token(&pool, &refresh_token_hash)
+        let _invalidation = self
+            .authentication_cache
+            .session_mutation(&user_id, &session_id);
+        let mut tx = pool
+            .begin()
             .await
-            .map_err(|e| map_auth_sqlx(e, "logout.fetch_refresh_token", Some(&pool)))?
-            .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Token expired or revoked"))?;
-
-        if stored.0 != user_id || stored.2.as_deref() != Some(session_id.as_str()) {
+            .map_err(|e| map_auth_sqlx(e, "logout.begin_tx", Some(&pool)))?;
+        let account = query("SELECT user_id FROM auth_users WHERE user_id = ? FOR UPDATE")
+            .bind(&user_id)
+            .fetch_optional(&mut *tx)
+            .await
+            .map_err(|e| map_auth_sqlx(e, "logout.lock_user", Some(&pool)))?;
+        if account.is_none() {
             return Err(error_response(
                 StatusCode::UNAUTHORIZED,
-                "Token session mismatch",
+                "Token expired or revoked",
             ));
         }
-        if stored.1 < Utc::now().format("%Y-%m-%dT%H:%M:%S").to_string() {
+        // Recheck after waiting for the lock. The signed session claim is the
+        // proof even when a consumed token's old hash has been replaced.
+        if claims.exp <= Utc::now().timestamp() {
             return Err(error_response(
                 StatusCode::UNAUTHORIZED,
                 "Token expired or revoked",
             ));
         }
-
-        let _invalidation = self.authentication_cache.mutation();
-        let mut tx = pool
-            .begin()
-            .await
-            .map_err(|e| map_auth_sqlx(e, "logout.begin_tx", Some(&pool)))?;
-        query("UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE token_hash = ? AND user_id = ?")
-            .bind(&refresh_token_hash)
+        query("UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE user_id = ? AND session_id = ? AND is_revoked = 0")
             .bind(&user_id)
+            .bind(&session_id)
             .execute(&mut *tx)
             .await
-            .map_err(|e| map_auth_sqlx(e, "logout.revoke_submitted_token", Some(&pool)))?;
+            .map_err(|e| map_auth_sqlx(e, "logout.revoke_session", Some(&pool)))?;
 
         tx.commit()
             .await

crates/services/tests/auth_cache_db_it.rs

@@ -29,7 +29,6 @@ struct SnapshotBarrier {
 impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for Queries {
     fn on_event(&self, event: &tracing::Event<'_>, _: tracing_subscriber::layer::Context<'_, S>) {
         if event.metadata().target() == "sqlx::query" {
-            self.count.fetch_add(1, Ordering::SeqCst);
             #[derive(Default)]
             struct Fields(String);
             impl tracing::field::Visit for Fields {
@@ -39,6 +38,11 @@ impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for Queries {
             }
             let mut fields = Fields::default();
             event.record(&mut fields);
+            // Pool expansion can initialize a new connection alongside an
+            // auth query; count application statements, not session setup.
+            if !fields.0.contains("SET sql_mode=") {
+                self.count.fetch_add(1, Ordering::SeqCst);
+            }
             if fields.0.contains("SELECT u.user_id, u.username")
                 && self.snapshot.armed.swap(false, Ordering::SeqCst)
             {
@@ -175,6 +179,26 @@ async fn jwt_authentication_uses_one_cold_query_and_zero_warm_queries_with_local
         .current_principal(&second)
         .await
         .expect("logout does not revoke another device");
+    assert_eq!(
+        count.swap(0, Ordering::SeqCst),
+        0,
+        "another session stays warm"
+    );
+    for _ in 0..3 {
+        service
+            .logout(AuthRefreshRequestData {
+                refresh_token: tokens[0].refresh_token.clone(),
+            })
+            .await
+            .unwrap();
+        count.store(0, Ordering::SeqCst);
+        service.current_principal(&second).await.unwrap();
+        assert_eq!(
+            count.swap(0, Ordering::SeqCst),
+            0,
+            "replayed logout must not flush unrelated sessions"
+        );
+    }
 
     sqlx::query("UPDATE auth_users SET is_active = 0 WHERE user_id = ?")
         .bind(&user.user_id)
@@ -246,8 +270,8 @@ async fn jwt_authentication_uses_one_cold_query_and_zero_warm_queries_with_local
         let subsequent = observer.current_principal(&second).await;
         assert_eq!(
             count.swap(0, Ordering::SeqCst),
-            1,
-            "invalidated fill cannot re-enter cache"
+            usize::from(same_session),
+            "only the revoked session's in-flight fill must be fenced"
         );
         if same_session {
             assert_eq!(