storage: allow filtering browser sessions by which upstream session

authd them

Author: sandhose

crates/storage-pg/src/iden.rs

@@ -18,6 +18,18 @@ pub enum UserSessions {
     LastActiveIp,
 }
 
+#[derive(sea_query::Iden)]
+#[expect(dead_code)]
+pub enum UserSessionAuthentications {
+    Table,
+    UserSessionAuthenticationId,
+    UserSessionId,
+    CreatedAt,
+    UserPasswordId,
+    #[iden = "upstream_oauth_authorization_session_id"]
+    UpstreamOAuthAuthorizationSessionId,
+}
+
 #[derive(sea_query::Iden)]
 pub enum Users {
     Table,

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

@@ -17,7 +17,7 @@ use mas_storage::{
     user::{BrowserSessionFilter, BrowserSessionRepository},
 };
 use rand::RngCore;
-use sea_query::{Expr, PostgresQueryBuilder};
+use sea_query::{Expr, PgFunc, 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::{UserSessions, Users},
+    iden::{UserSessionAuthentications, UserSessions, Users},
     pagination::QueryBuilderExt,
     tracing::ExecuteExt,
 };
@@ -145,6 +145,31 @@ 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| {
+                // 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();
+
+                Expr::col((UserSessions::Table, UserSessions::UserSessionId)).in_subquery(
+                    Query::select()
+                        .expr(Expr::col((
+                            UserSessionAuthentications::Table,
+                            UserSessionAuthentications::UserSessionId,
+                        )))
+                        .from(UserSessionAuthentications::Table)
+                        .and_where(
+                            Expr::col((
+                                UserSessionAuthentications::Table,
+                                UserSessionAuthentications::UpstreamOAuthAuthorizationSessionId,
+                            ))
+                            .eq(PgFunc::any(Expr::value(session_ids))),
+                        )
+                        .take(),
+                )
+            }))
     }
 }
 

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

@@ -5,14 +5,17 @@
 // Please see LICENSE files in the repository root for full details.
 
 use chrono::Duration;
+use mas_iana::jose::JsonWebSignatureAlg;
 use mas_storage::{
     Clock, Pagination, RepositoryAccess,
     clock::MockClock,
+    upstream_oauth2::UpstreamOAuthProviderParams,
     user::{
         BrowserSessionFilter, BrowserSessionRepository, UserEmailFilter, UserEmailRepository,
         UserFilter, UserPasswordRepository, UserRepository,
     },
 };
+use oauth2_types::scope::{OPENID, Scope};
 use rand::SeedableRng;
 use rand_chacha::ChaChaRng;
 use sqlx::PgPool;
@@ -717,6 +720,97 @@ async fn test_user_session(pool: PgPool) {
     assert_eq!(repo.browser_session().count(all_bob).await.unwrap(), 5);
     assert_eq!(repo.browser_session().count(active_bob).await.unwrap(), 0);
     assert_eq!(repo.browser_session().count(finished).await.unwrap(), 11);
+
+    // Checking the 'authenticaated by upstream sessions' filter
+    // We need a provider
+    let provider = repo
+        .upstream_oauth_provider()
+        .add(
+            &mut rng,
+            &clock,
+            UpstreamOAuthProviderParams {
+                issuer: None,
+                human_name: None,
+                brand_name: None,
+                scope: Scope::from_iter([OPENID]),
+                token_endpoint_auth_method:
+                    mas_data_model::UpstreamOAuthProviderTokenAuthMethod::None,
+                token_endpoint_signing_alg: None,
+                id_token_signed_response_alg: JsonWebSignatureAlg::Rs256,
+                fetch_userinfo: false,
+                userinfo_signed_response_alg: None,
+                client_id: "client".to_owned(),
+                encrypted_client_secret: None,
+                claims_imports: mas_data_model::UpstreamOAuthProviderClaimsImports::default(),
+                authorization_endpoint_override: None,
+                token_endpoint_override: None,
+                userinfo_endpoint_override: None,
+                jwks_uri_override: None,
+                discovery_mode: mas_data_model::UpstreamOAuthProviderDiscoveryMode::Disabled,
+                pkce_mode: mas_data_model::UpstreamOAuthProviderPkceMode::Disabled,
+                response_mode: None,
+                additional_authorization_parameters: Vec::new(),
+                forward_login_hint: false,
+                ui_order: 0,
+                on_backchannel_logout:
+                    mas_data_model::UpstreamOAuthProviderOnBackchannelLogout::DoNothing,
+            },
+        )
+        .await
+        .unwrap();
+
+    // Start a authorization session
+    let upstream_oauth_session = repo
+        .upstream_oauth_session()
+        .add(&mut rng, &clock, &provider, "state".to_owned(), None, None)
+        .await
+        .unwrap();
+
+    // Start a browser session
+    let session = repo
+        .browser_session()
+        .add(&mut rng, &clock, &alice, None)
+        .await
+        .unwrap();
+
+    // Make the session from alice authenticated by this session
+    repo.browser_session()
+        .authenticate_with_upstream(&mut rng, &clock, &session, &upstream_oauth_session)
+        .await
+        .unwrap();
+
+    let session_list = vec![upstream_oauth_session];
+    let filter = BrowserSessionFilter::new().authenticated_by_upstream_sessions_only(&session_list);
+
+    // Now try to look it up
+    let page = repo
+        .browser_session()
+        .list(filter, Pagination::first(10))
+        .await
+        .unwrap();
+    assert_eq!(page.edges.len(), 1);
+    assert_eq!(page.edges[0].id, session.id);
+
+    // Try counting
+    assert_eq!(repo.browser_session().count(filter).await.unwrap(), 1);
+
+    // Try finishing the session
+    let affected = repo
+        .browser_session()
+        .finish_bulk(&clock, filter)
+        .await
+        .unwrap();
+    assert_eq!(affected, 1);
+
+    // Lookup the session by its ID
+    let lookup = repo
+        .browser_session()
+        .lookup(session.id)
+        .await
+        .unwrap()
+        .expect("session to be found in the database");
+    // It should be finished
+    assert!(lookup.finished_at.is_some());
 }
 
 #[sqlx::test(migrator = "crate::MIGRATOR")]

crates/storage/src/user/session.rs

@@ -39,6 +39,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]>,
 }
 
 impl<'a> BrowserSessionFilter<'a> {
@@ -110,6 +111,25 @@ impl<'a> BrowserSessionFilter<'a> {
     pub fn state(&self) -> Option<BrowserSessionState> {
         self.state
     }
+
+    /// Only return browser sessions authenticated by the given upstream OAuth
+    /// sessions
+    #[must_use]
+    pub fn authenticated_by_upstream_sessions_only(
+        mut self,
+        upstream_oauth_sessions: &'a [UpstreamOAuthAuthorizationSession],
+    ) -> Self {
+        self.authenticated_by_upstream_sessions = Some(upstream_oauth_sessions);
+        self
+    }
+
+    /// Get the upstream OAuth session filter
+    #[must_use]
+    pub fn authenticated_by_upstream_sessions(
+        &self,
+    ) -> Option<&'a [UpstreamOAuthAuthorizationSession]> {
+        self.authenticated_by_upstream_sessions
+    }
 }
 
 /// A [`BrowserSessionRepository`] helps interacting with [`BrowserSession`]