diff --git a/crates/storage-pg/src/iden.rs b/crates/storage-pg/src/iden.rs index d401ca144..f7342adf5 100644 --- a/crates/storage-pg/src/iden.rs +++ b/crates/storage-pg/src/iden.rs @@ -140,6 +140,29 @@ pub enum UpstreamOAuthLinks { CreatedAt, } +#[derive(sea_query::Iden)] +#[iden = "upstream_oauth_authorization_sessions"] +pub enum UpstreamOAuthAuthorizationSessions { + Table, + #[iden = "upstream_oauth_authorization_session_id"] + UpstreamOAuthAuthorizationSessionId, + #[iden = "upstream_oauth_provider_id"] + UpstreamOAuthProviderId, + #[iden = "upstream_oauth_link_id"] + UpstreamOAuthLinkId, + State, + CodeChallengeVerifier, + Nonce, + IdToken, + IdTokenClaims, + ExtraCallbackParameters, + Userinfo, + CreatedAt, + CompletedAt, + ConsumedAt, + UnlinkedAt, +} + #[derive(sea_query::Iden)] pub enum UserRegistrationTokens { Table, diff --git a/crates/storage-pg/src/upstream_oauth2/mod.rs b/crates/storage-pg/src/upstream_oauth2/mod.rs index b77e0631d..2158d872c 100644 --- a/crates/storage-pg/src/upstream_oauth2/mod.rs +++ b/crates/storage-pg/src/upstream_oauth2/mod.rs @@ -29,7 +29,7 @@ mod tests { upstream_oauth2::{ UpstreamOAuthLinkFilter, UpstreamOAuthLinkRepository, UpstreamOAuthProviderFilter, UpstreamOAuthProviderParams, UpstreamOAuthProviderRepository, - UpstreamOAuthSessionRepository, + UpstreamOAuthSessionFilter, UpstreamOAuthSessionRepository, }, user::UserRepository, }; @@ -262,6 +262,29 @@ mod tests { 1 ); + // Test listing and counting sessions + let session_filter = UpstreamOAuthSessionFilter::new().for_provider(&provider); + + // Count the sessions for the provider + let session_count = repo + .upstream_oauth_session() + .count(session_filter) + .await + .unwrap(); + assert_eq!(session_count, 1); + + // List the sessions for the provider + let session_page = repo + .upstream_oauth_session() + .list(session_filter, Pagination::first(10)) + .await + .unwrap(); + + assert_eq!(session_page.edges.len(), 1); + assert_eq!(session_page.edges[0].id, session.id); + assert!(!session_page.has_next_page); + assert!(!session_page.has_previous_page); + // Try deleting the provider repo.upstream_oauth_provider() .delete(provider) @@ -423,4 +446,136 @@ mod tests { .is_empty() ); } + + /// Test that the pagination works as expected in the upstream OAuth + /// session repository + #[sqlx::test(migrator = "crate::MIGRATOR")] + async fn test_session_repository_pagination(pool: PgPool) { + let scope = Scope::from_iter([OPENID]); + + let mut rng = rand_chacha::ChaChaRng::seed_from_u64(42); + let clock = MockClock::default(); + let mut repo = PgRepository::from_pool(&pool).await.unwrap(); + + // Create a provider + let provider = repo + .upstream_oauth_provider() + .add( + &mut rng, + &clock, + UpstreamOAuthProviderParams { + issuer: Some("https://example.com/".to_owned()), + human_name: None, + brand_name: None, + scope, + token_endpoint_auth_method: UpstreamOAuthProviderTokenAuthMethod::None, + id_token_signed_response_alg: JsonWebSignatureAlg::Rs256, + fetch_userinfo: false, + userinfo_signed_response_alg: None, + token_endpoint_signing_alg: None, + client_id: "client-id".to_owned(), + encrypted_client_secret: None, + claims_imports: UpstreamOAuthProviderClaimsImports::default(), + token_endpoint_override: None, + authorization_endpoint_override: None, + userinfo_endpoint_override: None, + jwks_uri_override: None, + discovery_mode: mas_data_model::UpstreamOAuthProviderDiscoveryMode::Oidc, + pkce_mode: mas_data_model::UpstreamOAuthProviderPkceMode::Auto, + response_mode: None, + additional_authorization_parameters: Vec::new(), + forward_login_hint: false, + ui_order: 0, + }, + ) + .await + .unwrap(); + + let filter = UpstreamOAuthSessionFilter::new().for_provider(&provider); + + // Count the number of sessions before we start + assert_eq!( + repo.upstream_oauth_session().count(filter).await.unwrap(), + 0 + ); + + let mut ids = Vec::with_capacity(20); + // Create 20 sessions + for idx in 0..20 { + let state = format!("state-{idx}"); + let session = repo + .upstream_oauth_session() + .add(&mut rng, &clock, &provider, state, None, None) + .await + .unwrap(); + ids.push(session.id); + clock.advance(Duration::microseconds(10 * 1000 * 1000)); + } + + // Now we have 20 sessions + assert_eq!( + repo.upstream_oauth_session().count(filter).await.unwrap(), + 20 + ); + + // Lookup the first 10 items + let page = repo + .upstream_oauth_session() + .list(filter, Pagination::first(10)) + .await + .unwrap(); + + // It returned the first 10 items + assert!(page.has_next_page); + let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect(); + assert_eq!(&edge_ids, &ids[..10]); + + // Lookup the next 10 items + let page = repo + .upstream_oauth_session() + .list(filter, Pagination::first(10).after(ids[9])) + .await + .unwrap(); + + // It returned the next 10 items + assert!(!page.has_next_page); + let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect(); + assert_eq!(&edge_ids, &ids[10..]); + + // Lookup the last 10 items + let page = repo + .upstream_oauth_session() + .list(filter, Pagination::last(10)) + .await + .unwrap(); + + // It returned the last 10 items + assert!(page.has_previous_page); + let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect(); + assert_eq!(&edge_ids, &ids[10..]); + + // Lookup the previous 10 items + let page = repo + .upstream_oauth_session() + .list(filter, Pagination::last(10).before(ids[10])) + .await + .unwrap(); + + // It returned the previous 10 items + assert!(!page.has_previous_page); + let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect(); + assert_eq!(&edge_ids, &ids[..10]); + + // Lookup 5 items between two IDs + let page = repo + .upstream_oauth_session() + .list(filter, Pagination::first(10).after(ids[5]).before(ids[11])) + .await + .unwrap(); + + // It returned the items in between + assert!(!page.has_next_page); + let edge_ids: Vec<_> = page.edges.iter().map(|s| s.id).collect(); + assert_eq!(&edge_ids, &ids[6..11]); + } } diff --git a/crates/storage-pg/src/upstream_oauth2/session.rs b/crates/storage-pg/src/upstream_oauth2/session.rs index 7946d63f0..a595f600d 100644 --- a/crates/storage-pg/src/upstream_oauth2/session.rs +++ b/crates/storage-pg/src/upstream_oauth2/session.rs @@ -10,13 +10,36 @@ use mas_data_model::{ UpstreamOAuthAuthorizationSession, UpstreamOAuthAuthorizationSessionState, UpstreamOAuthLink, UpstreamOAuthProvider, }; -use mas_storage::{Clock, upstream_oauth2::UpstreamOAuthSessionRepository}; +use mas_storage::{ + Clock, Page, Pagination, + upstream_oauth2::{UpstreamOAuthSessionFilter, UpstreamOAuthSessionRepository}, +}; use rand::RngCore; +use sea_query::{Expr, PostgresQueryBuilder, Query, enum_def}; +use sea_query_binder::SqlxBinder; use sqlx::PgConnection; use ulid::Ulid; use uuid::Uuid; -use crate::{DatabaseError, DatabaseInconsistencyError, tracing::ExecuteExt}; +use crate::{ + DatabaseError, DatabaseInconsistencyError, + filter::{Filter, StatementExt}, + iden::UpstreamOAuthAuthorizationSessions, + pagination::QueryBuilderExt, + tracing::ExecuteExt, +}; + +impl Filter for UpstreamOAuthSessionFilter<'_> { + fn generate_condition(&self, _has_joins: bool) -> impl sea_query::IntoCondition { + sea_query::Condition::all().add_option(self.provider().map(|provider| { + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::UpstreamOAuthProviderId, + )) + .eq(Uuid::from(provider.id)) + })) + } +} /// An implementation of [`UpstreamOAuthSessionRepository`] for a PostgreSQL /// connection @@ -32,6 +55,8 @@ impl<'c> PgUpstreamOAuthSessionRepository<'c> { } } +#[derive(sqlx::FromRow)] +#[enum_def] struct SessionLookup { upstream_oauth_authorization_session_id: Uuid, upstream_oauth_provider_id: Uuid, @@ -346,4 +371,173 @@ impl UpstreamOAuthSessionRepository for PgUpstreamOAuthSessionRepository<'_> { Ok(upstream_oauth_authorization_session) } + + #[tracing::instrument( + name = "db.upstream_oauth_authorization_session.list", + skip_all, + fields( + db.query.text, + ), + err, + )] + async fn list( + &mut self, + filter: UpstreamOAuthSessionFilter<'_>, + pagination: Pagination, + ) -> Result, Self::Error> { + let (sql, arguments) = Query::select() + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::UpstreamOAuthAuthorizationSessionId, + )), + SessionLookupIden::UpstreamOauthAuthorizationSessionId, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::UpstreamOAuthProviderId, + )), + SessionLookupIden::UpstreamOauthProviderId, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::UpstreamOAuthLinkId, + )), + SessionLookupIden::UpstreamOauthLinkId, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::State, + )), + SessionLookupIden::State, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::CodeChallengeVerifier, + )), + SessionLookupIden::CodeChallengeVerifier, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::Nonce, + )), + SessionLookupIden::Nonce, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::IdToken, + )), + SessionLookupIden::IdToken, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::IdTokenClaims, + )), + SessionLookupIden::IdTokenClaims, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::ExtraCallbackParameters, + )), + SessionLookupIden::ExtraCallbackParameters, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::Userinfo, + )), + SessionLookupIden::Userinfo, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::CreatedAt, + )), + SessionLookupIden::CreatedAt, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::CompletedAt, + )), + SessionLookupIden::CompletedAt, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::ConsumedAt, + )), + SessionLookupIden::ConsumedAt, + ) + .expr_as( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::UnlinkedAt, + )), + SessionLookupIden::UnlinkedAt, + ) + .from(UpstreamOAuthAuthorizationSessions::Table) + .apply_filter(filter) + .generate_pagination( + ( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::UpstreamOAuthAuthorizationSessionId, + ), + pagination, + ) + .build_sqlx(PostgresQueryBuilder); + + let edges: Vec = sqlx::query_as_with(&sql, arguments) + .traced() + .fetch_all(&mut *self.conn) + .await?; + + let page = pagination + .process(edges) + .try_map(UpstreamOAuthAuthorizationSession::try_from)?; + + Ok(page) + } + + #[tracing::instrument( + name = "db.upstream_oauth_authorization_session.count", + skip_all, + fields( + db.query.text, + ), + err, + )] + async fn count( + &mut self, + filter: UpstreamOAuthSessionFilter<'_>, + ) -> Result { + let (sql, arguments) = Query::select() + .expr( + Expr::col(( + UpstreamOAuthAuthorizationSessions::Table, + UpstreamOAuthAuthorizationSessions::UpstreamOAuthAuthorizationSessionId, + )) + .count(), + ) + .from(UpstreamOAuthAuthorizationSessions::Table) + .apply_filter(filter) + .build_sqlx(PostgresQueryBuilder); + + let count: i64 = sqlx::query_scalar_with(&sql, arguments) + .traced() + .fetch_one(&mut *self.conn) + .await?; + + count + .try_into() + .map_err(DatabaseError::to_invalid_operation) + } } diff --git a/crates/storage/src/upstream_oauth2/mod.rs b/crates/storage/src/upstream_oauth2/mod.rs index 3fd6666b4..39fefffe8 100644 --- a/crates/storage/src/upstream_oauth2/mod.rs +++ b/crates/storage/src/upstream_oauth2/mod.rs @@ -16,5 +16,5 @@ pub use self::{ provider::{ UpstreamOAuthProviderFilter, UpstreamOAuthProviderParams, UpstreamOAuthProviderRepository, }, - session::UpstreamOAuthSessionRepository, + session::{UpstreamOAuthSessionFilter, UpstreamOAuthSessionRepository}, }; diff --git a/crates/storage/src/upstream_oauth2/session.rs b/crates/storage/src/upstream_oauth2/session.rs index 273370794..66fcb1ba8 100644 --- a/crates/storage/src/upstream_oauth2/session.rs +++ b/crates/storage/src/upstream_oauth2/session.rs @@ -9,7 +9,36 @@ use mas_data_model::{UpstreamOAuthAuthorizationSession, UpstreamOAuthLink, Upstr use rand_core::RngCore; use ulid::Ulid; -use crate::{Clock, repository_impl}; +use crate::{Clock, Pagination, pagination::Page, repository_impl}; + +/// Filter parameters for listing upstream OAuth sessions +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub struct UpstreamOAuthSessionFilter<'a> { + provider: Option<&'a UpstreamOAuthProvider>, +} + +impl<'a> UpstreamOAuthSessionFilter<'a> { + /// Create a new [`UpstreamOAuthSessionFilter`] with default values + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Set the upstream OAuth provider for which to list sessions + #[must_use] + pub fn for_provider(mut self, provider: &'a UpstreamOAuthProvider) -> Self { + self.provider = Some(provider); + self + } + + /// Get the upstream OAuth provider filter + /// + /// Returns [`None`] if no filter was set + #[must_use] + pub fn provider(&self) -> Option<&UpstreamOAuthProvider> { + self.provider + } +} /// An [`UpstreamOAuthSessionRepository`] helps interacting with /// [`UpstreamOAuthAuthorizationSession`] saved in the storage backend @@ -112,6 +141,36 @@ pub trait UpstreamOAuthSessionRepository: Send + Sync { clock: &dyn Clock, upstream_oauth_authorization_session: UpstreamOAuthAuthorizationSession, ) -> Result; + + /// List [`UpstreamOAuthAuthorizationSession`] with the given filter and + /// pagination + /// + /// # Parameters + /// + /// * `filter`: The filter to apply + /// * `pagination`: The pagination parameters + /// + /// # Errors + /// + /// Returns [`Self::Error`] if the underlying repository fails + async fn list( + &mut self, + filter: UpstreamOAuthSessionFilter<'_>, + pagination: Pagination, + ) -> Result, Self::Error>; + + /// Count the number of [`UpstreamOAuthAuthorizationSession`] with the given + /// filter + /// + /// # Parameters + /// + /// * `filter`: The filter to apply + /// + /// # Errors + /// + /// Returns [`Self::Error`] if the underlying repository fails + async fn count(&mut self, filter: UpstreamOAuthSessionFilter<'_>) + -> Result; } repository_impl!(UpstreamOAuthSessionRepository: @@ -146,4 +205,12 @@ repository_impl!(UpstreamOAuthSessionRepository: clock: &dyn Clock, upstream_oauth_authorization_session: UpstreamOAuthAuthorizationSession, ) -> Result; + + async fn list( + &mut self, + filter: UpstreamOAuthSessionFilter<'_>, + pagination: Pagination, + ) -> Result, Self::Error>; + + async fn count(&mut self, filter: UpstreamOAuthSessionFilter<'_>) -> Result; );