From 6f3e3a13b8454d7a19b494aa6bc8e92fa0947974 Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Tue, 10 Dec 2024 14:06:06 +0100
Subject: [PATCH 1/7] Record the next refresh token ID when refreshing

This will help us determine whether we had a double-refresh happening
---
 crates/data-model/src/tokens.rs               | 32 ++++++++++--
 crates/handlers/src/oauth2/token.rs           |  2 +-
 ...af36c103e030358262e46a4bd5e4006dc630.json} | 12 +++--
 ...019c36bf5983182b0871de6e85036d8df466.json} | 12 +++--
 ...29341848fce336d339b6bbf425956f5ed5032.json | 15 ------
 ...2980b5fe575ae8432a000e9c4e4307caa2d9b.json | 16 ++++++
 ...0115428_oauth_refresh_token_track_next.sql |  9 ++++
 crates/storage-pg/src/oauth2/mod.rs           | 15 +++++-
 crates/storage-pg/src/oauth2/refresh_token.rs | 49 ++++++++++++++-----
 crates/storage/src/oauth2/refresh_token.rs    |  3 ++
 10 files changed, 125 insertions(+), 40 deletions(-)
 rename crates/storage-pg/.sqlx/{query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json => query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json} (74%)
 rename crates/storage-pg/.sqlx/{query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json => query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json} (74%)
 delete mode 100644 crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json
 create mode 100644 crates/storage-pg/.sqlx/query-ffbfef8b7e72ec4bae02b6bbe862980b5fe575ae8432a000e9c4e4307caa2d9b.json
 create mode 100644 crates/storage-pg/migrations/20241210115428_oauth_refresh_token_track_next.sql

diff --git a/crates/data-model/src/tokens.rs b/crates/data-model/src/tokens.rs
index b4c8ede28..210eae0cc 100644
--- a/crates/data-model/src/tokens.rs
+++ b/crates/data-model/src/tokens.rs
@@ -109,6 +109,7 @@ pub enum RefreshTokenState {
     Valid,
     Consumed {
         consumed_at: DateTime<Utc>,
+        next_refresh_token_id: Option<Ulid>,
     },
 }
 
