From 7349da088990fe7c093ccd796944729a9ecd899c Mon Sep 17 00:00:00 2001 From: Quentin Gliech Date: Fri, 9 Jan 2026 18:36:58 +0100 Subject: [PATCH] Cleanup revoked refresh tokens --- ...e2777f9327708b450d048638a162343478cc6.json | 30 ++++++++++ ...9172537_oauth_refresh_token_revoked_at.sql | 9 +++ crates/storage-pg/src/oauth2/refresh_token.rs | 55 +++++++++++++++++++ crates/storage/src/oauth2/refresh_token.rs | 29 ++++++++++ crates/storage/src/queue/tasks.rs | 8 +++ crates/tasks/src/database.rs | 49 ++++++++++++++++- crates/tasks/src/lib.rs | 7 +++ 7 files changed, 186 insertions(+), 1 deletion(-) create mode 100644 crates/storage-pg/.sqlx/query-31e8bf68ff70a436fd0b6787ac8e2777f9327708b450d048638a162343478cc6.json create mode 100644 crates/storage-pg/migrations/20260109172537_oauth_refresh_token_revoked_at.sql diff --git a/crates/storage-pg/.sqlx/query-31e8bf68ff70a436fd0b6787ac8e2777f9327708b450d048638a162343478cc6.json b/crates/storage-pg/.sqlx/query-31e8bf68ff70a436fd0b6787ac8e2777f9327708b450d048638a162343478cc6.json new file mode 100644 index 000000000..84e37eb52 --- /dev/null +++ b/crates/storage-pg/.sqlx/query-31e8bf68ff70a436fd0b6787ac8e2777f9327708b450d048638a162343478cc6.json @@ -0,0 +1,30 @@ +{ + "db_name": "PostgreSQL", + "query": "\n WITH\n to_delete AS (\n SELECT oauth2_refresh_token_id\n FROM oauth2_refresh_tokens\n WHERE revoked_at IS NOT NULL\n AND ($1::timestamptz IS NULL OR revoked_at >= $1::timestamptz)\n AND revoked_at < $2::timestamptz\n ORDER BY revoked_at ASC\n LIMIT $3\n FOR UPDATE\n ),\n\n deleted AS (\n DELETE FROM oauth2_refresh_tokens\n USING to_delete\n WHERE oauth2_refresh_tokens.oauth2_refresh_token_id = to_delete.oauth2_refresh_token_id\n RETURNING oauth2_refresh_tokens.revoked_at\n )\n\n SELECT\n COUNT(*) as \"count!\",\n MAX(revoked_at) as last_revoked_at\n FROM deleted\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "count!", + "type_info": "Int8" + }, + { + "ordinal": 1, + "name": "last_revoked_at", + "type_info": "Timestamptz" + } + ], + "parameters": { + "Left": [ + "Timestamptz", + "Timestamptz", + "Int8" + ] + }, + "nullable": [ + null, + null + ] + }, + "hash": "31e8bf68ff70a436fd0b6787ac8e2777f9327708b450d048638a162343478cc6" +} diff --git a/crates/storage-pg/migrations/20260109172537_oauth_refresh_token_revoked_at.sql b/crates/storage-pg/migrations/20260109172537_oauth_refresh_token_revoked_at.sql new file mode 100644 index 000000000..6b982e835 --- /dev/null +++ b/crates/storage-pg/migrations/20260109172537_oauth_refresh_token_revoked_at.sql @@ -0,0 +1,9 @@ +-- no-transaction +-- Copyright 2026 Element Creations Ltd. +-- +-- SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial +-- Please see LICENSE in the repository root for full details. + +-- This adds an index on the revoked_at field on oauth2_refresh_tokens to speed up cleaning them up +CREATE INDEX CONCURRENTLY IF NOT EXISTS oauth_refresh_tokens_revoked_at_idx + ON oauth2_refresh_tokens (revoked_at) WHERE revoked_at IS NOT NULL; diff --git a/crates/storage-pg/src/oauth2/refresh_token.rs b/crates/storage-pg/src/oauth2/refresh_token.rs index 5b49e4bd7..d98e986c8 100644 --- a/crates/storage-pg/src/oauth2/refresh_token.rs +++ b/crates/storage-pg/src/oauth2/refresh_token.rs @@ -1,3 +1,4 @@ +// Copyright 2025, 2026 Element Creations Ltd. // Copyright 2024, 2025 New Vector Ltd. // Copyright 2021-2024 The Matrix.org Foundation C.I.C. // @@ -281,4 +282,58 @@ impl OAuth2RefreshTokenRepository for PgOAuth2RefreshTokenRepository<'_> { .revoke(revoked_at) .map_err(DatabaseError::to_invalid_operation) } + + #[tracing::instrument( + name = "db.oauth2_refresh_token.cleanup_revoked", + skip_all, + fields( + db.query.text, + ), + err, + )] + async fn cleanup_revoked( + &mut self, + since: Option>, + until: DateTime, + limit: usize, + ) -> Result<(usize, Option>), Self::Error> { + let res = sqlx::query!( + r#" + WITH + to_delete AS ( + SELECT oauth2_refresh_token_id + FROM oauth2_refresh_tokens + WHERE revoked_at IS NOT NULL + AND ($1::timestamptz IS NULL OR revoked_at >= $1::timestamptz) + AND revoked_at < $2::timestamptz + ORDER BY revoked_at ASC + LIMIT $3 + FOR UPDATE + ), + + deleted AS ( + DELETE FROM oauth2_refresh_tokens + USING to_delete + WHERE oauth2_refresh_tokens.oauth2_refresh_token_id = to_delete.oauth2_refresh_token_id + RETURNING oauth2_refresh_tokens.revoked_at + ) + + SELECT + COUNT(*) as "count!", + MAX(revoked_at) as last_revoked_at + FROM deleted + "#, + since, + until, + i64::try_from(limit).unwrap_or(i64::MAX), + ) + .traced() + .fetch_one(&mut *self.conn) + .await?; + + Ok(( + res.count.try_into().unwrap_or(usize::MAX), + res.last_revoked_at, + )) + } } diff --git a/crates/storage/src/oauth2/refresh_token.rs b/crates/storage/src/oauth2/refresh_token.rs index b81420a08..1dd55db16 100644 --- a/crates/storage/src/oauth2/refresh_token.rs +++ b/crates/storage/src/oauth2/refresh_token.rs @@ -1,3 +1,4 @@ +// Copyright 2025, 2026 Element Creations Ltd. // Copyright 2024, 2025 New Vector Ltd. // Copyright 2021-2024 The Matrix.org Foundation C.I.C. // @@ -111,6 +112,27 @@ pub trait OAuth2RefreshTokenRepository: Send + Sync { clock: &dyn Clock, refresh_token: RefreshToken, ) -> Result; + + /// Cleanup revoked refresh tokens that were revoked before a certain time + /// + /// Returns the number of deleted tokens and the last `revoked_at` timestamp + /// processed + /// + /// # Parameters + /// + /// * `since`: An optional timestamp to start from + /// * `until`: The timestamp before which to revoke tokens + /// * `limit`: The maximum number of tokens to revoke + /// + /// # Errors + /// + /// Returns [`Self::Error`] if the underlying repository fails + async fn cleanup_revoked( + &mut self, + since: Option>, + until: chrono::DateTime, + limit: usize, + ) -> Result<(usize, Option>), Self::Error>; } repository_impl!(OAuth2RefreshTokenRepository: @@ -142,4 +164,11 @@ repository_impl!(OAuth2RefreshTokenRepository: clock: &dyn Clock, refresh_token: RefreshToken, ) -> Result; + + async fn cleanup_revoked( + &mut self, + since: Option>, + until: chrono::DateTime, + limit: usize, + ) -> Result<(usize, Option>), Self::Error>; ); diff --git a/crates/storage/src/queue/tasks.rs b/crates/storage/src/queue/tasks.rs index 4a1373fe0..b033f68f7 100644 --- a/crates/storage/src/queue/tasks.rs +++ b/crates/storage/src/queue/tasks.rs @@ -334,6 +334,14 @@ impl InsertableJob for CleanupExpiredOAuthAccessTokensJob { const QUEUE_NAME: &'static str = "cleanup-expired-oauth-access-tokens"; } +/// Cleanup revoked OAuth 2.0 refresh tokens +#[derive(Serialize, Deserialize, Debug, Clone, Default)] +pub struct CleanupRevokedOAuthRefreshTokensJob; + +impl InsertableJob for CleanupRevokedOAuthRefreshTokensJob { + const QUEUE_NAME: &'static str = "cleanup-revoked-oauth-refresh-tokens"; +} + /// Scheduled job to expire inactive sessions /// /// This job will trigger jobs to expire inactive compat, oauth and user diff --git a/crates/tasks/src/database.rs b/crates/tasks/src/database.rs index 798141b76..f1cffea3f 100644 --- a/crates/tasks/src/database.rs +++ b/crates/tasks/src/database.rs @@ -11,7 +11,8 @@ use std::time::Duration; use async_trait::async_trait; use mas_storage::queue::{ - CleanupExpiredOAuthAccessTokensJob, CleanupRevokedOAuthAccessTokensJob, PruneStalePolicyDataJob, + CleanupExpiredOAuthAccessTokensJob, CleanupRevokedOAuthAccessTokensJob, + CleanupRevokedOAuthRefreshTokensJob, PruneStalePolicyDataJob, }; use tracing::{debug, info}; @@ -120,6 +121,52 @@ impl RunnableJob for CleanupExpiredOAuthAccessTokensJob { } } +#[async_trait] +impl RunnableJob for CleanupRevokedOAuthRefreshTokensJob { + #[tracing::instrument(name = "job.cleanup_revoked_oauth_refresh_tokens", skip_all)] + async fn run(&self, state: &State, context: JobContext) -> Result<(), JobError> { + // Cleanup tokens that were revoked more than an hour ago + let until = state.clock.now() - chrono::Duration::hours(1); + 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 tokens, and the last revoked_at timestamp + let (count, last_revoked_at) = repo + .oauth2_refresh_token() + .cleanup_revoked(since, until, BATCH_SIZE) + .await + .map_err(JobError::retry)?; + repo.save().await.map_err(JobError::retry)?; + + since = last_revoked_at; + 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 token to clean up"); + } else { + info!(count = total, "cleaned up revoked tokens"); + } + + Ok(()) + } + + fn timeout(&self) -> Option { + Some(Duration::from_secs(60)) + } +} + #[async_trait] impl RunnableJob for PruneStalePolicyDataJob { #[tracing::instrument(name = "job.prune_stale_policy_data", skip_all)] diff --git a/crates/tasks/src/lib.rs b/crates/tasks/src/lib.rs index b0d2b3d72..eb1ee3bc9 100644 --- a/crates/tasks/src/lib.rs +++ b/crates/tasks/src/lib.rs @@ -131,6 +131,7 @@ pub async fn init( worker .register_handler::() .register_handler::() + .register_handler::() .register_handler::() .register_handler::() .register_handler::() @@ -152,6 +153,12 @@ pub async fn init( "0 0 * * * *".parse()?, mas_storage::queue::CleanupRevokedOAuthAccessTokensJob, ) + .add_schedule( + "cleanup-revoked-oauth-refresh-tokens", + // Run this job every hour + "0 10 * * * *".parse()?, + mas_storage::queue::CleanupRevokedOAuthRefreshTokensJob, + ) .add_schedule( "cleanup-expired-oauth-access-tokens", // Run this job every 4 hours