Save which user session created a compat session

This also exposes the user session in the GraphQL API, and allow
filtering on browser session ID on the app session list.

crates/cli/src/commands/manage.rs

@@ -184,7 +184,7 @@ impl Options {
 
                 let compat_session = repo
                     .compat_session()
-                    .add(&mut rng, &clock, &user, device, admin)
+                    .add(&mut rng, &clock, &user, device, None, admin)
                     .await?;
 
                 let token = TokenType::CompatAccessToken.generate(&mut rng);

crates/data-model/src/compat/session.rs

@@ -80,6 +80,7 @@ pub struct CompatSession {
     pub state: CompatSessionState,
     pub user_id: Ulid,
     pub device: Device,
+    pub user_session_id: Option<Ulid>,
     pub created_at: DateTime<Utc>,
     pub is_synapse_admin: bool,
     pub last_active_at: Option<DateTime<Utc>>,

crates/graphql/src/model/browser_sessions.rs

@@ -12,11 +12,20 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use async_graphql::{Context, Description, Object, ID};
+use async_graphql::{
+    connection::{query, Connection, Edge, OpaqueCursor},
+    Context, Description, Object, ID,
+};
 use chrono::{DateTime, Utc};
-use mas_storage::{user::BrowserSessionRepository, RepositoryAccess};
-
-use super::{NodeType, SessionState, User};
+use mas_data_model::Device;
+use mas_storage::{
+    app_session::AppSessionFilter, user::BrowserSessionRepository, Pagination, RepositoryAccess,
+};
+
+use super::{
+    AppSession, CompatSession, Cursor, NodeCursor, NodeType, OAuth2Session, PreloadedTotalCount,
+    SessionState, User,
+};
 use crate::state::ContextExt;
 
 /// A browser session represents a logged in user in a browser.
@@ -92,6 +101,97 @@ impl BrowserSession {
     pub async fn last_active_at(&self) -> Option<DateTime<Utc>> {
         self.0.last_active_at
     }
+
+    /// Get the list of both compat and OAuth 2.0 sessions started by this
+    /// browser session, chronologically sorted
+    #[allow(clippy::too_many_arguments)]
+    async fn app_sessions(
+        &self,
+        ctx: &Context<'_>,
+
+        #[graphql(name = "state", desc = "List only sessions in the given state.")]
+        state_param: Option<SessionState>,
+
+        #[graphql(name = "device", desc = "List only sessions for the given device.")]
+        device_param: Option<String>,
+
+        #[graphql(desc = "Returns the elements in the list that come after the cursor.")]
+        after: Option<String>,
+        #[graphql(desc = "Returns the elements in the list that come before the cursor.")]
+        before: Option<String>,
+        #[graphql(desc = "Returns the first *n* elements from the list.")] first: Option<i32>,
+        #[graphql(desc = "Returns the last *n* elements from the list.")] last: Option<i32>,
+    ) -> Result<Connection<Cursor, AppSession, PreloadedTotalCount>, async_graphql::Error> {
+        let state = ctx.state();
+        let mut repo = state.repository().await?;
+
+        query(
+            after,
+            before,
+            first,
+            last,
+            |after, before, first, last| async move {
+                let after_id = after
+                    .map(|x: OpaqueCursor<NodeCursor>| {
+                        x.extract_for_types(&[NodeType::OAuth2Session, NodeType::CompatSession])
+                    })
+                    .transpose()?;
+                let before_id = before
+                    .map(|x: OpaqueCursor<NodeCursor>| {
+                        x.extract_for_types(&[NodeType::OAuth2Session, NodeType::CompatSession])
+                    })
+                    .transpose()?;
+                let pagination = Pagination::try_new(before_id, after_id, first, last)?;
+
+                let device_param = device_param.map(Device::try_from).transpose()?;
+
+                let filter = AppSessionFilter::new().for_browser_session(&self.0);
+
+                let filter = match state_param {
+                    Some(SessionState::Active) => filter.active_only(),
+                    Some(SessionState::Finished) => filter.finished_only(),
+                    None => filter,
+                };
+
+                let filter = match device_param.as_ref() {
+                    Some(device) => filter.for_device(device),
+                    None => filter,
+                };
+
+                let page = repo.app_session().list(filter, pagination).await?;
+
+                let count = if ctx.look_ahead().field("totalCount").exists() {
+                    Some(repo.app_session().count(filter).await?)
+                } else {
+                    None
+                };
+
+                repo.cancel().await?;
+
+                let mut connection = Connection::with_additional_fields(
+                    page.has_previous_page,
+                    page.has_next_page,
+                    PreloadedTotalCount(count),
+                );
+
+                connection
+                    .edges
+                    .extend(page.edges.into_iter().map(|s| match s {
+                        mas_storage::app_session::AppSession::Compat(session) => Edge::new(
+                            OpaqueCursor(NodeCursor(NodeType::CompatSession, session.id)),
+                            AppSession::CompatSession(Box::new(CompatSession::new(*session))),
+                        ),
+                        mas_storage::app_session::AppSession::OAuth2(session) => Edge::new(
+                            OpaqueCursor(NodeCursor(NodeType::OAuth2Session, session.id)),
+                            AppSession::OAuth2Session(Box::new(OAuth2Session(*session))),
+                        ),
+                    }));
+
+                Ok::<_, async_graphql::Error>(connection)
+            },
+        )
+        .await
+    }
 }
 
 /// An authentication records when a user enter their credential in a browser

crates/graphql/src/model/compat_sessions.rs

@@ -18,7 +18,7 @@ use chrono::{DateTime, Utc};
 use mas_storage::{compat::CompatSessionRepository, user::UserRepository};
 use url::Url;
 
-use super::{NodeType, SessionState, User};
+use super::{BrowserSession, NodeType, SessionState, User};
 use crate::state::ContextExt;
 
 /// Lazy-loaded reverse reference.
@@ -125,6 +125,27 @@ impl CompatSession {
         Ok(sso_login.map(CompatSsoLogin))
     }
 
+    /// The browser session which started this session, if any.
+    pub async fn browser_session(
+        &self,
+        ctx: &Context<'_>,
+    ) -> Result<Option<BrowserSession>, async_graphql::Error> {
+        let Some(user_session_id) = self.session.user_session_id else {
+            return Ok(None);
+        };
+
+        let state = ctx.state();
+        let mut repo = state.repository().await?;
+        let browser_session = repo
+            .browser_session()
+            .lookup(user_session_id)
+            .await?
+            .context("Could not load browser session")?;
+        repo.cancel().await?;
+
+        Ok(Some(BrowserSession(browser_session)))
+    }
+
     /// The state of the session.
     pub async fn state(&self) -> SessionState {
         match &self.session.state {

crates/graphql/src/model/mod.rs

@@ -32,7 +32,7 @@ pub use self::{
     node::{Node, NodeType},
     oauth::{OAuth2Client, OAuth2Session},
     upstream_oauth::{UpstreamOAuth2Link, UpstreamOAuth2Provider},
-    users::{User, UserEmail},
+    users::{AppSession, User, UserEmail},
     viewer::{Anonymous, Viewer, ViewerSession},
 };
 

crates/graphql/src/model/users.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use anyhow::Context as _;
 use async_graphql::{
     connection::{query, Connection, Edge, OpaqueCursor},
     Context, Description, Enum, Object, Union, ID,
@@ -544,6 +545,12 @@ impl User {
         #[graphql(name = "device", desc = "List only sessions for the given device.")]
         device_param: Option<String>,
 
+        #[graphql(
+            name = "browserSession",
+            desc = "List only sessions for the given session."
+        )]
+        browser_session_param: Option<ID>,
+
         #[graphql(desc = "Returns the elements in the list that come after the cursor.")]
         after: Option<String>,
         #[graphql(desc = "Returns the elements in the list that come before the cursor.")]
@@ -552,6 +559,7 @@ impl User {
         #[graphql(desc = "Returns the last *n* elements from the list.")] last: Option<i32>,
     ) -> Result<Connection<Cursor, AppSession, PreloadedTotalCount>, async_graphql::Error> {
         let state = ctx.state();
+        let requester = ctx.requester();
         let mut repo = state.repository().await?;
 
         query(
@@ -587,6 +595,38 @@ impl User {
                     None => filter,
                 };
 
+                let maybe_session = match browser_session_param {
+                    Some(id) => {
+                        // This might fail, but we're probably alright with it
+                        let id = NodeType::BrowserSession
+                            .extract_ulid(&id)
+                            .context("Invalid browser_session parameter")?;
+
+                        let Some(session) = repo
+                            .browser_session()
+                            .lookup(id)
+                            .await?
+                            .filter(|u| requester.is_owner_or_admin(u))
+                        else {
+                            // If we couldn't find the session or if the requester can't access it,
+                            // return an empty list
+                            return Ok(Connection::with_additional_fields(
+                                false,
+                                false,
+                                PreloadedTotalCount(Some(0)),
+                            ));
+                        };
+
+                        Some(session)
+                    }
+                    None => None,
+                };
+
+                let filter = match maybe_session {
+                    Some(ref session) => filter.for_browser_session(session),
+                    None => filter,
+                };
+
                 let page = repo.app_session().list(filter, pagination).await?;
 
                 let count = if ctx.look_ahead().field("totalCount").exists() {
@@ -625,7 +665,7 @@ impl User {
 
 /// A session in an application, either a compatibility or an OAuth 2.0 one
 #[derive(Union)]
-enum AppSession {
+pub enum AppSession {
     CompatSession(Box<CompatSession>),
     OAuth2Session(Box<OAuth2Session>),
 }

crates/handlers/src/compat/login.rs

@@ -411,7 +411,7 @@ async fn user_password_login(
 
     let session = repo
         .compat_session()
-        .add(&mut rng, clock, &user, device, false)
+        .add(&mut rng, clock, &user, device, None, false)
         .await?;
 
     Ok((session, user))
@@ -738,7 +738,14 @@ mod tests {
         // Complete the flow by fulfilling it with a session
         let compat_session = repo
             .compat_session()
-            .add(&mut state.rng(), &state.clock, user, device.clone(), false)
+            .add(
+                &mut state.rng(),
+                &state.clock,
+                user,
+                device.clone(),
+                None,
+                false,
+            )
             .await
             .unwrap();
 

crates/handlers/src/compat/login_sso_complete.rs

@@ -208,7 +208,14 @@ pub async fn post(
 
     let compat_session = repo
         .compat_session()
-        .add(&mut rng, &clock, &session.user, device, false)
+        .add(
+            &mut rng,
+            &clock,
+            &session.user,
+            device,
+            Some(&session),
+            false,
+        )
         .await?;
 
     repo.compat_sso_login()