Make sure refreshing OAuth token is idempotent

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)
     }
 }

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
@@ -109,6 +116,10 @@ pub enum RefreshTokenState {
     Valid,
     Consumed {
         consumed_at: DateTime<Utc>,
+        next_refresh_token_id: Option<Ulid>,
+    },
+    Revoked {
+        revoked_at: DateTime<Utc>,
     },
 }
 
@@ -117,11 +128,30 @@ impl RefreshTokenState {
     ///
     /// # Errors
     ///
-    /// Returns an error if the refresh token is already consumed.
-    fn consume(self, consumed_at: DateTime<Utc>) -> Result<Self, InvalidTransitionError> {
+    /// 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 | Self::Consumed { .. } => Ok(Self::Consumed {
+                consumed_at,
+                next_refresh_token_id: Some(replaced_by.id),
+            }),
+            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::Consumed { consumed_at }),
-            Self::Consumed { .. } => Err(InvalidTransitionError),
+            Self::Valid => Ok(Self::Revoked { revoked_at }),
+            Self::Consumed { .. } | Self::Revoked { .. } => Err(InvalidTransitionError),
         }
     }
 
@@ -133,12 +163,16 @@ impl RefreshTokenState {
         matches!(self, Self::Valid)
     }
 
-    /// Returns `true` if the refresh token state is [`Consumed`].
-    ///
-    /// [`Consumed`]: RefreshTokenState::Consumed
+    /// Returns the next refresh token ID, if any.
     #[must_use]
-    pub fn is_consumed(&self) -> bool {
-        matches!(self, Self::Consumed { .. })
+    pub fn next_refresh_token_id(&self) -> Option<Ulid> {
+        match self {
+            Self::Valid | Self::Revoked { .. } => None,
+            Self::Consumed {
+                next_refresh_token_id,
+                ..
+            } => *next_refresh_token_id,
+        }
     }
 }
 
@@ -170,9 +204,23 @@ 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)?;
+    /// Returns an error if the refresh token is revoked.
+    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)
+    }
+
+    /// 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)
     }
 }

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