Fix compat token refresh giving back a consumed token

crates/storage-pg/.sqlx/query-4e64540bbffe5f4b9c4a6589012cf69eb67adaa4d40fc1910dfcd2640e32ab37.json

@@ -1,15 +1,16 @@
 {
   "db_name": "PostgreSQL",
-  "query": "\n                UPDATE compat_refresh_tokens\n                SET consumed_at = $2\n                WHERE compat_session_id = $1\n                  AND consumed_at IS NULL\n            ",
+  "query": "\n                UPDATE compat_refresh_tokens\n                SET consumed_at = $2\n                WHERE compat_session_id = $1\n                  AND consumed_at IS NULL\n                  AND compat_refresh_token_id <> $3\n            ",
   "describe": {
     "columns": [],
     "parameters": {
       "Left": [
         "Uuid",
-        "Timestamptz"
+        "Timestamptz",
+        "Uuid"
       ]
     },
     "nullable": []
   },
-  "hash": "f75e44b528234dac708640ad9a111f3f6b468a91bf0d5b574795bf8c80605f19"
+  "hash": "4e64540bbffe5f4b9c4a6589012cf69eb67adaa4d40fc1910dfcd2640e32ab37"
 }

crates/storage-pg/src/compat/mod.rs

@@ -437,6 +437,7 @@ mod tests {
     async fn test_refresh_token_repository(pool: PgPool) {
         const ACCESS_TOKEN: &str = "access_token";
         const REFRESH_TOKEN: &str = "refresh_token";
+        const REFRESH_TOKEN2: &str = "refresh_token2";
         let mut rng = ChaChaRng::seed_from_u64(42);
         let clock = MockClock::default();
         let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed();
@@ -508,16 +509,28 @@ mod tests {
         assert!(refresh_token_lookup.is_valid());
         assert!(!refresh_token_lookup.is_consumed());
 
-        // Consume it
+        // Consume the first token, but to do so we need a 2nd to replace it with
+        let refresh_token2 = repo
+            .compat_refresh_token()
+            .add(
+                &mut rng,
+                &clock,
+                &session,
+                &access_token,
+                REFRESH_TOKEN2.to_owned(),
+            )
+            .await
+            .unwrap();
+
         let refresh_token = repo
             .compat_refresh_token()
-            .consume(&clock, refresh_token)
+            .consume_and_replace(&clock, refresh_token, &refresh_token2)
             .await
             .unwrap();
         assert!(!refresh_token.is_valid());
         assert!(refresh_token.is_consumed());
 
-        // Reload it and check again
+        // Reload the first token and check again
         let refresh_token_lookup = repo
             .compat_refresh_token()
             .find_by_token(REFRESH_TOKEN)
@@ -530,7 +543,7 @@ mod tests {
         // Consuming it again should not work
         assert!(
             repo.compat_refresh_token()
-                .consume(&clock, refresh_token)
+                .consume_and_replace(&clock, refresh_token, &refresh_token2)
                 .await
                 .is_err()
         );

crates/storage-pg/src/compat/refresh_token.rs

@@ -185,20 +185,26 @@ impl CompatRefreshTokenRepository for PgCompatRefreshTokenRepository<'_> {
     }
 
     #[tracing::instrument(
-        name = "db.compat_refresh_token.consume",
+        name = "db.compat_refresh_token.consume_and_replace",
         skip_all,
         fields(
             db.query.text,
             %compat_refresh_token.id,
+            %successor_compat_refresh_token.id,
             compat_session.id = %compat_refresh_token.session_id,
         ),
         err,
     )]
-    async fn consume(
+    async fn consume_and_replace(
         &mut self,
         clock: &dyn Clock,
         compat_refresh_token: CompatRefreshToken,
+        successor_compat_refresh_token: &CompatRefreshToken,
     ) -> Result<CompatRefreshToken, Self::Error> {
+        if compat_refresh_token.session_id != successor_compat_refresh_token.session_id {
+            return Err(DatabaseError::invalid_operation());
+        }
+
         let consumed_at = clock.now();
         let res = sqlx::query!(
             r#"
@@ -206,9 +212,11 @@ impl CompatRefreshTokenRepository for PgCompatRefreshTokenRepository<'_> {
                 SET consumed_at = $2
                 WHERE compat_session_id = $1
                   AND consumed_at IS NULL
+                  AND compat_refresh_token_id <> $3
             "#,
             Uuid::from(compat_refresh_token.session_id),
             consumed_at,
+            Uuid::from(successor_compat_refresh_token.id),
         )
         .traced()
         .execute(&mut *self.conn)