graphql: API to query client sessions out of a device_id and a user ID

crates/graphql/src/query/mod.rs

@@ -20,14 +20,15 @@ use crate::{
     UserId,
 };
 
+mod session;
 mod upstream_oauth;
 mod viewer;
 
-use self::{upstream_oauth::UpstreamOAuthQuery, viewer::ViewerQuery};
+use self::{session::SessionQuery, upstream_oauth::UpstreamOAuthQuery, viewer::ViewerQuery};
 
 /// The query root of the GraphQL interface.
 #[derive(Default, MergedObject)]
-pub struct Query(BaseQuery, UpstreamOAuthQuery, ViewerQuery);
+pub struct Query(BaseQuery, UpstreamOAuthQuery, SessionQuery, ViewerQuery);
 
 impl Query {
     #[must_use]

crates/graphql/src/query/session.rs

@@ -0,0 +1,106 @@
+// Copyright 2023 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use async_graphql::{Context, Object, Union, ID};
+use mas_data_model::Device;
+use mas_storage::{
+    compat::CompatSessionRepository, oauth2::OAuth2SessionFilter, Pagination, RepositoryAccess,
+};
+use oauth2_types::scope::Scope;
+
+use crate::{
+    model::{CompatSession, NodeType, OAuth2Session},
+    state::ContextExt,
+    UserId,
+};
+
+#[derive(Default)]
+pub struct SessionQuery;
+
+/// A client session, either compat or OAuth 2.0
+#[derive(Union)]
+enum Session {
+    CompatSession(Box<CompatSession>),
+    OAuth2Session(Box<OAuth2Session>),
+}
+
+#[Object]
+impl SessionQuery {
+    /// Lookup a compat or OAuth 2.0 session
+    async fn session(
+        &self,
+        ctx: &Context<'_>,
+        user_id: ID,
+        device_id: String,
+    ) -> Result<Option<Session>, async_graphql::Error> {
+        let user_id = NodeType::User.extract_ulid(&user_id)?;
+        let requester = ctx.requester();
+        if !requester.is_owner_or_admin(&UserId(user_id)) {
+            return Ok(None);
+        }
+
+        let Ok(device) = Device::try_from(device_id) else {
+            return Ok(None);
+        };
+        let state = ctx.state();
+        let mut repo = state.repository().await?;
+
+        // Lookup the user
+        let Some(user) = repo.user().lookup(user_id).await? else {
+            return Ok(None);
+        };
+
+        // First, try to find a compat session
+        let compat_session = repo.compat_session().find_by_device(&user, &device).await?;
+        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(
+                compat_session,
+                None,
+            )))));
+        }
+
+        // Then, try to find an OAuth 2.0 session. Because we don't have any dedicated
+        // device column, we're looking up using the device scope.
+        let scope = Scope::from_iter([device.to_scope_token()]);
+        let filter = OAuth2SessionFilter::new()
+            .for_user(&user)
+            .active_only()
+            .with_scope(&scope);
+        // We only want most recent session
+        let pagination = Pagination::last(1);
+        let sessions = repo.oauth2_session().list(filter, pagination).await?;
+
+        // It's technically possible to have multiple active OAuth 2.0 sessions. For
+        // now, we just log it if it is the case
+        if sessions.has_next_page {
+            // XXX: should we bail out?
+            tracing::warn!(
+                "Found more than one active session with device {device} for user {user_id}"
+            );
+        }
+
+        if let Some(session) = sessions.edges.into_iter().next() {
+            repo.cancel().await?;
+            return Ok(Some(Session::OAuth2Session(Box::new(OAuth2Session(
+                session,
+            )))));
+        }
+        repo.cancel().await?;
+
+        Ok(None)
+    }
+}

frontend/schema.graphql

@@ -702,6 +702,10 @@ type Query {
     last: Int
   ): UpstreamOAuth2ProviderConnection!
   """
+  Lookup a compat or OAuth 2.0 session
+  """
+  session(userId: ID!, deviceId: String!): Session
+  """
   Get the viewer
   """
   viewer: Viewer!
@@ -799,6 +803,11 @@ enum SendVerificationEmailStatus {
   ALREADY_VERIFIED
 }
 
+"""
+A client session, either compat or OAuth 2.0
+"""
+union Session = CompatSession | Oauth2Session
+
 """
 The input for the `addEmail` mutation
 """

frontend/src/gql/graphql.ts

@@ -505,6 +505,8 @@ export type Query = {
   node?: Maybe<Node>;
   /** Fetch an OAuth 2.0 client by its ID. */
   oauth2Client?: Maybe<Oauth2Client>;
+  /** Lookup a compat or OAuth 2.0 session */
+  session?: Maybe<Session>;
   /** Fetch an upstream OAuth 2.0 link by its ID. */
   upstreamOauth2Link?: Maybe<UpstreamOAuth2Link>;
   /** Fetch an upstream OAuth 2.0 provider by its ID. */
@@ -536,6 +538,12 @@ export type QueryOauth2ClientArgs = {
   id: Scalars["ID"]["input"];
 };
 
+/** The query root of the GraphQL interface. */
+export type QuerySessionArgs = {
+  deviceId: Scalars["String"]["input"];
+  userId: Scalars["ID"]["input"];
+};
+
 /** The query root of the GraphQL interface. */
 export type QueryUpstreamOauth2LinkArgs = {
   id: Scalars["ID"]["input"];
@@ -616,6 +624,9 @@ export enum SendVerificationEmailStatus {
   Sent = "SENT",
 }
 
+/** A client session, either compat or OAuth 2.0 */
+export type Session = CompatSession | Oauth2Session;
+
 /** The input for the `addEmail` mutation */
 export type SetDisplayNameInput = {
   /** The display name to set. If `None`, the display name will be removed. */

frontend/src/gql/schema.ts

@@ -1504,6 +1504,36 @@ export default {
               },
             ],
           },
+          {
+            name: "session",
+            type: {
+              kind: "UNION",
+              name: "Session",
+              ofType: null,
+            },
+            args: [
+              {
+                name: "deviceId",
+                type: {
+                  kind: "NON_NULL",
+                  ofType: {
+                    kind: "SCALAR",
+                    name: "Any",
+                  },
+                },
+              },
+              {
+                name: "userId",
+                type: {
+                  kind: "NON_NULL",
+                  ofType: {
+                    kind: "SCALAR",
+                    name: "Any",
+                  },
+                },
+              },
+            ],
+          },
           {
             name: "upstreamOauth2Link",
             type: {
@@ -1730,6 +1760,20 @@ export default {
         ],
         interfaces: [],
       },
+      {
+        kind: "UNION",
+        name: "Session",
+        possibleTypes: [
+          {
+            kind: "OBJECT",
+            name: "CompatSession",
+          },
+          {
+            kind: "OBJECT",
+            name: "Oauth2Session",
+          },
+        ],
+      },
       {
         kind: "OBJECT",
         name: "SetDisplayNamePayload",