@@ -118,9 +119,16 @@ impl RefreshTokenState {
     /// # Errors
     ///
     /// Returns an error if the refresh token is already consumed.
-    fn consume(self, consumed_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+    fn consume(
+        self,
+        consumed_at: DateTime<Utc>,
+        replaced_by: &RefreshToken,
+    ) -> Result<Self, InvalidTransitionError> {
         match self {
-            Self::Valid => Ok(Self::Consumed { consumed_at }),
+            Self::Valid => Ok(Self::Consumed {
+                consumed_at,
+                next_refresh_token_id: Some(replaced_by.id),
+            }),
             Self::Consumed { .. } => Err(InvalidTransitionError),
         }
     }
@@ -140,6 +148,18 @@ impl RefreshTokenState {
     pub fn is_consumed(&self) -> bool {
         matches!(self, Self::Consumed { .. })
     }
+
+    /// Returns the next refresh token ID, if any.
+    #[must_use]
+    pub fn next_refresh_token_id(&self) -> Option<Ulid> {
+        match self {
+            Self::Valid => None,
+            Self::Consumed {
+                next_refresh_token_id,
+                ..
+            } => *next_refresh_token_id,
+        }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -171,8 +191,12 @@ impl RefreshToken {
     /// # Errors
     ///
     /// Returns an error if the refresh token is already consumed.
-    pub fn consume(mut self, consumed_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
-        self.state = self.state.consume(consumed_at)?;
+    pub fn consume(
+        mut self,
+        consumed_at: DateTime<Utc>,
+        replaced_by: &Self,
+    ) -> Result<Self, InvalidTransitionError> {
+        self.state = self.state.consume(consumed_at, replaced_by)?;
         Ok(self)
     }
 }
diff --git a/crates/handlers/src/oauth2/token.rs b/crates/handlers/src/oauth2/token.rs
index a228f746b..1222b32e5 100644
--- a/crates/handlers/src/oauth2/token.rs
+++ b/crates/handlers/src/oauth2/token.rs
@@ -544,7 +544,7 @@ async fn refresh_token_grant(
 
     let refresh_token = repo
         .oauth2_refresh_token()
-        .consume(clock, refresh_token)
+        .consume(clock, refresh_token, &new_refresh_token)
         .await?;
 
     if let Some(access_token_id) = refresh_token.access_token_id {
diff --git a/crates/storage-pg/.sqlx/query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json b/crates/storage-pg/.sqlx/query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json
similarity index 74%
rename from crates/storage-pg/.sqlx/query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json
rename to crates/storage-pg/.sqlx/query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json
index 411b02cbe..df5d58132 100644
--- a/crates/storage-pg/.sqlx/query-e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671.json
+++ b/crates/storage-pg/.sqlx/query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json
@@ -1,6 +1,6 @@
 {
   "db_name": "PostgreSQL",
-  "query": "\n                SELECT oauth2_refresh_token_id\n                     , refresh_token\n                     , created_at\n                     , consumed_at\n                     , oauth2_access_token_id\n                     , oauth2_session_id\n                FROM oauth2_refresh_tokens\n\n                WHERE refresh_token = $1\n            ",
+  "query": "\n                SELECT oauth2_refresh_token_id\n                     , refresh_token\n                     , created_at\n                     , consumed_at\n                     , oauth2_access_token_id\n                     , oauth2_session_id\n                     , next_oauth2_refresh_token_id\n                FROM oauth2_refresh_tokens\n\n                WHERE refresh_token = $1\n            ",
   "describe": {
     "columns": [
       {
@@ -32,6 +32,11 @@
         "ordinal": 5,
         "name": "oauth2_session_id",
         "type_info": "Uuid"
+      },
+      {
+        "ordinal": 6,
+        "name": "next_oauth2_refresh_token_id",
+        "type_info": "Uuid"
       }
     ],
     "parameters": {
@@ -45,8 +50,9 @@
       false,
       true,
       true,
-      false
+      false,
+      true
     ]
   },
-  "hash": "e709869c062ac50248b1f9f8f808cc2f5e7bef58a6c2e42a7bb0c1cb8f508671"
+  "hash": "265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630"
 }
diff --git a/crates/storage-pg/.sqlx/query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json b/crates/storage-pg/.sqlx/query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json
similarity index 74%
rename from crates/storage-pg/.sqlx/query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json
rename to crates/storage-pg/.sqlx/query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json
index bfa6485e8..be9fa3e32 100644
--- a/crates/storage-pg/.sqlx/query-a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20.json
+++ b/crates/storage-pg/.sqlx/query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json
@@ -1,6 +1,6 @@
 {
   "db_name": "PostgreSQL",
-  "query": "\n                SELECT oauth2_refresh_token_id\n                     , refresh_token\n                     , created_at\n                     , consumed_at\n                     , oauth2_access_token_id\n                     , oauth2_session_id\n                FROM oauth2_refresh_tokens\n\n                WHERE oauth2_refresh_token_id = $1\n            ",
+  "query": "\n                SELECT oauth2_refresh_token_id\n                     , refresh_token\n                     , created_at\n                     , consumed_at\n                     , oauth2_access_token_id\n                     , oauth2_session_id\n                     , next_oauth2_refresh_token_id\n                FROM oauth2_refresh_tokens\n\n                WHERE oauth2_refresh_token_id = $1\n            ",
   "describe": {
     "columns": [
       {
@@ -32,6 +32,11 @@
         "ordinal": 5,
         "name": "oauth2_session_id",
         "type_info": "Uuid"
+      },
+      {
+        "ordinal": 6,
+        "name": "next_oauth2_refresh_token_id",
+        "type_info": "Uuid"
       }
     ],
     "parameters": {
@@ -45,8 +50,9 @@
       false,
       true,
       true,
-      false
+      false,
+      true
     ]
   },
-  "hash": "a6fa7811d0a7c62c7cccff96dc82db5b25462fa7669fde1941ccab4712585b20"
+  "hash": "5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466"
 }
diff --git a/crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json b/crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json
deleted file mode 100644
index 30aa4caf3..000000000
--- a/crates/storage-pg/.sqlx/query-b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032.json
+++ /dev/null
@@ -1,15 +0,0 @@
-{
-  "db_name": "PostgreSQL",
-  "query": "\n                UPDATE oauth2_refresh_tokens\n                SET consumed_at = $2\n                WHERE oauth2_refresh_token_id = $1\n            ",
-  "describe": {
-    "columns": [],
-    "parameters": {
-      "Left": [
-        "Uuid",
-        "Timestamptz"
-      ]
-    },
-    "nullable": []
-  },
-  "hash": "b6a6f5386dc89e4bc2ce56d578a29341848fce336d339b6bbf425956f5ed5032"
-}
diff --git a/crates/storage-pg/.sqlx/query-ffbfef8b7e72ec4bae02b6bbe862980b5fe575ae8432a000e9c4e4307caa2d9b.json b/crates/storage-pg/.sqlx/query-ffbfef8b7e72ec4bae02b6bbe862980b5fe575ae8432a000e9c4e4307caa2d9b.json
new file mode 100644
index 000000000..77f807125
--- /dev/null
+++ b/crates/storage-pg/.sqlx/query-ffbfef8b7e72ec4bae02b6bbe862980b5fe575ae8432a000e9c4e4307caa2d9b.json
@@ -0,0 +1,16 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "\n                UPDATE oauth2_refresh_tokens\n                SET consumed_at = $2,\n                    next_oauth2_refresh_token_id = $3\n                WHERE oauth2_refresh_token_id = $1\n            ",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Left": [
+        "Uuid",
+        "Timestamptz",
+        "Uuid"
+      ]
+    },
+    "nullable": []
+  },
+  "hash": "ffbfef8b7e72ec4bae02b6bbe862980b5fe575ae8432a000e9c4e4307caa2d9b"
+}
diff --git a/crates/storage-pg/migrations/20241210115428_oauth_refresh_token_track_next.sql b/crates/storage-pg/migrations/20241210115428_oauth_refresh_token_track_next.sql
new file mode 100644
index 000000000..b1bd2c3f6
--- /dev/null
+++ b/crates/storage-pg/migrations/20241210115428_oauth_refresh_token_track_next.sql
@@ -0,0 +1,9 @@
+-- Copyright 2024 New Vector Ltd.
+--
+-- SPDX-License-Identifier: AGPL-3.0-only
+-- Please see LICENSE in the repository root for full details.
+
+-- Add a reference to the 'next' refresh token when it was consumed and replaced
+ALTER TABLE oauth2_refresh_tokens
+  ADD COLUMN "next_oauth2_refresh_token_id" UUID
+    REFERENCES oauth2_refresh_tokens (oauth2_refresh_token_id);
diff --git a/crates/storage-pg/src/oauth2/mod.rs b/crates/storage-pg/src/oauth2/mod.rs
index c55b4d70e..54ce32fb7 100644
--- a/crates/storage-pg/src/oauth2/mod.rs
+++ b/crates/storage-pg/src/oauth2/mod.rs
@@ -341,6 +341,19 @@ mod tests {
         clock.advance(Duration::try_minutes(-6).unwrap()); // Go back in time
         assert!(access_token.is_valid(clock.now()));
 
+        // Create a new refresh token to be able to consume the old one
+        let new_refresh_token = repo
+            .oauth2_refresh_token()
+            .add(
+                &mut rng,
+                &clock,
+                &session,
+                &access_token,
+                "ddeeff".to_owned(),
+            )
+            .await
+            .unwrap();
+
         // Mark the access token as revoked
         let access_token = repo
             .oauth2_access_token()
@@ -353,7 +366,7 @@ mod tests {
         assert!(refresh_token.is_valid());
         let refresh_token = repo
             .oauth2_refresh_token()
-            .consume(&clock, refresh_token)
+            .consume(&clock, refresh_token, &new_refresh_token)
             .await
             .unwrap();
         assert!(!refresh_token.is_valid());
diff --git a/crates/storage-pg/src/oauth2/refresh_token.rs b/crates/storage-pg/src/oauth2/refresh_token.rs
index bc0480a67..1f0ddf0f4 100644
--- a/crates/storage-pg/src/oauth2/refresh_token.rs
+++ b/crates/storage-pg/src/oauth2/refresh_token.rs
@@ -13,7 +13,7 @@ use sqlx::PgConnection;
 use ulid::Ulid;
 use uuid::Uuid;
 
-use crate::{tracing::ExecuteExt, DatabaseError};
+use crate::{tracing::ExecuteExt, DatabaseError, DatabaseInconsistencyError};
 
 /// An implementation of [`OAuth2RefreshTokenRepository`] for a PostgreSQL
 /// connection
@@ -36,23 +36,39 @@ struct OAuth2RefreshTokenLookup {
     consumed_at: Option<DateTime<Utc>>,
     oauth2_access_token_id: Option<Uuid>,
     oauth2_session_id: Uuid,
+    next_oauth2_refresh_token_id: Option<Uuid>,
 }
 
-impl From<OAuth2RefreshTokenLookup> for RefreshToken {
-    fn from(value: OAuth2RefreshTokenLookup) -> Self {
-        let state = match value.consumed_at {
-            None => RefreshTokenState::Valid,
-            Some(consumed_at) => RefreshTokenState::Consumed { consumed_at },
+impl TryFrom<OAuth2RefreshTokenLookup> for RefreshToken {
+    type Error = DatabaseInconsistencyError;
+
+    fn try_from(value: OAuth2RefreshTokenLookup) -> Result<Self, Self::Error> {
+        let id = value.oauth2_refresh_token_id.into();
+        let state = match (value.consumed_at, value.next_oauth2_refresh_token_id) {
+            (None, None) => RefreshTokenState::Valid,
+            (Some(consumed_at), None) => RefreshTokenState::Consumed {
+                consumed_at,
+                next_refresh_token_id: None,
+            },
+            (Some(consumed_at), Some(id)) => RefreshTokenState::Consumed {
+                consumed_at,
+                next_refresh_token_id: Some(Ulid::from(id)),
+            },
+            (None, Some(_)) => {
+                return Err(DatabaseInconsistencyError::on("oauth2_refresh_tokens")
+                    .column("next_oauth2_refresh_token_id")
+                    .row(id))
+            }
         };
 
-        RefreshToken {
-            id: value.oauth2_refresh_token_id.into(),
+        Ok(RefreshToken {
+            id,
             state,
             session_id: value.oauth2_session_id.into(),
             refresh_token: value.refresh_token,
             created_at: value.created_at,
             access_token_id: value.oauth2_access_token_id.map(Ulid::from),
-        }
+        })
     }
 }
 
@@ -79,18 +95,20 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> {
                      , consumed_at
                      , oauth2_access_token_id
                      , oauth2_session_id
+                     , next_oauth2_refresh_token_id
                 FROM oauth2_refresh_tokens
 
                 WHERE oauth2_refresh_token_id = $1
             "#,
             Uuid::from(id),
         )
+        .traced()
         .fetch_optional(&mut *self.conn)
         .await?;
 
         let Some(res) = res else { return Ok(None) };
 
-        Ok(Some(res.into()))
+        Ok(Some(res.try_into()?))
     }
 
     #[tracing::instrument(
@@ -114,6 +132,7 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> {
                      , consumed_at
                      , oauth2_access_token_id
                      , oauth2_session_id
+                     , next_oauth2_refresh_token_id
                 FROM oauth2_refresh_tokens
 
                 WHERE refresh_token = $1
@@ -126,7 +145,7 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> {
 
         let Some(res) = res else { return Ok(None) };
 
-        Ok(Some(res.into()))
+        Ok(Some(res.try_into()?))
     }
 
     #[tracing::instrument(
@@ -194,24 +213,28 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> {
         &mut self,
         clock: &dyn Clock,
         refresh_token: RefreshToken,
+        replaced_by: &RefreshToken,
     ) -> Result<RefreshToken, Self::Error> {
         let consumed_at = clock.now();
         let res = sqlx::query!(
             r#"
                 UPDATE oauth2_refresh_tokens
-                SET consumed_at = $2
+                SET consumed_at = $2,
+                    next_oauth2_refresh_token_id = $3
                 WHERE oauth2_refresh_token_id = $1
             "#,
             Uuid::from(refresh_token.id),
             consumed_at,
+            Uuid::from(replaced_by.id),
         )
+        .traced()
         .execute(&mut *self.conn)
         .await?;
 
         DatabaseError::ensure_affected_rows(&res, 1)?;
 
         refresh_token
-            .consume(consumed_at)
+            .consume(consumed_at, replaced_by)
             .map_err(DatabaseError::to_invalid_operation)
     }
 }
diff --git a/crates/storage/src/oauth2/refresh_token.rs b/crates/storage/src/oauth2/refresh_token.rs
index c003e9270..5612b3647 100644
--- a/crates/storage/src/oauth2/refresh_token.rs
+++ b/crates/storage/src/oauth2/refresh_token.rs
@@ -80,6 +80,7 @@ pub trait OAuth2RefreshTokenRepository: Send + Sync {
     ///
     /// * `clock`: The clock used to generate timestamps
     /// * `refresh_token`: The [`RefreshToken`] to consume
+    /// * `replaced_by`: The [`RefreshToken`] which replaced this one
     ///
     /// # Errors
     ///
@@ -89,6 +90,7 @@ pub trait OAuth2RefreshTokenRepository: Send + Sync {
         &mut self,
         clock: &dyn Clock,
         refresh_token: RefreshToken,
+        replaced_by: &RefreshToken,
     ) -> Result<RefreshToken, Self::Error>;
 }
 
@@ -113,5 +115,6 @@ repository_impl!(OAuth2RefreshTokenRepository:
         &mut self,
         clock: &dyn Clock,
         refresh_token: RefreshToken,
+        replaced_by: &RefreshToken,
     ) -> Result<RefreshToken, Self::Error>;
 );

From 45f578017994fb906ff92eddc43176b6dcf130ed Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Tue, 10 Dec 2024 15:17:36 +0100
Subject: [PATCH 2/7] Record when access tokens are first used

---
 crates/data-model/src/tokens.rs               |  7 +++
 crates/handlers/src/oauth2/introspection.rs   | 22 ++++++-
 ...fb2254d1bf9e8a97cd7d32ba789c740e0fbdb.json | 15 +++++
 ...745cc1c57bbe3c3c69110592a07400116c7f.json} | 12 +++-
 ...fab6ea87589819128e6fc9ed5da11dfc2770.json} | 12 +++-
 ...0133651_oauth2_access_token_first_used.sql |  8 +++
 crates/storage-pg/src/oauth2/access_token.rs  | 60 +++++++++++++++++++
 crates/storage/src/oauth2/access_token.rs     | 22 +++++++
 8 files changed, 151 insertions(+), 7 deletions(-)
 create mode 100644 crates/storage-pg/.sqlx/query-7189b6136fd08ac9ae7c51bff06fb2254d1bf9e8a97cd7d32ba789c740e0fbdb.json
 rename crates/storage-pg/.sqlx/{query-dd16942318bf38d9a245b2c86fedd3cbd6b65e7a13465552d79cd3c022122fd4.json => query-875294dc5cf87bcf302fb9e87933745cc1c57bbe3c3c69110592a07400116c7f.json} (71%)
 rename crates/storage-pg/.sqlx/{query-477f79556e5777b38feb85013b4f04dbb8230e4b0b0bcc45f669d7b8d0b91db4.json => query-c09e0bb0378d9dfb15de7f2f1209fab6ea87589819128e6fc9ed5da11dfc2770.json} (72%)
 create mode 100644 crates/storage-pg/migrations/20241210133651_oauth2_access_token_first_used.sql

diff --git a/crates/data-model/src/tokens.rs b/crates/data-model/src/tokens.rs
index 210eae0cc..bb0b16abe 100644
--- a/crates/data-model/src/tokens.rs
+++ b/crates/data-model/src/tokens.rs
@@ -55,6 +55,7 @@ pub struct AccessToken {
     pub access_token: String,
     pub created_at: DateTime<Utc>,
     pub expires_at: Option<DateTime<Utc>>,
+    pub first_used_at: Option<DateTime<Utc>>,
 }
 
 impl AccessToken {
@@ -88,6 +89,12 @@ impl AccessToken {
         }
     }
 
+    /// Whether the access token was used at least once
+    #[must_use]
+    pub fn is_used(&self) -> bool {
+        self.first_used_at.is_some()
+    }
+
     /// Mark the access token as revoked
     ///
     /// # Parameters
diff --git a/crates/handlers/src/oauth2/introspection.rs b/crates/handlers/src/oauth2/introspection.rs
index 4c534f8d3..fcb0a205f 100644
--- a/crates/handlers/src/oauth2/introspection.rs
+++ b/crates/handlers/src/oauth2/introspection.rs
@@ -207,7 +207,7 @@ pub(crate) async fn post(
 
     let reply = match token_type {
         TokenType::AccessToken => {
-            let access_token = repo
+            let mut access_token = repo
                 .oauth2_access_token()
                 .find_by_token(token)
                 .await?
@@ -227,6 +227,14 @@ pub(crate) async fn post(
                 return Err(RouteError::InvalidOAuthSession);
             }
 
+            // If this is the first time we're using this token, mark it as used
+            if !access_token.is_used() {
+                access_token = repo
+                    .oauth2_access_token()
+                    .mark_used(&clock, access_token)
+                    .await?;
+            }
+
             // The session might not have a user on it (for Client Credentials grants for
             // example), so we're optionally fetching the user
             let (sub, username) = if let Some(user_id) = session.user_id {
@@ -443,6 +451,8 @@ pub(crate) async fn post(
         }
     };
 
+    repo.save().await?;
+
     Ok(Json(reply))
 }
 
@@ -625,6 +635,16 @@ mod tests {
             .unwrap()
             .unwrap();
         assert_eq!(session.last_active_at, Some(state.clock.now()));
+
+        // And recorded the access token as used
+        let access_token_lookup = repo
+            .oauth2_access_token()
+            .find_by_token(&access_token)
+            .await
+            .unwrap()
+            .unwrap();
+        assert!(access_token_lookup.is_used());
+        assert_eq!(access_token_lookup.first_used_at, Some(state.clock.now()));
         repo.cancel().await.unwrap();
 
         // Advance the clock to invalidate the access token
diff --git a/crates/storage-pg/.sqlx/query-7189b6136fd08ac9ae7c51bff06fb2254d1bf9e8a97cd7d32ba789c740e0fbdb.json b/crates/storage-pg/.sqlx/query-7189b6136fd08ac9ae7c51bff06fb2254d1bf9e8a97cd7d32ba789c740e0fbdb.json
new file mode 100644
index 000000000..a0695db46
--- /dev/null
+++ b/crates/storage-pg/.sqlx/query-7189b6136fd08ac9ae7c51bff06fb2254d1bf9e8a97cd7d32ba789c740e0fbdb.json
@@ -0,0 +1,15 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "\n                UPDATE oauth2_access_tokens\n                SET first_used_at = $2\n                WHERE oauth2_access_token_id = $1\n            ",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Left": [
+        "Uuid",
+        "Timestamptz"
+      ]
+    },
+    "nullable": []
+  },
+  "hash": "7189b6136fd08ac9ae7c51bff06fb2254d1bf9e8a97cd7d32ba789c740e0fbdb"
+}
diff --git a/crates/storage-pg/.sqlx/query-dd16942318bf38d9a245b2c86fedd3cbd6b65e7a13465552d79cd3c022122fd4.json b/crates/storage-pg/.sqlx/query-875294dc5cf87bcf302fb9e87933745cc1c57bbe3c3c69110592a07400116c7f.json
similarity index 71%
rename from crates/storage-pg/.sqlx/query-dd16942318bf38d9a245b2c86fedd3cbd6b65e7a13465552d79cd3c022122fd4.json
rename to crates/storage-pg/.sqlx/query-875294dc5cf87bcf302fb9e87933745cc1c57bbe3c3c69110592a07400116c7f.json
index c9bd70555..b82d7c463 100644
--- a/crates/storage-pg/.sqlx/query-dd16942318bf38d9a245b2c86fedd3cbd6b65e7a13465552d79cd3c022122fd4.json
+++ b/crates/storage-pg/.sqlx/query-875294dc5cf87bcf302fb9e87933745cc1c57bbe3c3c69110592a07400116c7f.json
@@ -1,6 +1,6 @@
 {
   "db_name": "PostgreSQL",
-  "query": "\n                SELECT oauth2_access_token_id\n                     , access_token\n                     , created_at\n                     , expires_at\n                     , revoked_at\n                     , oauth2_session_id\n\n                FROM oauth2_access_tokens\n\n                WHERE oauth2_access_token_id = $1\n            ",
+  "query": "\n                SELECT oauth2_access_token_id\n                     , access_token\n                     , created_at\n                     , expires_at\n                     , revoked_at\n                     , oauth2_session_id\n                     , first_used_at\n\n                FROM oauth2_access_tokens\n\n                WHERE oauth2_access_token_id = $1\n            ",
   "describe": {
     "columns": [
       {
@@ -32,6 +32,11 @@
         "ordinal": 5,
         "name": "oauth2_session_id",
         "type_info": "Uuid"
+      },
+      {
+        "ordinal": 6,
+        "name": "first_used_at",
+        "type_info": "Timestamptz"
       }
     ],
     "parameters": {
@@ -45,8 +50,9 @@
       false,
       true,
       true,
-      false
+      false,
+      true
     ]
   },
-  "hash": "dd16942318bf38d9a245b2c86fedd3cbd6b65e7a13465552d79cd3c022122fd4"
+  "hash": "875294dc5cf87bcf302fb9e87933745cc1c57bbe3c3c69110592a07400116c7f"
 }
diff --git a/crates/storage-pg/.sqlx/query-477f79556e5777b38feb85013b4f04dbb8230e4b0b0bcc45f669d7b8d0b91db4.json b/crates/storage-pg/.sqlx/query-c09e0bb0378d9dfb15de7f2f1209fab6ea87589819128e6fc9ed5da11dfc2770.json
similarity index 72%
rename from crates/storage-pg/.sqlx/query-477f79556e5777b38feb85013b4f04dbb8230e4b0b0bcc45f669d7b8d0b91db4.json
rename to crates/storage-pg/.sqlx/query-c09e0bb0378d9dfb15de7f2f1209fab6ea87589819128e6fc9ed5da11dfc2770.json
index bc6c2fb72..2ef7b86cf 100644
--- a/crates/storage-pg/.sqlx/query-477f79556e5777b38feb85013b4f04dbb8230e4b0b0bcc45f669d7b8d0b91db4.json
+++ b/crates/storage-pg/.sqlx/query-c09e0bb0378d9dfb15de7f2f1209fab6ea87589819128e6fc9ed5da11dfc2770.json
@@ -1,6 +1,6 @@
 {
   "db_name": "PostgreSQL",
-  "query": "\n                SELECT oauth2_access_token_id\n                     , access_token\n                     , created_at\n                     , expires_at\n                     , revoked_at\n                     , oauth2_session_id\n\n                FROM oauth2_access_tokens\n\n                WHERE access_token = $1\n            ",
+  "query": "\n                SELECT oauth2_access_token_id\n                     , access_token\n                     , created_at\n                     , expires_at\n                     , revoked_at\n                     , oauth2_session_id\n                     , first_used_at\n\n                FROM oauth2_access_tokens\n\n                WHERE access_token = $1\n            ",
   "describe": {
     "columns": [
       {
@@ -32,6 +32,11 @@
         "ordinal": 5,
         "name": "oauth2_session_id",
         "type_info": "Uuid"
+      },
+      {
+        "ordinal": 6,
+        "name": "first_used_at",
+        "type_info": "Timestamptz"
       }
     ],
     "parameters": {
@@ -45,8 +50,9 @@
       false,
       true,
       true,
-      false
+      false,
+      true
     ]
   },
-  "hash": "477f79556e5777b38feb85013b4f04dbb8230e4b0b0bcc45f669d7b8d0b91db4"
+  "hash": "c09e0bb0378d9dfb15de7f2f1209fab6ea87589819128e6fc9ed5da11dfc2770"
 }
diff --git a/crates/storage-pg/migrations/20241210133651_oauth2_access_token_first_used.sql b/crates/storage-pg/migrations/20241210133651_oauth2_access_token_first_used.sql
new file mode 100644
index 000000000..8d9507a69
--- /dev/null
+++ b/crates/storage-pg/migrations/20241210133651_oauth2_access_token_first_used.sql
@@ -0,0 +1,8 @@
+-- Copyright 2024 New Vector Ltd.
+--
+-- SPDX-License-Identifier: AGPL-3.0-only
+-- Please see LICENSE in the repository root for full details.
+
+-- Track when the access token was first used. A NULL value means it was never used.
+ALTER TABLE oauth2_access_tokens
+  ADD COLUMN "first_used_at" TIMESTAMP WITH TIME ZONE;
diff --git a/crates/storage-pg/src/oauth2/access_token.rs b/crates/storage-pg/src/oauth2/access_token.rs
index ea8be74d3..6fdcf41ef 100644
--- a/crates/storage-pg/src/oauth2/access_token.rs
+++ b/crates/storage-pg/src/oauth2/access_token.rs
@@ -36,6 +36,7 @@ struct OAuth2AccessTokenLookup {
     created_at: DateTime<Utc>,
     expires_at: Option<DateTime<Utc>>,
     revoked_at: Option<DateTime<Utc>>,
+    first_used_at: Option<DateTime<Utc>>,
 }
 
 impl From<OAuth2AccessTokenLookup> for AccessToken {
@@ -52,6 +53,7 @@ impl From<OAuth2AccessTokenLookup> for AccessToken {
             access_token: value.access_token,
             created_at: value.created_at,
             expires_at: value.expires_at,
+            first_used_at: value.first_used_at,
         }
     }
 }
@@ -70,6 +72,7 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> {
                      , expires_at
                      , revoked_at
                      , oauth2_session_id
+                     , first_used_at
 
                 FROM oauth2_access_tokens
 
@@ -106,6 +109,7 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> {
                      , expires_at
                      , revoked_at
                      , oauth2_session_id
+                     , first_used_at
 
                 FROM oauth2_access_tokens
 
@@ -170,9 +174,20 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> {
             session_id: session.id,
             created_at,
             expires_at,
+            first_used_at: None,
         })
     }
 
+    #[tracing::instrument(
+        name = "db.oauth2_access_token.revoked",
+        skip_all,
+        fields(
+            db.query.text,
+            session.id = %access_token.session_id,
+            %access_token.id,
+        ),
+        err,
+    )]
     async fn revoke(
         &mut self,
         clock: &dyn Clock,
@@ -188,6 +203,7 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> {
             Uuid::from(access_token.id),
             revoked_at,
         )
+        .traced()
         .execute(&mut *self.conn)
         .await?;
 
@@ -198,6 +214,49 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> {
             .map_err(DatabaseError::to_invalid_operation)
     }
 
+    #[tracing::instrument(
+        name = "db.oauth2_access_token.mark_used",
+        skip_all,
+        fields(
+            db.query.text,
+            session.id = %access_token.session_id,
+            %access_token.id,
+        ),
+        err,
+    )]
+    async fn mark_used(
+        &mut self,
+        clock: &dyn Clock,
+        mut access_token: AccessToken,
+    ) -> Result<AccessToken, Self::Error> {
+        let now = clock.now();
+        let res = sqlx::query!(
+            r#"
+                UPDATE oauth2_access_tokens
+                SET first_used_at = $2
+                WHERE oauth2_access_token_id = $1
+            "#,
+            Uuid::from(access_token.id),
+            now,
+        )
+        .execute(&mut *self.conn)
+        .await?;
+
+        DatabaseError::ensure_affected_rows(&res, 1)?;
+
+        access_token.first_used_at = Some(now);
+
+        Ok(access_token)
+    }
+
+    #[tracing::instrument(
+        name = "db.oauth2_access_token.cleanup_expired",
+        skip_all,
+        fields(
+            db.query.text,
+        ),
+        err,
+    )]
     async fn cleanup_expired(&mut self, clock: &dyn Clock) -> Result<usize, Self::Error> {
         // Cleanup token which expired more than 15 minutes ago
         let threshold = clock.now() - Duration::microseconds(15 * 60 * 1000 * 1000);
@@ -208,6 +267,7 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> {
             "#,
             threshold,
         )
+        .traced()
         .execute(&mut *self.conn)
         .await?;
 
diff --git a/crates/storage/src/oauth2/access_token.rs b/crates/storage/src/oauth2/access_token.rs
index 79354ef6b..83eefee54 100644
--- a/crates/storage/src/oauth2/access_token.rs
+++ b/crates/storage/src/oauth2/access_token.rs
@@ -91,6 +91,22 @@ pub trait OAuth2AccessTokenRepository: Send + Sync {
         access_token: AccessToken,
     ) -> Result<AccessToken, Self::Error>;
 
+    /// Mark the access token as used, to track when it was first used
+    ///
+    /// # Parameters
+    ///
+    /// * `clock`: The clock used to generate timestamps
+    /// * `access_token`: The access token to mark as used
+    ///
+    /// # Errors
+    ///
+    /// Returns [`Self::Error`] if the underlying repository fails
+    async fn mark_used(
+        &mut self,
+        clock: &dyn Clock,
+        access_token: AccessToken,
+    ) -> Result<AccessToken, Self::Error>;
+
     /// Cleanup expired access tokens
     ///
     /// Returns the number of access tokens that were cleaned up
@@ -128,5 +144,11 @@ repository_impl!(OAuth2AccessTokenRepository:
         access_token: AccessToken,
     ) -> Result<AccessToken, Self::Error>;
 
+    async fn mark_used(
+        &mut self,
+        clock: &dyn Clock,
+        access_token: AccessToken,
+    ) -> Result<AccessToken, Self::Error>;
+
     async fn cleanup_expired(&mut self, clock: &dyn Clock) -> Result<usize, Self::Error>;
 );

