Compose filters for batch logging out of browser sessions

Instead of having to load all authentication sessions in memory, we
allow composing browser session filters with a upstream auth sessions
filter

Author: sandhose

crates/handlers/src/upstream_oauth2/backchannel_logout.rs

@@ -22,7 +22,7 @@ use mas_oidc_client::{
     requests::jose::{JwtVerificationData, verify_signed_jwt},
 };
 use mas_storage::{
-    BoxClock, BoxRepository, Pagination, upstream_oauth2::UpstreamOAuthSessionFilter,
+    BoxClock, BoxRepository, upstream_oauth2::UpstreamOAuthSessionFilter,
     user::BrowserSessionFilter,
 };
 use oauth2_types::errors::{ClientError, ClientErrorCode};
@@ -222,41 +222,27 @@ pub(crate) async fn post(
     claims::NONCE.assert_absent(&claims)?; // (7)
 
     // Find the corresponding upstream OAuth 2.0 sessions
-    let mut filter = UpstreamOAuthSessionFilter::new().for_provider(&provider);
+    let mut auth_session_filter = UpstreamOAuthSessionFilter::new().for_provider(&provider);
     if let Some(sub) = &sub {
-        filter = filter.with_sub_claim(sub);
+        auth_session_filter = auth_session_filter.with_sub_claim(sub);
     }
     if let Some(sid) = &sid {
-        filter = filter.with_sid_claim(sid);
+        auth_session_filter = auth_session_filter.with_sid_claim(sid);
     }
+    let count = repo
+        .upstream_oauth_session()
+        .count(auth_session_filter)
+        .await?;
 
-    // Load the corresponding authentication sessions, by batches of 100s. It's
-    // VERY unlikely that we'll ever have more that 100 sessions for a single
-    // logout notification, but we'll handle it anyway.
-    let mut cursor = Pagination::first(100);
-    let mut sessions = Vec::new();
-    loop {
-        let page = repo.upstream_oauth_session().list(filter, cursor).await?;
-
-        for session in page.edges {
-            cursor = cursor.after(session.id);
-            sessions.push(session);
-        }
-
-        if !page.has_next_page {
-            break;
-        }
-    }
-
-    tracing::info!(sub, sid, %provider.id, "Backchannel logout received, found {} corresponding authentication sessions", sessions.len());
+    tracing::info!(sub, sid, %provider.id, "Backchannel logout received, found {count} corresponding authentication sessions");
 
     match provider.on_backchannel_logout {
         UpstreamOAuthProviderOnBackchannelLogout::DoNothing => {
             tracing::warn!(%provider.id, "Provider configured to do nothing on backchannel logout");
         }
         UpstreamOAuthProviderOnBackchannelLogout::LogoutBrowserOnly => {
-            let filter =
-                BrowserSessionFilter::new().authenticated_by_upstream_sessions_only(&sessions);
+            let filter = BrowserSessionFilter::new()
+                .authenticated_by_upstream_sessions_only(auth_session_filter);
             let affected = repo.browser_session().finish_bulk(&clock, filter).await?;
             tracing::info!("Finished {affected} browser sessions");
         }

crates/storage-pg/src/user/session.rs

@@ -17,7 +17,7 @@ use mas_storage::{
     user::{BrowserSessionFilter, BrowserSessionRepository},
 };
 use rand::RngCore;
-use sea_query::{Expr, PgFunc, PostgresQueryBuilder, Query};
+use sea_query::{Expr, PostgresQueryBuilder, Query};
 use sea_query_binder::SqlxBinder;
 use sqlx::PgConnection;
 use ulid::Ulid;
@@ -26,7 +26,7 @@ use uuid::Uuid;
 use crate::{
     DatabaseError, DatabaseInconsistencyError,
     filter::StatementExt,
-    iden::{UserSessionAuthentications, UserSessions, Users},
+    iden::{UpstreamOAuthAuthorizationSessions, UserSessionAuthentications, UserSessions, Users},
     pagination::QueryBuilderExt,
     tracing::ExecuteExt,
 };
@@ -145,13 +145,17 @@ impl crate::filter::Filter for BrowserSessionFilter<'_> {
             .add_option(self.last_active_before().map(|last_active_before| {
                 Expr::col((UserSessions::Table, UserSessions::LastActiveAt)).lt(last_active_before)
             }))
-            .add_option(self.authenticated_by_upstream_sessions().map(|sessions| {
+            .add_option(self.authenticated_by_upstream_sessions().map(|filter| {
                 // For filtering by upstream sessions, we need to hop over the
                 // `user_session_authentications` table
-                let session_ids: Vec<_> = sessions
-                    .iter()
-                    .map(|session| Uuid::from(session.id))
-                    .collect();
+                let join_expr = Expr::col((
+                    UserSessionAuthentications::Table,
+                    UserSessionAuthentications::UpstreamOAuthAuthorizationSessionId,
+                ))
+                .eq(Expr::col((
+                    UpstreamOAuthAuthorizationSessions::Table,
+                    UpstreamOAuthAuthorizationSessions::UpstreamOAuthAuthorizationSessionId,
+                )));
 
                 Expr::col((UserSessions::Table, UserSessions::UserSessionId)).in_subquery(
                     Query::select()
@@ -160,13 +164,8 @@ impl crate::filter::Filter for BrowserSessionFilter<'_> {
                             UserSessionAuthentications::UserSessionId,
                         )))
                         .from(UserSessionAuthentications::Table)
-                        .and_where(
-                            Expr::col((
-                                UserSessionAuthentications::Table,
-                                UserSessionAuthentications::UpstreamOAuthAuthorizationSessionId,
-                            ))
-                            .eq(PgFunc::any(Expr::value(session_ids))),
-                        )
+                        .inner_join(UpstreamOAuthAuthorizationSessions::Table, join_expr)
+                        .apply_filter(filter)
                         .take(),
                 )
             }))

crates/storage-pg/src/user/tests.rs

@@ -9,7 +9,7 @@ use mas_iana::jose::JsonWebSignatureAlg;
 use mas_storage::{
     Clock, Pagination, RepositoryAccess,
     clock::MockClock,
-    upstream_oauth2::UpstreamOAuthProviderParams,
+    upstream_oauth2::{UpstreamOAuthProviderParams, UpstreamOAuthSessionFilter},
     user::{
         BrowserSessionFilter, BrowserSessionRepository, UserEmailFilter, UserEmailRepository,
         UserFilter, UserPasswordRepository, UserRepository,
@@ -779,8 +779,11 @@ async fn test_user_session(pool: PgPool) {
         .await
         .unwrap();
 
-    let session_list = vec![upstream_oauth_session];
-    let filter = BrowserSessionFilter::new().authenticated_by_upstream_sessions_only(&session_list);
+    // This will match all authorization sessions, which matches exactly that one
+    // authorization session
+    let upstream_oauth_session_filter = UpstreamOAuthSessionFilter::new();
+    let filter = BrowserSessionFilter::new()
+        .authenticated_by_upstream_sessions_only(upstream_oauth_session_filter);
 
     // Now try to look it up
     let page = repo

crates/storage/src/user/session.rs

@@ -14,7 +14,10 @@ use mas_data_model::{
 use rand_core::RngCore;
 use ulid::Ulid;
 
-use crate::{Clock, Pagination, pagination::Page, repository_impl};
+use crate::{
+    Clock, Pagination, pagination::Page, repository_impl,
+    upstream_oauth2::UpstreamOAuthSessionFilter,
+};
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum BrowserSessionState {
@@ -39,7 +42,7 @@ pub struct BrowserSessionFilter<'a> {
     state: Option<BrowserSessionState>,
     last_active_before: Option<DateTime<Utc>>,
     last_active_after: Option<DateTime<Utc>>,
-    authenticated_by_upstream_sessions: Option<&'a [UpstreamOAuthAuthorizationSession]>,
+    authenticated_by_upstream_sessions: Option<UpstreamOAuthSessionFilter<'a>>,
 }
 
 impl<'a> BrowserSessionFilter<'a> {
@@ -117,17 +120,15 @@ impl<'a> BrowserSessionFilter<'a> {
     #[must_use]
     pub fn authenticated_by_upstream_sessions_only(
         mut self,
-        upstream_oauth_sessions: &'a [UpstreamOAuthAuthorizationSession],
+        filter: UpstreamOAuthSessionFilter<'a>,
     ) -> Self {
-        self.authenticated_by_upstream_sessions = Some(upstream_oauth_sessions);
+        self.authenticated_by_upstream_sessions = Some(filter);
         self
     }
 
     /// Get the upstream OAuth session filter
     #[must_use]
-    pub fn authenticated_by_upstream_sessions(
-        &self,
-    ) -> Option<&'a [UpstreamOAuthAuthorizationSession]> {
+    pub fn authenticated_by_upstream_sessions(&self) -> Option<UpstreamOAuthSessionFilter<'a>> {
         self.authenticated_by_upstream_sessions
     }
 }