diff --git a/crates/graphql/src/model/compat_sessions.rs b/crates/graphql/src/model/compat_sessions.rs index 441efd67..25e147ef 100644 --- a/crates/graphql/src/model/compat_sessions.rs +++ b/crates/graphql/src/model/compat_sessions.rs @@ -21,13 +21,41 @@ use url::Url; use super::{NodeType, User}; use crate::state::ContextExt; +/// Lazy-loaded reverse reference. +/// +/// XXX: maybe we want to stick that in a utility module +#[derive(Clone, Debug, Default)] +enum ReverseReference { + Loaded(T), + #[default] + Lazy, +} + /// A compat session represents a client session which used the legacy Matrix /// login API. #[derive(Description)] -pub struct CompatSession( - pub mas_data_model::CompatSession, - pub Option, -); +pub struct CompatSession { + session: mas_data_model::CompatSession, + sso_login: ReverseReference>, +} + +impl CompatSession { + pub fn new(session: mas_data_model::CompatSession) -> Self { + Self { + session, + sso_login: ReverseReference::Lazy, + } + } + + /// Save an eagerly loaded SSO login. + pub fn with_loaded_sso_login( + mut self, + sso_login: Option, + ) -> Self { + self.sso_login = ReverseReference::Loaded(sso_login); + self + } +} /// The state of a compatibility session. #[derive(Enum, Copy, Clone, Eq, PartialEq)] @@ -53,7 +81,7 @@ pub enum CompatSessionType { impl CompatSession { /// ID of the object. pub async fn id(&self) -> ID { - NodeType::CompatSession.id(self.0.id) + NodeType::CompatSession.id(self.session.id) } /// The user authorized for this session. @@ -62,7 +90,7 @@ impl CompatSession { let mut repo = state.repository().await?; let user = repo .user() - .lookup(self.0.user_id) + .lookup(self.session.user_id) .await? .context("Could not load user")?; repo.cancel().await?; @@ -72,27 +100,44 @@ impl CompatSession { /// The Matrix Device ID of this session. async fn device_id(&self) -> &str { - self.0.device.as_str() + self.session.device.as_str() } /// When the object was created. pub async fn created_at(&self) -> DateTime { - self.0.created_at + self.session.created_at } /// When the session ended. pub async fn finished_at(&self) -> Option> { - self.0.finished_at() + self.session.finished_at() } /// The associated SSO login, if any. - pub async fn sso_login(&self) -> Option { - self.1.as_ref().map(|l| CompatSsoLogin(l.clone())) + pub async fn sso_login( + &self, + ctx: &Context<'_>, + ) -> Result, async_graphql::Error> { + if let ReverseReference::Loaded(sso_login) = &self.sso_login { + return Ok(sso_login.clone().map(CompatSsoLogin)); + } + + // We need to load it on the fly + let state = ctx.state(); + let mut repo = state.repository().await?; + let sso_login = repo + .compat_sso_login() + .find_for_session(&self.session) + .await + .context("Could not load SSO login")?; + repo.cancel().await?; + + Ok(sso_login.map(CompatSsoLogin)) } /// The state of the session. pub async fn state(&self) -> CompatSessionState { - match &self.0.state { + match &self.session.state { mas_data_model::CompatSessionState::Valid => CompatSessionState::Active, mas_data_model::CompatSessionState::Finished { .. } => CompatSessionState::Finished, } @@ -150,6 +195,8 @@ impl CompatSsoLogin { .context("Could not load compat session")?; repo.cancel().await?; - Ok(Some(CompatSession(session, Some(self.0.clone())))) + Ok(Some( + CompatSession::new(session).with_loaded_sso_login(Some(self.0.clone())), + )) } } diff --git a/crates/graphql/src/model/users.rs b/crates/graphql/src/model/users.rs index 7fbfdd8d..780d8b0f 100644 --- a/crates/graphql/src/model/users.rs +++ b/crates/graphql/src/model/users.rs @@ -213,7 +213,7 @@ impl User { .extend(page.edges.into_iter().map(|(session, sso_login)| { Edge::new( OpaqueCursor(NodeCursor(NodeType::CompatSession, session.id)), - CompatSession(session, sso_login), + CompatSession::new(session).with_loaded_sso_login(sso_login), ) })); diff --git a/crates/graphql/src/mutations/compat_session.rs b/crates/graphql/src/mutations/compat_session.rs index 75dfef5f..877a08ff 100644 --- a/crates/graphql/src/mutations/compat_session.rs +++ b/crates/graphql/src/mutations/compat_session.rs @@ -66,8 +66,7 @@ impl EndCompatSessionPayload { /// Returns the ended session. async fn compat_session(&self) -> Option { match self { - // XXX: the SSO login is not returned here. - Self::Ended(session) => Some(CompatSession(session.clone(), None)), + Self::Ended(session) => Some(CompatSession::new(session.clone())), Self::NotFound => None, } } diff --git a/crates/graphql/src/query/session.rs b/crates/graphql/src/query/session.rs index 245b6e90..4eee094a 100644 --- a/crates/graphql/src/query/session.rs +++ b/crates/graphql/src/query/session.rs @@ -66,10 +66,8 @@ impl SessionQuery { if let Some(compat_session) = compat_session { repo.cancel().await?; - // XXX: we should load the compat SSO login as well - return Ok(Some(Session::CompatSession(Box::new(CompatSession( + return Ok(Some(Session::CompatSession(Box::new(CompatSession::new( compat_session, - None, ))))); } diff --git a/crates/storage-pg/.sqlx/query-1787a5e86b60f57295fe5111259a29ffb15aa31e707cb7f2ad4269d125f6d8c9.json b/crates/storage-pg/.sqlx/query-1787a5e86b60f57295fe5111259a29ffb15aa31e707cb7f2ad4269d125f6d8c9.json new file mode 100644 index 00000000..56b805e1 --- /dev/null +++ b/crates/storage-pg/.sqlx/query-1787a5e86b60f57295fe5111259a29ffb15aa31e707cb7f2ad4269d125f6d8c9.json @@ -0,0 +1,58 @@ +{ + "db_name": "PostgreSQL", + "query": "\n SELECT compat_sso_login_id\n , login_token\n , redirect_uri\n , created_at\n , fulfilled_at\n , exchanged_at\n , compat_session_id\n\n FROM compat_sso_logins\n WHERE compat_session_id = $1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "compat_sso_login_id", + "type_info": "Uuid" + }, + { + "ordinal": 1, + "name": "login_token", + "type_info": "Text" + }, + { + "ordinal": 2, + "name": "redirect_uri", + "type_info": "Text" + }, + { + "ordinal": 3, + "name": "created_at", + "type_info": "Timestamptz" + }, + { + "ordinal": 4, + "name": "fulfilled_at", + "type_info": "Timestamptz" + }, + { + "ordinal": 5, + "name": "exchanged_at", + "type_info": "Timestamptz" + }, + { + "ordinal": 6, + "name": "compat_session_id", + "type_info": "Uuid" + } + ], + "parameters": { + "Left": [ + "Uuid" + ] + }, + "nullable": [ + false, + false, + false, + false, + true, + true, + true + ] + }, + "hash": "1787a5e86b60f57295fe5111259a29ffb15aa31e707cb7f2ad4269d125f6d8c9" +} diff --git a/crates/storage-pg/src/compat/sso_login.rs b/crates/storage-pg/src/compat/sso_login.rs index fe5b3b46..ca7b4433 100644 --- a/crates/storage-pg/src/compat/sso_login.rs +++ b/crates/storage-pg/src/compat/sso_login.rs @@ -137,6 +137,44 @@ impl<'c> CompatSsoLoginRepository for PgCompatSsoLoginRepository<'c> { Ok(Some(res.try_into()?)) } + #[tracing::instrument( + name = "db.compat_sso_login.find_for_session", + skip_all, + fields( + db.statement, + %compat_session.id, + ), + err, + )] + async fn find_for_session( + &mut self, + compat_session: &CompatSession, + ) -> Result, Self::Error> { + let res = sqlx::query_as!( + CompatSsoLoginLookup, + r#" + SELECT compat_sso_login_id + , login_token + , redirect_uri + , created_at + , fulfilled_at + , exchanged_at + , compat_session_id + + FROM compat_sso_logins + WHERE compat_session_id = $1 + "#, + Uuid::from(compat_session.id), + ) + .traced() + .fetch_optional(&mut *self.conn) + .await?; + + let Some(res) = res else { return Ok(None) }; + + Ok(Some(res.try_into()?)) + } + #[tracing::instrument( name = "db.compat_sso_login.find_by_token", skip_all, diff --git a/crates/storage/src/compat/sso_login.rs b/crates/storage/src/compat/sso_login.rs index 0782f01b..282048c0 100644 --- a/crates/storage/src/compat/sso_login.rs +++ b/crates/storage/src/compat/sso_login.rs @@ -122,6 +122,22 @@ pub trait CompatSsoLoginRepository: Send + Sync { /// Returns [`Self::Error`] if the underlying repository fails async fn lookup(&mut self, id: Ulid) -> Result, Self::Error>; + /// Find a compat SSO login by its session + /// + /// Returns the compat SSO login if it exists, `None` otherwise + /// + /// # Parameters + /// + /// * `session`: The session of the compat SSO login to lookup + /// + /// # Errors + /// + /// Returns [`Self::Error`] if the underlying repository fails + async fn find_for_session( + &mut self, + session: &CompatSession, + ) -> Result, Self::Error>; + /// Find a compat SSO login by its login token /// /// Returns the compat SSO login if found, `None` otherwise @@ -232,6 +248,11 @@ pub trait CompatSsoLoginRepository: Send + Sync { repository_impl!(CompatSsoLoginRepository: async fn lookup(&mut self, id: Ulid) -> Result, Self::Error>; + async fn find_for_session( + &mut self, + session: &CompatSession, + ) -> Result, Self::Error>; + async fn find_by_token( &mut self, login_token: &str,