From 44c27e895731432033e054704dd7988e96cec933 Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Tue, 10 Dec 2024 16:15:16 +0100
Subject: [PATCH 3/7] Cleanup revoked tokens instead of expired ones

If we continue deleting expired tokens, we might not record whether the
token was used or not, and not know what to do in case of
a double-refresh.

Revoked tokens are safe to delete.

This also reduces the frequency of the cleanup job to once an hour.
---
 ...09403caafe1cc5267780232cd615468d50f5c0af2bedb.json} |  4 ++--
 crates/storage-pg/src/oauth2/access_token.rs           | 10 +++++-----
 crates/storage/src/oauth2/access_token.rs              |  6 +++---
 crates/tasks/src/database.rs                           |  4 ++--
 crates/tasks/src/lib.rs                                |  2 +-
 5 files changed, 13 insertions(+), 13 deletions(-)
 rename crates/storage-pg/.sqlx/{query-1a8701f5672de052bb766933f60b93249acc7237b996e8b93cd61b9f69c902ff.json => query-c68b28232b42a62907709403caafe1cc5267780232cd615468d50f5c0af2bedb.json} (63%)

diff --git a/crates/storage-pg/.sqlx/query-1a8701f5672de052bb766933f60b93249acc7237b996e8b93cd61b9f69c902ff.json b/crates/storage-pg/.sqlx/query-c68b28232b42a62907709403caafe1cc5267780232cd615468d50f5c0af2bedb.json
similarity index 63%
rename from crates/storage-pg/.sqlx/query-1a8701f5672de052bb766933f60b93249acc7237b996e8b93cd61b9f69c902ff.json
rename to crates/storage-pg/.sqlx/query-c68b28232b42a62907709403caafe1cc5267780232cd615468d50f5c0af2bedb.json
index 07ab1afa8..5da8bf5bc 100644
--- a/crates/storage-pg/.sqlx/query-1a8701f5672de052bb766933f60b93249acc7237b996e8b93cd61b9f69c902ff.json
+++ b/crates/storage-pg/.sqlx/query-c68b28232b42a62907709403caafe1cc5267780232cd615468d50f5c0af2bedb.json
@@ -1,6 +1,6 @@
 {
   "db_name": "PostgreSQL",
-  "query": "\n                DELETE FROM oauth2_access_tokens\n                WHERE expires_at < $1\n            ",
+  "query": "\n                DELETE FROM oauth2_access_tokens\n                WHERE revoked_at < $1\n            ",
   "describe": {
     "columns": [],
     "parameters": {
@@ -10,5 +10,5 @@
     },
     "nullable": []
   },
-  "hash": "1a8701f5672de052bb766933f60b93249acc7237b996e8b93cd61b9f69c902ff"
+  "hash": "c68b28232b42a62907709403caafe1cc5267780232cd615468d50f5c0af2bedb"
 }
