Setup a job to expire OAuth 2.0 sessions

Author: sandhose

crates/storage/src/queue/tasks.rs

@@ -3,11 +3,13 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Please see LICENSE in the repository root for full details.
 
-use mas_data_model::{Device, User, UserEmailAuthentication, UserRecoverySession};
+use chrono::{DateTime, Utc};
+use mas_data_model::{Device, Session, User, UserEmailAuthentication, UserRecoverySession};
 use serde::{Deserialize, Serialize};
 use ulid::Ulid;
 
 use super::InsertableJob;
+use crate::{Page, Pagination};
 
 /// This is the previous iteration of the email verification job. It has been
 /// replaced by [`SendEmailAuthenticationCodeJob`]. This struct is kept to be
@@ -193,6 +195,15 @@ impl SyncDevicesJob {
         Self { user_id: user.id }
     }
 
+    /// Create a new job to sync the list of devices of a user with the
+    /// homeserver for the given user ID
+    ///
+    /// This is useful to use in cases where the [`User`] object isn't loaded
+    #[must_use]
+    pub fn new_for_id(user_id: Ulid) -> Self {
+        Self { user_id }
+    }
+
     /// The ID of the user to sync the devices for
     #[must_use]
     pub fn user_id(&self) -> Ulid {
@@ -310,3 +321,60 @@ pub struct CleanupExpiredTokensJob;
 impl InsertableJob for CleanupExpiredTokensJob {
     const QUEUE_NAME: &'static str = "cleanup-expired-tokens";
 }
+
+/// Expire inactive OAuth 2.0 sessions
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct ExpireInactiveOAuthSessionsJob {
+    threshold: DateTime<Utc>,
+    after: Option<Ulid>,
+}
+
+impl ExpireInactiveOAuthSessionsJob {
+    /// Create a new job to expire inactive OAuth 2.0 sessions
+    ///
+    /// # Parameters
+    ///
+    /// * `threshold` - The threshold to expire sessions at
+    #[must_use]
+    pub fn new(threshold: DateTime<Utc>) -> Self {
+        Self {
+            threshold,
+            after: None,
+        }
+    }
+
+    /// Get the threshold to expire sessions at
+    #[must_use]
+    pub fn threshold(&self) -> DateTime<Utc> {
+        self.threshold
+    }
+
+    /// Get the pagination cursor
+    #[must_use]
+    pub fn pagination(&self, batch_size: usize) -> Pagination {
+        let pagination = Pagination::first(batch_size);
+        if let Some(after) = self.after {
+            pagination.after(after)
+        } else {
+            pagination
+        }
+    }
+
+    /// Get the next job given the page returned by the database
+    #[must_use]
+    pub fn next(&self, page: &Page<Session>) -> Option<Self> {
+        if !page.has_next_page {
+            return None;
+        }
+
+        let last_edge = page.edges.last()?;
+        Some(Self {
+            threshold: self.threshold,
+            after: Some(last_edge.id),
+        })
+    }
+}
+
+impl InsertableJob for ExpireInactiveOAuthSessionsJob {
+    const QUEUE_NAME: &'static str = "expire-inactive-oauth-sessions";
+}

crates/tasks/src/lib.rs

@@ -22,6 +22,7 @@ mod email;
 mod matrix;
 mod new_queue;
 mod recovery;
+mod sessions;
 mod user;
 
 static METER: LazyLock<Meter> = LazyLock::new(|| {
@@ -128,6 +129,7 @@ pub async fn init(
         .register_handler::<mas_storage::queue::SendEmailAuthenticationCodeJob>()
         .register_handler::<mas_storage::queue::SyncDevicesJob>()
         .register_handler::<mas_storage::queue::VerifyEmailJob>()
+        .register_handler::<mas_storage::queue::ExpireInactiveOAuthSessionsJob>()
         .add_schedule(
             "cleanup-expired-tokens",
             "0 0 * * * *".parse()?,

crates/tasks/src/sessions.rs

@@ -0,0 +1,82 @@
+// Copyright 2025 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+
+use std::collections::HashSet;
+
+use async_trait::async_trait;
+use chrono::Duration;
+use mas_storage::{
+    oauth2::OAuth2SessionFilter,
+    queue::{ExpireInactiveOAuthSessionsJob, QueueJobRepositoryExt, SyncDevicesJob},
+};
+
+use crate::{
+    new_queue::{JobContext, JobError, RunnableJob},
+    State,
+};
+
+#[async_trait]
+impl RunnableJob for ExpireInactiveOAuthSessionsJob {
+    async fn run(&self, state: &State, _context: JobContext) -> Result<(), JobError> {
+        let mut repo = state.repository().await.map_err(JobError::retry)?;
+        let clock = state.clock();
+        let mut rng = state.rng();
+        let mut users_synced = HashSet::new();
+
+        // This delay is used to space out the device sync jobs
+        // We add 10 seconds between each device sync, meaning that it will spread out
+        // the syncs over ~16 minutes max if we get a full batch of 100 users
+        let mut delay = Duration::minutes(1);
+
+        let filter = OAuth2SessionFilter::new()
+            .with_last_active_before(self.threshold())
+            .for_any_user()
+            .active_only();
+
+        let pagination = self.pagination(100);
+
+        let page = repo
+            .oauth2_session()
+            .list(filter, pagination)
+            .await
+            .map_err(JobError::retry)?;
+
+        if let Some(job) = self.next(&page) {
+            tracing::info!("Scheduling job to expire the next batch of inactive sessions");
+            repo.queue_job()
+                .schedule_job(&mut rng, &clock, job)
+                .await
+                .map_err(JobError::retry)?;
+        }
+
+        for edge in page.edges {
+            if let Some(user_id) = edge.user_id {
+                let inserted = users_synced.insert(user_id);
+                if inserted {
+                    tracing::info!(user.id = %user_id, "Scheduling devices sync for user");
+                    repo.queue_job()
+                        .schedule_job_later(
+                            &mut rng,
+                            &clock,
+                            SyncDevicesJob::new_for_id(user_id),
+                            clock.now() + delay,
+                        )
+                        .await
+                        .map_err(JobError::retry)?;
+                    delay += Duration::seconds(10);
+                }
+            }
+
+            repo.oauth2_session()
+                .finish(&clock, edge)
+                .await
+                .map_err(JobError::retry)?;
+        }
+
+        repo.save().await.map_err(JobError::retry)?;
+
+        Ok(())
+    }
+}