Record when access tokens are first used

Author: sandhose

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

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

crates/storage-pg/.sqlx/query-7189b6136fd08ac9ae7c51bff06fb2254d1bf9e8a97cd7d32ba789c740e0fbdb.json

@@ -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;

crates/storage-pg/.sqlx/query-dd16942318bf38d9a245b2c86fedd3cbd6b65e7a13465552d79cd3c022122fd4.json renamed to crates/storage-pg/.sqlx/query-875294dc5cf87bcf302fb9e87933745cc1c57bbe3c3c69110592a07400116c7f.json

@@ -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?;
 

crates/storage-pg/.sqlx/query-477f79556e5777b38feb85013b4f04dbb8230e4b0b0bcc45f669d7b8d0b91db4.json renamed to crates/storage-pg/.sqlx/query-c09e0bb0378d9dfb15de7f2f1209fab6ea87589819128e6fc9ed5da11dfc2770.json

@@ -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>;
 );