diff --git a/crates/storage-pg/src/oauth2/access_token.rs b/crates/storage-pg/src/oauth2/access_token.rs
index 6fdcf41ef..0320761d0 100644
--- a/crates/storage-pg/src/oauth2/access_token.rs
+++ b/crates/storage-pg/src/oauth2/access_token.rs
@@ -250,20 +250,20 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> {
     }
 
     #[tracing::instrument(
-        name = "db.oauth2_access_token.cleanup_expired",
+        name = "db.oauth2_access_token.cleanup_revoked",
         skip_all,
         fields(
             db.query.text,
         ),
         err,
     )]
-    async fn cleanup_expired(&mut self, clock: &dyn Clock) -> Result<usize, Self::Error> {
-        // Cleanup token which expired more than 15 minutes ago
-        let threshold = clock.now() - Duration::microseconds(15 * 60 * 1000 * 1000);
+    async fn cleanup_revoked(&mut self, clock: &dyn Clock) -> Result<usize, Self::Error> {
+        // Cleanup token that were revoked more than an hour ago
+        let threshold = clock.now() - Duration::microseconds(60 * 60 * 1000 * 1000);
         let res = sqlx::query!(
             r#"
                 DELETE FROM oauth2_access_tokens
-                WHERE expires_at < $1
+                WHERE revoked_at < $1
             "#,
             threshold,
         )
diff --git a/crates/storage/src/oauth2/access_token.rs b/crates/storage/src/oauth2/access_token.rs
index 83eefee54..989dd9861 100644
--- a/crates/storage/src/oauth2/access_token.rs
+++ b/crates/storage/src/oauth2/access_token.rs
@@ -107,7 +107,7 @@ pub trait OAuth2AccessTokenRepository: Send + Sync {
         access_token: AccessToken,
     ) -> Result<AccessToken, Self::Error>;
 
-    /// Cleanup expired access tokens
+    /// Cleanup revoked access tokens
     ///
     /// Returns the number of access tokens that were cleaned up
     ///
@@ -118,7 +118,7 @@ pub trait OAuth2AccessTokenRepository: Send + Sync {
     /// # Errors
     ///
     /// Returns [`Self::Error`] if the underlying repository fails
-    async fn cleanup_expired(&mut self, clock: &dyn Clock) -> Result<usize, Self::Error>;
+    async fn cleanup_revoked(&mut self, clock: &dyn Clock) -> Result<usize, Self::Error>;
 }
 
 repository_impl!(OAuth2AccessTokenRepository:
@@ -150,5 +150,5 @@ repository_impl!(OAuth2AccessTokenRepository:
         access_token: AccessToken,
     ) -> Result<AccessToken, Self::Error>;
 
-    async fn cleanup_expired(&mut self, clock: &dyn Clock) -> Result<usize, Self::Error>;
+    async fn cleanup_revoked(&mut self, clock: &dyn Clock) -> Result<usize, Self::Error>;
 );
diff --git a/crates/tasks/src/database.rs b/crates/tasks/src/database.rs
index b64e1defb..74903fcd1 100644
--- a/crates/tasks/src/database.rs
+++ b/crates/tasks/src/database.rs
@@ -24,7 +24,7 @@ impl RunnableJob for CleanupExpiredTokensJob {
 
         let count = repo
             .oauth2_access_token()
-            .cleanup_expired(&clock)
+            .cleanup_revoked(&clock)
             .await
             .map_err(JobError::retry)?;
         repo.save().await.map_err(JobError::retry)?;
@@ -32,7 +32,7 @@ impl RunnableJob for CleanupExpiredTokensJob {
         if count == 0 {
             debug!("no token to clean up");
         } else {
-            info!(count, "cleaned up expired tokens");
+            info!(count, "cleaned up revoked tokens");
         }
 
         Ok(())
diff --git a/crates/tasks/src/lib.rs b/crates/tasks/src/lib.rs
index ecfd1ca18..eb9f19306 100644
--- a/crates/tasks/src/lib.rs
+++ b/crates/tasks/src/lib.rs
@@ -121,7 +121,7 @@ pub async fn init(
         .register_handler::<mas_storage::queue::VerifyEmailJob>()
         .add_schedule(
             "cleanup-expired-tokens",
-            "*/15 * * * * *".parse()?,
+            "0 0 * * * *".parse()?,
             mas_storage::queue::CleanupExpiredTokensJob,
         );
 

From 1f555b896415112d69dd93cbc65e1104baea7e63 Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Tue, 10 Dec 2024 17:44:07 +0100
Subject: [PATCH 4/7] Allow revoking refresh tokens

This lets us track 'revoked' tokens separately from 'consumed' tokens.
---
 crates/data-model/src/tokens.rs               | 43 ++++++++++-----
 ...7a709b06455e08b9fd75cc08f142070f330b3.json | 15 ++++++
 ...6d3b08fd5d736ecf36845e6fd4bfc515b2cf.json} | 14 +++--
 ...009785f358b334f5c586c2e358f0d0b4d856.json} | 14 +++--
 crates/storage-pg/src/oauth2/access_token.rs  |  2 +-
 crates/storage-pg/src/oauth2/refresh_token.rs | 54 +++++++++++++++++--
 crates/storage/src/oauth2/refresh_token.rs    | 27 +++++++++-
 7 files changed, 141 insertions(+), 28 deletions(-)
 create mode 100644 crates/storage-pg/.sqlx/query-66693f31eff5673e88ca516ee727a709b06455e08b9fd75cc08f142070f330b3.json
 rename crates/storage-pg/.sqlx/{query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json => query-6d71188dffc492ddc8f7f21476516d3b08fd5d736ecf36845e6fd4bfc515b2cf.json} (72%)
 rename crates/storage-pg/.sqlx/{query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json => query-a75a6a08c9639053cfc3cffa9d4a009785f358b334f5c586c2e358f0d0b4d856.json} (72%)

diff --git a/crates/data-model/src/tokens.rs b/crates/data-model/src/tokens.rs
index bb0b16abe..9e94f7171 100644
--- a/crates/data-model/src/tokens.rs
+++ b/crates/data-model/src/tokens.rs
@@ -118,6 +118,9 @@ pub enum RefreshTokenState {
         consumed_at: DateTime<Utc>,
         next_refresh_token_id: Option<Ulid>,
     },
+    Revoked {
+        revoked_at: DateTime<Utc>,
+    },
 }
 
 impl RefreshTokenState {
@@ -125,18 +128,30 @@ impl RefreshTokenState {
     ///
     /// # Errors
     ///
-    /// Returns an error if the refresh token is already consumed.
+    /// Returns an error if the refresh token is revoked.
     fn consume(
         self,
         consumed_at: DateTime<Utc>,
         replaced_by: &RefreshToken,
     ) -> Result<Self, InvalidTransitionError> {
         match self {
-            Self::Valid => Ok(Self::Consumed {
+            Self::Valid | Self::Consumed { .. } => Ok(Self::Consumed {
                 consumed_at,
                 next_refresh_token_id: Some(replaced_by.id),
             }),
-            Self::Consumed { .. } => Err(InvalidTransitionError),
+            Self::Revoked { .. } => Err(InvalidTransitionError),
+        }
+    }
+
+    /// Revoke the refresh token, returning a new state.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the refresh token is already consumed or revoked.
+    pub fn revoke(self, revoked_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+        match self {
+            Self::Valid => Ok(Self::Revoked { revoked_at }),
+            Self::Consumed { .. } | Self::Revoked { .. } => Err(InvalidTransitionError),
         }
     }
 
@@ -148,19 +163,11 @@ impl RefreshTokenState {
         matches!(self, Self::Valid)
     }
 
-    /// Returns `true` if the refresh token state is [`Consumed`].
-    ///
-    /// [`Consumed`]: RefreshTokenState::Consumed
-    #[must_use]
-    pub fn is_consumed(&self) -> bool {
-        matches!(self, Self::Consumed { .. })
-    }
-
     /// Returns the next refresh token ID, if any.
     #[must_use]
     pub fn next_refresh_token_id(&self) -> Option<Ulid> {
         match self {
-            Self::Valid => None,
+            Self::Valid | Self::Revoked { .. } => None,
             Self::Consumed {
                 next_refresh_token_id,
                 ..
@@ -197,7 +204,7 @@ impl RefreshToken {
     ///
     /// # Errors
     ///
-    /// Returns an error if the refresh token is already consumed.
+    /// Returns an error if the refresh token is revoked.
     pub fn consume(
         mut self,
         consumed_at: DateTime<Utc>,
@@ -206,6 +213,16 @@ impl RefreshToken {
         self.state = self.state.consume(consumed_at, replaced_by)?;
         Ok(self)
     }
+
+    /// Revokes the refresh token and returns a new revoked token
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the refresh token is already revoked.
+    pub fn revoke(mut self, revoked_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+        self.state = self.state.revoke(revoked_at)?;
+        Ok(self)
+    }
 }
 
 /// Type of token to generate or validate
diff --git a/crates/storage-pg/.sqlx/query-66693f31eff5673e88ca516ee727a709b06455e08b9fd75cc08f142070f330b3.json b/crates/storage-pg/.sqlx/query-66693f31eff5673e88ca516ee727a709b06455e08b9fd75cc08f142070f330b3.json
new file mode 100644
index 000000000..716bbf8c7
--- /dev/null
+++ b/crates/storage-pg/.sqlx/query-66693f31eff5673e88ca516ee727a709b06455e08b9fd75cc08f142070f330b3.json
@@ -0,0 +1,15 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "\n                UPDATE oauth2_refresh_tokens\n                SET revoked_at = $2\n                WHERE oauth2_refresh_token_id = $1\n            ",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Left": [
+        "Uuid",
+        "Timestamptz"
+      ]
+    },
+    "nullable": []
+  },
+  "hash": "66693f31eff5673e88ca516ee727a709b06455e08b9fd75cc08f142070f330b3"
+}
diff --git a/crates/storage-pg/.sqlx/query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json b/crates/storage-pg/.sqlx/query-6d71188dffc492ddc8f7f21476516d3b08fd5d736ecf36845e6fd4bfc515b2cf.json
similarity index 72%
rename from crates/storage-pg/.sqlx/query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json
rename to crates/storage-pg/.sqlx/query-6d71188dffc492ddc8f7f21476516d3b08fd5d736ecf36845e6fd4bfc515b2cf.json
index be9fa3e32..1acfa2784 100644
--- a/crates/storage-pg/.sqlx/query-5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466.json
+++ b/crates/storage-pg/.sqlx/query-6d71188dffc492ddc8f7f21476516d3b08fd5d736ecf36845e6fd4bfc515b2cf.json
@@ -1,6 +1,6 @@
 {
   "db_name": "PostgreSQL",
-  "query": "\n                SELECT oauth2_refresh_token_id\n                     , refresh_token\n                     , created_at\n                     , consumed_at\n                     , oauth2_access_token_id\n                     , oauth2_session_id\n                     , next_oauth2_refresh_token_id\n                FROM oauth2_refresh_tokens\n\n                WHERE oauth2_refresh_token_id = $1\n            ",
+  "query": "\n                SELECT oauth2_refresh_token_id\n                     , refresh_token\n                     , created_at\n                     , consumed_at\n                     , revoked_at\n                     , oauth2_access_token_id\n                     , oauth2_session_id\n                     , next_oauth2_refresh_token_id\n                FROM oauth2_refresh_tokens\n\n                WHERE oauth2_refresh_token_id = $1\n            ",
   "describe": {
     "columns": [
       {
@@ -25,16 +25,21 @@
       },
       {
         "ordinal": 4,
+        "name": "revoked_at",
+        "type_info": "Timestamptz"
+      },
+      {
+        "ordinal": 5,
         "name": "oauth2_access_token_id",
         "type_info": "Uuid"
       },
       {
-        "ordinal": 5,
+        "ordinal": 6,
         "name": "oauth2_session_id",
         "type_info": "Uuid"
       },
       {
-        "ordinal": 6,
+        "ordinal": 7,
         "name": "next_oauth2_refresh_token_id",
         "type_info": "Uuid"
       }
@@ -50,9 +55,10 @@
       false,
       true,
       true,
+      true,
       false,
       true
     ]
   },
-  "hash": "5b94c22d44692a16fa5a6edd5dac019c36bf5983182b0871de6e85036d8df466"
+  "hash": "6d71188dffc492ddc8f7f21476516d3b08fd5d736ecf36845e6fd4bfc515b2cf"
 }
diff --git a/crates/storage-pg/.sqlx/query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json b/crates/storage-pg/.sqlx/query-a75a6a08c9639053cfc3cffa9d4a009785f358b334f5c586c2e358f0d0b4d856.json
similarity index 72%
rename from crates/storage-pg/.sqlx/query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json
rename to crates/storage-pg/.sqlx/query-a75a6a08c9639053cfc3cffa9d4a009785f358b334f5c586c2e358f0d0b4d856.json
index df5d58132..c7146b7d1 100644
--- a/crates/storage-pg/.sqlx/query-265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630.json
+++ b/crates/storage-pg/.sqlx/query-a75a6a08c9639053cfc3cffa9d4a009785f358b334f5c586c2e358f0d0b4d856.json
@@ -1,6 +1,6 @@
 {
   "db_name": "PostgreSQL",
-  "query": "\n                SELECT oauth2_refresh_token_id\n                     , refresh_token\n                     , created_at\n                     , consumed_at\n                     , oauth2_access_token_id\n                     , oauth2_session_id\n                     , next_oauth2_refresh_token_id\n                FROM oauth2_refresh_tokens\n\n                WHERE refresh_token = $1\n            ",
+  "query": "\n                SELECT oauth2_refresh_token_id\n                     , refresh_token\n                     , created_at\n                     , consumed_at\n                     , revoked_at\n                     , oauth2_access_token_id\n                     , oauth2_session_id\n                     , next_oauth2_refresh_token_id\n                FROM oauth2_refresh_tokens\n\n                WHERE refresh_token = $1\n            ",
   "describe": {
     "columns": [
       {
@@ -25,16 +25,21 @@
       },
       {
         "ordinal": 4,
+        "name": "revoked_at",
+        "type_info": "Timestamptz"
+      },
+      {
+        "ordinal": 5,
         "name": "oauth2_access_token_id",
         "type_info": "Uuid"
       },
       {
-        "ordinal": 5,
+        "ordinal": 6,
         "name": "oauth2_session_id",
         "type_info": "Uuid"
       },
       {
-        "ordinal": 6,
+        "ordinal": 7,
         "name": "next_oauth2_refresh_token_id",
         "type_info": "Uuid"
       }
@@ -50,9 +55,10 @@
       false,
       true,
       true,
+      true,
       false,
       true
     ]
   },
-  "hash": "265a981142194216593cd09cd8f6af36c103e030358262e46a4bd5e4006dc630"
+  "hash": "a75a6a08c9639053cfc3cffa9d4a009785f358b334f5c586c2e358f0d0b4d856"
 }
diff --git a/crates/storage-pg/src/oauth2/access_token.rs b/crates/storage-pg/src/oauth2/access_token.rs
index 0320761d0..48801ead6 100644
--- a/crates/storage-pg/src/oauth2/access_token.rs
+++ b/crates/storage-pg/src/oauth2/access_token.rs
@@ -179,7 +179,7 @@ impl OAuth2AccessTokenRepository for PgOAuth2AccessTokenRepository<'_> {
     }
 
     #[tracing::instrument(
-        name = "db.oauth2_access_token.revoked",
+        name = "db.oauth2_access_token.revoke",
         skip_all,
         fields(
             db.query.text,
diff --git a/crates/storage-pg/src/oauth2/refresh_token.rs b/crates/storage-pg/src/oauth2/refresh_token.rs
index 1f0ddf0f4..ee4412ffd 100644
--- a/crates/storage-pg/src/oauth2/refresh_token.rs
+++ b/crates/storage-pg/src/oauth2/refresh_token.rs
@@ -34,6 +34,7 @@ struct OAuth2RefreshTokenLookup {
     refresh_token: String,
     created_at: DateTime<Utc>,
     consumed_at: Option<DateTime<Utc>>,
+    revoked_at: Option<DateTime<Utc>>,
     oauth2_access_token_id: Option<Uuid>,
     oauth2_session_id: Uuid,
     next_oauth2_refresh_token_id: Option<Uuid>,
@@ -44,17 +45,22 @@ impl TryFrom<OAuth2RefreshTokenLookup> for RefreshToken {
 
     fn try_from(value: OAuth2RefreshTokenLookup) -> Result<Self, Self::Error> {
         let id = value.oauth2_refresh_token_id.into();
-        let state = match (value.consumed_at, value.next_oauth2_refresh_token_id) {
-            (None, None) => RefreshTokenState::Valid,
-            (Some(consumed_at), None) => RefreshTokenState::Consumed {
+        let state = match (
+            value.revoked_at,
+            value.consumed_at,
+            value.next_oauth2_refresh_token_id,
+        ) {
+            (None, None, None) => RefreshTokenState::Valid,
+            (Some(revoked_at), None, None) => RefreshTokenState::Revoked { revoked_at },
+            (None, Some(consumed_at), None) => RefreshTokenState::Consumed {
                 consumed_at,
                 next_refresh_token_id: None,
             },
-            (Some(consumed_at), Some(id)) => RefreshTokenState::Consumed {
+            (None, Some(consumed_at), Some(id)) => RefreshTokenState::Consumed {
                 consumed_at,
                 next_refresh_token_id: Some(Ulid::from(id)),
             },
-            (None, Some(_)) => {
+            _ => {
                 return Err(DatabaseInconsistencyError::on("oauth2_refresh_tokens")
                     .column("next_oauth2_refresh_token_id")
                     .row(id))
@@ -93,6 +99,7 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> {
                      , refresh_token
                      , created_at
                      , consumed_at
+                     , revoked_at
                      , oauth2_access_token_id
                      , oauth2_session_id
                      , next_oauth2_refresh_token_id
@@ -130,6 +137,7 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> {
                      , refresh_token
                      , created_at
                      , consumed_at
+                     , revoked_at
                      , oauth2_access_token_id
                      , oauth2_session_id
                      , next_oauth2_refresh_token_id
@@ -237,4 +245,40 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> {
             .consume(consumed_at, replaced_by)
             .map_err(DatabaseError::to_invalid_operation)
     }
+
+    #[tracing::instrument(
+        name = "db.oauth2_refresh_token.revoke",
+        skip_all,
+        fields(
+            db.query.text,
+            %refresh_token.id,
+            session.id = %refresh_token.session_id,
+        ),
+        err,
+    )]
+    async fn revoke(
+        &mut self,
+        clock: &dyn Clock,
+        refresh_token: RefreshToken,
+    ) -> Result<RefreshToken, Self::Error> {
+        let revoked_at = clock.now();
+        let res = sqlx::query!(
+            r#"
+                UPDATE oauth2_refresh_tokens
+                SET revoked_at = $2
+                WHERE oauth2_refresh_token_id = $1
+            "#,
+            Uuid::from(refresh_token.id),
+            revoked_at,
+        )
+        .traced()
+        .execute(&mut *self.conn)
+        .await?;
+
+        DatabaseError::ensure_affected_rows(&res, 1)?;
+
+        refresh_token
+            .revoke(revoked_at)
+            .map_err(DatabaseError::to_invalid_operation)
+    }
 }
diff --git a/crates/storage/src/oauth2/refresh_token.rs b/crates/storage/src/oauth2/refresh_token.rs
index 5612b3647..6e240f7b9 100644
--- a/crates/storage/src/oauth2/refresh_token.rs
+++ b/crates/storage/src/oauth2/refresh_token.rs
@@ -85,13 +85,32 @@ pub trait OAuth2RefreshTokenRepository: Send + Sync {
     /// # Errors
     ///
     /// Returns [`Self::Error`] if the underlying repository fails, or if the
-    /// token was already consumed
+    /// token was already consumed or revoked
     async fn consume(
         &mut self,
         clock: &dyn Clock,
         refresh_token: RefreshToken,
         replaced_by: &RefreshToken,
     ) -> Result<RefreshToken, Self::Error>;
+
+    /// Revoke a refresh token
+    ///
+    /// Returns the updated [`RefreshToken`]
+    ///
+    /// # Parameters
+    ///
+    /// * `clock`: The clock used to generate timestamps
+    /// * `refresh_token`: The [`RefreshToken`] to revoke
+    ///
+    /// # Errors
+    ///
+    /// Returns [`Self::Error`] if the underlying repository fails, or if the
+    /// token was already revoked or consumed
+    async fn revoke(
+        &mut self,
+        clock: &dyn Clock,
+        refresh_token: RefreshToken,
+    ) -> Result<RefreshToken, Self::Error>;
 }
 
 repository_impl!(OAuth2RefreshTokenRepository:
@@ -117,4 +136,10 @@ repository_impl!(OAuth2RefreshTokenRepository:
         refresh_token: RefreshToken,
         replaced_by: &RefreshToken,
     ) -> Result<RefreshToken, Self::Error>;
+
+    async fn revoke(
+        &mut self,
+        clock: &dyn Clock,
+        refresh_token: RefreshToken,
+    ) -> Result<RefreshToken, Self::Error>;
 );

From e6ff9077b8b9c4d9fe1c10a7d494ec29fbe0f057 Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Tue, 10 Dec 2024 17:45:52 +0100
Subject: [PATCH 5/7] Mark access token as used when calling the userinfo
 endpoint

---
 crates/axum-utils/src/user_authorization.rs | 5 +++++
 crates/handlers/src/oauth2/userinfo.rs      | 2 ++
 2 files changed, 7 insertions(+)

diff --git a/crates/axum-utils/src/user_authorization.rs b/crates/axum-utils/src/user_authorization.rs
index 66ec7cd3d..eb71ad5a4 100644
--- a/crates/axum-utils/src/user_authorization.rs
+++ b/crates/axum-utils/src/user_authorization.rs
@@ -117,6 +117,11 @@ impl<F: Send> UserAuthorization<F> {
             return Err(AuthorizationVerificationError::InvalidToken);
         }
 
+        if !token.is_used() {
+            // Mark the token as used
+            repo.oauth2_access_token().mark_used(clock, token).await?;
+        }
+
         Ok(session)
     }
 }
diff --git a/crates/handlers/src/oauth2/userinfo.rs b/crates/handlers/src/oauth2/userinfo.rs
index 6618f063b..c3223482b 100644
--- a/crates/handlers/src/oauth2/userinfo.rs
+++ b/crates/handlers/src/oauth2/userinfo.rs
@@ -142,6 +142,8 @@ pub async fn get(
         .await?
         .ok_or(RouteError::NoSuchClient)?;
 
+    repo.save().await?;
+
     if let Some(alg) = client.userinfo_signed_response_alg {
         let key = key_store
             .signing_key_for_algorithm(&alg)

From 76137c4193b94d0f5b8da664c9cc6620dbe9a0e7 Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Tue, 10 Dec 2024 17:46:12 +0100
Subject: [PATCH 6/7] Make sure the refresh token is idempotent

This allows using a refresh token multiple times, as long as the new
pair of tokens were not used in the meantime.
---
 crates/handlers/src/oauth2/token.rs | 275 +++++++++++++++++++++++++++-
 1 file changed, 266 insertions(+), 9 deletions(-)

diff --git a/crates/handlers/src/oauth2/token.rs b/crates/handlers/src/oauth2/token.rs
index 1222b32e5..9e8e201a0 100644
--- a/crates/handlers/src/oauth2/token.rs
+++ b/crates/handlers/src/oauth2/token.rs
@@ -39,7 +39,7 @@ use oauth2_types::{
     scope,
 };
 use thiserror::Error;
-use tracing::debug;
+use tracing::{debug, info};
 use ulid::Ulid;
 
 use super::{generate_id_token, generate_token_pair};
@@ -98,6 +98,20 @@ pub(crate) enum RouteError {
     #[error("failed to load oauth session")]
     NoSuchOAuthSession,
 
+    #[error(
+        "failed to load the next refresh token ({next:?}) from the previous one ({previous:?})"
+    )]
+    NoSuchNextRefreshToken { next: Ulid, previous: Ulid },
+
+    #[error("failed to load the access token ({access_token:?}) associated with the next refresh token ({refresh_token:?})")]
+    NoSuchNextAccessToken {
+        access_token: Ulid,
+        refresh_token: Ulid,
+    },
+
+    #[error("no access token associated with the refresh token {refresh_token:?}")]
+    NoAccessTokenOnRefreshToken { refresh_token: Ulid },
+
     #[error("device code grant expired")]
     DeviceCodeExpired,
 
@@ -122,7 +136,10 @@ impl IntoResponse for RouteError {
             Self::Internal(_)
             | Self::NoSuchBrowserSession
             | Self::NoSuchOAuthSession
-            | Self::ProvisionDeviceFailed(_) => (
+            | Self::ProvisionDeviceFailed(_)
+            | Self::NoSuchNextRefreshToken { .. }
+            | Self::NoSuchNextAccessToken { .. }
+            | Self::NoAccessTokenOnRefreshToken { .. } => (
                 StatusCode::INTERNAL_SERVER_ERROR,
                 Json(ClientError::from(ClientErrorCode::ServerError)),
             ),
@@ -482,6 +499,7 @@ async fn authorization_code_grant(
     Ok((params, repo))
 }
 
+#[allow(clippy::too_many_lines)]
 async fn refresh_token_grant(
     rng: &mut BoxRng,
     clock: &impl Clock,
@@ -518,10 +536,6 @@ async fn refresh_token_grant(
             .await?;
     }
 
-    if !refresh_token.is_valid() {
-        return Err(RouteError::RefreshTokenInvalid(refresh_token.id));
-    }
-
     if !session.is_valid() {
         return Err(RouteError::SessionInvalid(session.id));
     }
@@ -534,6 +548,77 @@ async fn refresh_token_grant(
         });
     }
 
+    if !refresh_token.is_valid() {
+        // We're seing a refresh token that already has been consumed, this might be a
+        // double-refresh or a replay attack
+
+        // First, get the next refresh token
+        let Some(next_refresh_token_id) = refresh_token.next_refresh_token_id() else {
+            // If we don't have a 'next' refresh token, it may just be because this was
+            // before we were recording those. Let's just treat it as a replay.
+            return Err(RouteError::RefreshTokenInvalid(refresh_token.id));
+        };
+
+        let Some(next_refresh_token) = repo
+            .oauth2_refresh_token()
+            .lookup(next_refresh_token_id)
+            .await?
+        else {
+            return Err(RouteError::NoSuchNextRefreshToken {
+                next: next_refresh_token_id,
+                previous: refresh_token.id,
+            });
+        };
+
+        // Check if the next refresh token was already consumed or not
+        if !next_refresh_token.is_valid() {
+            // XXX: This is a replay, we *may* want to invalidate the session
+            return Err(RouteError::RefreshTokenInvalid(next_refresh_token.id));
+        }
+
+        // Check if the associated access token was already used
+        let Some(access_token_id) = next_refresh_token.access_token_id else {
+            // This should in theory not happen: this means an access token got cleaned up,
+            // but the refresh token was still valid.
+            return Err(RouteError::NoAccessTokenOnRefreshToken {
+                refresh_token: next_refresh_token.id,
+            });
+        };
+
+        // Load it
+        let next_access_token = repo
+            .oauth2_access_token()
+            .lookup(access_token_id)
+            .await?
+            .ok_or(RouteError::NoSuchNextAccessToken {
+                access_token: access_token_id,
+                refresh_token: next_refresh_token_id,
+            })?;
+
+        if next_access_token.is_used() {
+            // XXX: This is a replay, we *may* want to invalidate the session
+            return Err(RouteError::RefreshTokenInvalid(next_refresh_token.id));
+        }
+
+        // Looks like it's a double-refresh, client lost their refresh token on
+        // the way back. Let's revoke the two new access and refresh tokens, and
+        // issue new ones
+        info!(
+            oauth2_session.id = %session.id,
+            oauth2_client.id = %client.id,
+            %refresh_token.id,
+            "A refresh token was used twice, but the new refresh token was lost. Revoking the old ones and issuing new ones."
+        );
+
+        repo.oauth2_access_token()
+            .revoke(clock, next_access_token)
+            .await?;
+
+        repo.oauth2_refresh_token()
+            .revoke(clock, next_refresh_token)
+            .await?;
+    }
+
     activity_tracker
         .record_oauth2_session(clock, &session)
         .await;
@@ -550,9 +635,12 @@ async fn refresh_token_grant(
     if let Some(access_token_id) = refresh_token.access_token_id {
         let access_token = repo.oauth2_access_token().lookup(access_token_id).await?;
         if let Some(access_token) = access_token {
-            repo.oauth2_access_token()
-                .revoke(clock, access_token)
-                .await?;
+            // If it is a double-refresh, it might already be revoked
+            if !access_token.state.is_revoked() {
+                repo.oauth2_access_token()
+                    .revoke(clock, access_token)
+                    .await?;
+            }
         }
     }
 
@@ -1123,6 +1211,175 @@ mod tests {
         let _: AccessTokenResponse = response.json();
     }
 
+    #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
+    async fn test_double_refresh(pool: PgPool) {
+        setup();
+        let state = TestState::from_pool(pool).await.unwrap();
+
+        // Provision a client
+        let request =
+            Request::post(mas_router::OAuth2RegistrationEndpoint::PATH).json(serde_json::json!({
+                "client_uri": "https://example.com/",
+                "redirect_uris": ["https://example.com/callback"],
+                "token_endpoint_auth_method": "none",
+                "response_types": ["code"],
+                "grant_types": ["authorization_code", "refresh_token"],
+            }));
+
+        let response = state.request(request).await;
+        response.assert_status(StatusCode::CREATED);
+
+        let ClientRegistrationResponse { client_id, .. } = response.json();
+
+        // Let's provision a user and create a session for them. This part is hard to
+        // test with just HTTP requests, so we'll use the repository directly.
+        let mut repo = state.repository().await.unwrap();
+
+        let user = repo
+            .user()
+            .add(&mut state.rng(), &state.clock, "alice".to_owned())
+            .await
+            .unwrap();
+
+        let browser_session = repo
+            .browser_session()
+            .add(&mut state.rng(), &state.clock, &user, None)
+            .await
+            .unwrap();
+
+        // Lookup the client in the database.
+        let client = repo
+            .oauth2_client()
+            .find_by_client_id(&client_id)
+            .await
+            .unwrap()
+            .unwrap();
+
+        // Get a token pair
+        let session = repo
+            .oauth2_session()
+            .add_from_browser_session(
+                &mut state.rng(),
+                &state.clock,
+                &client,
+                &browser_session,
+                Scope::from_iter([OPENID]),
+            )
+            .await
+            .unwrap();
+
+        let (AccessToken { access_token, .. }, RefreshToken { refresh_token, .. }) =
+            generate_token_pair(
+                &mut state.rng(),
+                &state.clock,
+                &mut repo,
+                &session,
+                Duration::microseconds(5 * 60 * 1000 * 1000),
+            )
+            .await
+            .unwrap();
+
+        repo.save().await.unwrap();
+
+        // First check that the token is valid
+        assert!(state.is_access_token_valid(&access_token).await);
+
+        // Now call the token endpoint to get an access token.
+        let request =
+            Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({
+                "grant_type": "refresh_token",
+                "refresh_token": refresh_token,
+                "client_id": client.client_id,
+            }));
+
+        let first_response = state.request(request).await;
+        first_response.assert_status(StatusCode::OK);
+        let first_response: AccessTokenResponse = first_response.json();
+
+        // Call a second time, it should work, as we haven't done anything yet with the
+        // token
+        let request =
+            Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({
+                "grant_type": "refresh_token",
+                "refresh_token": refresh_token,
+                "client_id": client.client_id,
+            }));
+
+        let second_response = state.request(request).await;
+        second_response.assert_status(StatusCode::OK);
+        let second_response: AccessTokenResponse = second_response.json();
+
+        // Check that we got new tokens
+        assert_ne!(first_response.access_token, second_response.access_token);
+        assert_ne!(first_response.refresh_token, second_response.refresh_token);
+
+        // Check that the old-new token is invalid
+        assert!(
+            !state
+                .is_access_token_valid(&first_response.access_token)
+                .await
+        );
+
+        // Check that the new-new token is valid
+        assert!(
+            state
+                .is_access_token_valid(&second_response.access_token)
+                .await
+        );
+
+        // Do a third refresh, this one should not work, as we've used the new
+        // access token
+        let request =
+            Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({
+                "grant_type": "refresh_token",
+                "refresh_token": refresh_token,
+                "client_id": client.client_id,
+            }));
+
+        let third_response = state.request(request).await;
+        third_response.assert_status(StatusCode::BAD_REQUEST);
+
+        // The other reason we consider a new refresh token to be 'used' is if
+        // it was already used in a refresh
+        // So, if we do a refresh with the second_response.refresh_token, then
+        // another refresh with the result, redoing one with
+        // second_response.refresh_token again should fail
+        let request =
+            Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({
+                "grant_type": "refresh_token",
+                "refresh_token": second_response.refresh_token,
+                "client_id": client.client_id,
+            }));
+
+        // This one is fine
+        let fourth_response = state.request(request).await;
+        fourth_response.assert_status(StatusCode::OK);
+        let fourth_response: AccessTokenResponse = fourth_response.json();
+
+        // Do another one, it should be fine as well
+        let request =
+            Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({
+                "grant_type": "refresh_token",
+                "refresh_token": fourth_response.refresh_token,
+                "client_id": client.client_id,
+            }));
+
+        let fifth_response = state.request(request).await;
+        fifth_response.assert_status(StatusCode::OK);
+
+        // But now, if we re-do with the second_response.refresh_token, it should
+        // fail
+        let request =
+            Request::post(mas_router::OAuth2TokenEndpoint::PATH).form(serde_json::json!({
+                "grant_type": "refresh_token",
+                "refresh_token": second_response.refresh_token,
+                "client_id": client.client_id,
+            }));
+
+        let sixth_response = state.request(request).await;
+        sixth_response.assert_status(StatusCode::BAD_REQUEST);
+    }
+
     #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
     async fn test_client_credentials(pool: PgPool) {
         setup();

From e64e5f5c4c2ae77c60404b05ab12b4c9dd75cb0b Mon Sep 17 00:00:00 2001
From: Quentin Gliech <quenting@element.io>
Date: Wed, 11 Dec 2024 13:55:46 +0100
Subject: [PATCH 7/7] Apply suggestions from code review

Co-authored-by: reivilibre <oliverw@element.io>
---
 crates/handlers/src/oauth2/token.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/crates/handlers/src/oauth2/token.rs b/crates/handlers/src/oauth2/token.rs
index 9e8e201a0..ac64da50d 100644
--- a/crates/handlers/src/oauth2/token.rs
+++ b/crates/handlers/src/oauth2/token.rs
@@ -601,13 +601,13 @@ async fn refresh_token_grant(
         }
 
         // Looks like it's a double-refresh, client lost their refresh token on
-        // the way back. Let's revoke the two new access and refresh tokens, and
+        // the way back. Let's revoke the unused access and refresh tokens, and
         // issue new ones
         info!(
             oauth2_session.id = %session.id,
             oauth2_client.id = %client.id,
             %refresh_token.id,
-            "A refresh token was used twice, but the new refresh token was lost. Revoking the old ones and issuing new ones."
+            "Refresh token already used, but issued refresh and access tokens are unused. Assuming those were lost; revoking those and reissuing new ones."
         );
 
         repo.oauth2_access_token()