From 2ae95e30ec273e7f4e06c20eca9dcb5739252ec5 Mon Sep 17 00:00:00 2001 From: Quentin Gliech Date: Fri, 16 Jan 2026 13:25:58 +0100 Subject: [PATCH] Implement cleanup job for user recovery sessions Add scheduled cleanup job that removes old user recovery sessions after 7 days. Runs hourly. Implementation uses ULID cursor-based pagination with no additional indexes needed. Child tickets cascade-delete automatically. --- ...47e411b96b2376288a90c242034295e1a147e.json | 24 ++++++++ crates/storage-pg/src/user/recovery.rs | 50 +++++++++++++++++ crates/storage/src/queue/tasks.rs | 8 +++ crates/storage/src/user/recovery.rs | 33 +++++++++++ crates/tasks/src/database.rs | 56 ++++++++++++++++++- crates/tasks/src/lib.rs | 7 +++ 6 files changed, 177 insertions(+), 1 deletion(-) create mode 100644 crates/storage-pg/.sqlx/query-8ef977487429f84c557dc62272c47e411b96b2376288a90c242034295e1a147e.json diff --git a/crates/storage-pg/.sqlx/query-8ef977487429f84c557dc62272c47e411b96b2376288a90c242034295e1a147e.json b/crates/storage-pg/.sqlx/query-8ef977487429f84c557dc62272c47e411b96b2376288a90c242034295e1a147e.json new file mode 100644 index 000000000..aa2d4a1c8 --- /dev/null +++ b/crates/storage-pg/.sqlx/query-8ef977487429f84c557dc62272c47e411b96b2376288a90c242034295e1a147e.json @@ -0,0 +1,24 @@ +{ + "db_name": "PostgreSQL", + "query": "\n WITH to_delete AS (\n SELECT user_recovery_session_id\n FROM user_recovery_sessions\n WHERE ($1::uuid IS NULL OR user_recovery_session_id > $1)\n AND user_recovery_session_id <= $2\n ORDER BY user_recovery_session_id\n LIMIT $3\n )\n DELETE FROM user_recovery_sessions\n USING to_delete\n WHERE user_recovery_sessions.user_recovery_session_id = to_delete.user_recovery_session_id\n RETURNING user_recovery_sessions.user_recovery_session_id\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "user_recovery_session_id", + "type_info": "Uuid" + } + ], + "parameters": { + "Left": [ + "Uuid", + "Uuid", + "Int8" + ] + }, + "nullable": [ + false + ] + }, + "hash": "8ef977487429f84c557dc62272c47e411b96b2376288a90c242034295e1a147e" +} diff --git a/crates/storage-pg/src/user/recovery.rs b/crates/storage-pg/src/user/recovery.rs index 30ab64bcb..800cdcb70 100644 --- a/crates/storage-pg/src/user/recovery.rs +++ b/crates/storage-pg/src/user/recovery.rs @@ -1,3 +1,4 @@ +// Copyright 2025, 2026 Element Creations Ltd. // Copyright 2024, 2025 New Vector Ltd. // Copyright 2024 The Matrix.org Foundation C.I.C. // @@ -326,4 +327,53 @@ impl UserRecoveryRepository for PgUserRecoveryRepository<'_> { Ok(user_recovery_session) } + + #[tracing::instrument( + name = "db.user_recovery.cleanup", + skip_all, + fields( + db.query.text, + since = since.map(tracing::field::display), + until = %until, + limit = limit, + ), + err, + )] + async fn cleanup( + &mut self, + since: Option, + until: Ulid, + limit: usize, + ) -> Result<(usize, Option), Self::Error> { + // Use ULID cursor-based pagination. Since ULIDs contain a timestamp, + // we can efficiently delete old sessions without needing an index. + // `MAX(uuid)` isn't a thing in Postgres, so we aggregate on the client side. + let res = sqlx::query_scalar!( + r#" + WITH to_delete AS ( + SELECT user_recovery_session_id + FROM user_recovery_sessions + WHERE ($1::uuid IS NULL OR user_recovery_session_id > $1) + AND user_recovery_session_id <= $2 + ORDER BY user_recovery_session_id + LIMIT $3 + ) + DELETE FROM user_recovery_sessions + USING to_delete + WHERE user_recovery_sessions.user_recovery_session_id = to_delete.user_recovery_session_id + RETURNING user_recovery_sessions.user_recovery_session_id + "#, + since.map(Uuid::from), + Uuid::from(until), + i64::try_from(limit).unwrap_or(i64::MAX) + ) + .traced() + .fetch_all(&mut *self.conn) + .await?; + + let count = res.len(); + let max_id = res.into_iter().max(); + + Ok((count, max_id.map(Ulid::from))) + } } diff --git a/crates/storage/src/queue/tasks.rs b/crates/storage/src/queue/tasks.rs index 2607cb765..c331fff4b 100644 --- a/crates/storage/src/queue/tasks.rs +++ b/crates/storage/src/queue/tasks.rs @@ -382,6 +382,14 @@ impl InsertableJob for CleanupOAuthDeviceCodeGrantsJob { const QUEUE_NAME: &'static str = "cleanup-oauth-device-code-grants"; } +/// Cleanup old user recovery sessions +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +pub struct CleanupUserRecoverySessionsJob; + +impl InsertableJob for CleanupUserRecoverySessionsJob { + const QUEUE_NAME: &'static str = "cleanup-user-recovery-sessions"; +} + /// Scheduled job to expire inactive sessions /// /// This job will trigger jobs to expire inactive compat, oauth and user diff --git a/crates/storage/src/user/recovery.rs b/crates/storage/src/user/recovery.rs index 2bed10dfc..47c6bef59 100644 --- a/crates/storage/src/user/recovery.rs +++ b/crates/storage/src/user/recovery.rs @@ -1,3 +1,4 @@ +// Copyright 2025, 2026 Element Creations Ltd. // Copyright 2024, 2025 New Vector Ltd. // Copyright 2024 The Matrix.org Foundation C.I.C. // @@ -121,6 +122,31 @@ pub trait UserRecoveryRepository: Send + Sync { user_recovery_ticket: UserRecoveryTicket, user_recovery_session: UserRecoverySession, ) -> Result; + + /// Cleanup old recovery sessions + /// + /// This will delete recovery sessions with IDs up to and including `until`. + /// Uses ULID cursor-based pagination for efficiency. + /// Tickets will cascade-delete automatically. + /// + /// Returns the number of sessions deleted and the cursor for the next batch + /// + /// # Parameters + /// + /// * `since`: The cursor to start from (exclusive), or `None` to start from + /// the beginning + /// * `until`: The maximum ULID to delete (inclusive upper bound) + /// * `limit`: The maximum number of sessions to delete in this batch + /// + /// # Errors + /// + /// Returns [`Self::Error`] if the underlying repository fails + async fn cleanup( + &mut self, + since: Option, + until: Ulid, + limit: usize, + ) -> Result<(usize, Option), Self::Error>; } repository_impl!(UserRecoveryRepository: @@ -156,4 +182,11 @@ repository_impl!(UserRecoveryRepository: user_recovery_ticket: UserRecoveryTicket, user_recovery_session: UserRecoverySession, ) -> Result; + + async fn cleanup( + &mut self, + since: Option, + until: Ulid, + limit: usize, + ) -> Result<(usize, Option), Self::Error>; ); diff --git a/crates/tasks/src/database.rs b/crates/tasks/src/database.rs index f0c9f2dc7..a58885329 100644 --- a/crates/tasks/src/database.rs +++ b/crates/tasks/src/database.rs @@ -14,7 +14,8 @@ use mas_storage::queue::{ CleanupConsumedOAuthRefreshTokensJob, CleanupExpiredOAuthAccessTokensJob, CleanupFinishedCompatSessionsJob, CleanupOAuthAuthorizationGrantsJob, CleanupOAuthDeviceCodeGrantsJob, CleanupRevokedOAuthAccessTokensJob, - CleanupRevokedOAuthRefreshTokensJob, CleanupUserRegistrationsJob, PruneStalePolicyDataJob, + CleanupRevokedOAuthRefreshTokensJob, CleanupUserRecoverySessionsJob, + CleanupUserRegistrationsJob, PruneStalePolicyDataJob, }; use tracing::{debug, info}; use ulid::Ulid; @@ -219,6 +220,59 @@ impl RunnableJob for CleanupConsumedOAuthRefreshTokensJob { } } +#[async_trait] +impl RunnableJob for CleanupUserRecoverySessionsJob { + #[tracing::instrument(name = "job.cleanup_user_recovery_sessions", skip_all)] + async fn run(&self, state: &State, context: JobContext) -> Result<(), JobError> { + // Remove recovery sessions after 7 days. They are in practice only + // valid for a short time (tickets expire after 10 minutes), but keeping + // them around helps investigate abuse patterns. + let until = state.clock.now() - chrono::Duration::days(7); + // We use the fact that ULIDs include the creation time in their first 48 bits + // as a cursor + let until = Ulid::from_parts( + u64::try_from(until.timestamp_millis()).unwrap_or(u64::MIN), + u128::MAX, + ); + let mut total = 0; + + // Run until we get cancelled. We don't schedule a retry if we get cancelled, as + // this is a scheduled job and it will end up being rescheduled later anyway. + let mut since = None; + while !context.cancellation_token.is_cancelled() { + let mut repo = state.repository().await.map_err(JobError::retry)?; + // This returns the number of deleted sessions, and the greatest ULID processed + let (count, cursor) = repo + .user_recovery() + .cleanup(since, until, BATCH_SIZE) + .await + .map_err(JobError::retry)?; + repo.save().await.map_err(JobError::retry)?; + since = cursor; + total += count; + + // Check how many we deleted. If we deleted exactly BATCH_SIZE, + // there might be more to delete + if count != BATCH_SIZE { + break; + } + } + + if total == 0 { + debug!("no user recovery sessions to clean up"); + } else { + info!(count = total, "cleaned up user recovery sessions"); + } + + Ok(()) + } + + fn timeout(&self) -> Option { + // This job runs every hour, so having it running it for 10 minutes is fine + Some(Duration::from_secs(10 * 60)) + } +} + #[async_trait] impl RunnableJob for CleanupUserRegistrationsJob { #[tracing::instrument(name = "job.cleanup_user_registrations", skip_all)] diff --git a/crates/tasks/src/lib.rs b/crates/tasks/src/lib.rs index 53d780faa..fa5d4f188 100644 --- a/crates/tasks/src/lib.rs +++ b/crates/tasks/src/lib.rs @@ -137,6 +137,7 @@ pub async fn init( .register_handler::() .register_handler::() .register_handler::() + .register_handler::() .register_handler::() .register_handler::() .register_handler::() @@ -194,6 +195,12 @@ pub async fn init( "0 55 * * * *".parse()?, mas_storage::queue::CleanupOAuthDeviceCodeGrantsJob, ) + .add_schedule( + "cleanup-user-recovery-sessions", + // Run this job every hour + "0 56 * * * *".parse()?, + mas_storage::queue::CleanupUserRecoverySessionsJob, + ) .add_schedule( "cleanup-expired-oauth-access-tokens", // Run this job every 4 hours