Add a GraphQL mutation to create arbitrary OAuth2 sessions.

Author: sandhose

crates/graphql/src/lib.rs

@@ -131,6 +131,13 @@ impl Requester {
         }
     }
 
+    fn oauth2_session(&self) -> Option<&Session> {
+        match self {
+            Self::OAuth2Session(session, _) => Some(session),
+            Self::BrowserSession(_) | Self::Anonymous => None,
+        }
+    }
+
     /// Returns true if the requester can access the resource.
     fn is_owner_or_admin(&self, resource: &impl OwnerId) -> bool {
         // If the requester is an admin, they can do anything.

crates/graphql/src/mutations/oauth2_session.rs

@@ -13,13 +13,19 @@
 // limitations under the License.
 
 use anyhow::Context as _;
-use async_graphql::{Context, Enum, InputObject, Object, ID};
-use mas_data_model::Device;
+use async_graphql::{Context, Description, Enum, InputObject, Object, ID};
+use chrono::Duration;
+use mas_data_model::{Device, TokenType};
 use mas_storage::{
-    job::{DeleteDeviceJob, JobRepositoryExt},
-    oauth2::OAuth2SessionRepository,
+    job::{DeleteDeviceJob, JobRepositoryExt, ProvisionDeviceJob},
+    oauth2::{
+        OAuth2AccessTokenRepository, OAuth2ClientRepository, OAuth2RefreshTokenRepository,
+        OAuth2SessionRepository,
+    },
+    user::UserRepository,
     RepositoryAccess,
 };
+use oauth2_types::scope::Scope;
 
 use crate::{
     model::{NodeType, OAuth2Session},
@@ -31,6 +37,45 @@ pub struct OAuth2SessionMutations {
     _private: (),
 }
 
+/// The input of the `createOauth2Session` mutation.
+#[derive(InputObject)]
+pub struct CreateOAuth2SessionInput {
+    /// The scope of the session
+    scope: String,
+
+    /// The ID of the user for which to create the session
+    user_id: ID,
+
+    /// Whether the session should issue a never-expiring access token
+    permanent: Option<bool>,
+}
+
+/// The payload of the `createOauth2Session` mutation.
+#[derive(Description)]
+pub struct CreateOAuth2SessionPayload {
+    access_token: String,
+    refresh_token: Option<String>,
+    session: mas_data_model::Session,
+}
+
+#[Object(use_type_description)]
+impl CreateOAuth2SessionPayload {
+    /// Access token for this session
+    pub async fn access_token(&self) -> &str {
+        &self.access_token
+    }
+
+    /// Refresh token for this session, if it is not a permanent session
+    pub async fn refresh_token(&self) -> Option<&str> {
+        self.refresh_token.as_deref()
+    }
+
+    /// The OAuth 2.0 session which was just created
+    pub async fn oauth2_session(&self) -> OAuth2Session {
+        OAuth2Session(self.session.clone())
+    }
+}
+
 /// The input of the `endOauth2Session` mutation.
 #[derive(InputObject)]
 pub struct EndOAuth2SessionInput {
@@ -75,6 +120,93 @@ impl EndOAuth2SessionPayload {
 
 #[Object]
 impl OAuth2SessionMutations {
+    /// Create a new arbitrary OAuth 2.0 Session.
+    ///
+    /// Only available for administrators.
+    async fn create_oauth2_session(
+        &self,
+        ctx: &Context<'_>,
+        input: CreateOAuth2SessionInput,
+    ) -> Result<CreateOAuth2SessionPayload, async_graphql::Error> {
+        let state = ctx.state();
+        let user_id = NodeType::User.extract_ulid(&input.user_id)?;
+        let scope: Scope = input.scope.parse().context("Invalid scope")?;
+        let permanent = input.permanent.unwrap_or(false);
+        let requester = ctx.requester();
+
+        if !requester.is_admin() {
+            return Err(async_graphql::Error::new("Unauthorized"));
+        }
+
+        let session = requester
+            .oauth2_session()
+            .context("Requester should be a OAuth 2.0 client")?;
+
+        let mut repo = state.repository().await?;
+        let clock = state.clock();
+        let mut rng = state.rng();
+
+        let client = repo
+            .oauth2_client()
+            .lookup(session.client_id)
+            .await?
+            .context("Client not found")?;
+
+        let user = repo
+            .user()
+            .lookup(user_id)
+            .await?
+            .context("User not found")?;
+
+        // Generate a new access token
+        let access_token = TokenType::AccessToken.generate(&mut rng);
+
+        // Create the OAuth 2.0 Session
+        let session = repo
+            .oauth2_session()
+            .add(&mut rng, &clock, &client, Some(&user), None, scope)
+            .await?;
+
+        // Look for devices to provision
+        for scope in &*session.scope {
+            if let Some(device) = Device::from_scope_token(scope) {
+                repo.job()
+                    .schedule_job(ProvisionDeviceJob::new(&user, &device))
+                    .await?;
+            }
+        }
+
+        let ttl = if permanent {
+            // XXX: that's lazy
+            Duration::days(365 * 50)
+        } else {
+            Duration::minutes(5)
+        };
+        let access_token = repo
+            .oauth2_access_token()
+            .add(&mut rng, &clock, &session, access_token, ttl)
+            .await?;
+
+        let refresh_token = if !permanent {
+            let refresh_token = TokenType::RefreshToken.generate(&mut rng);
+
+            let refresh_token = repo
+                .oauth2_refresh_token()
+                .add(&mut rng, &clock, &session, &access_token, refresh_token)
+                .await?;
+
+            Some(refresh_token)
+        } else {
+            None
+        };
+
+        Ok(CreateOAuth2SessionPayload {
+            session,
+            access_token: access_token.access_token,
+            refresh_token: refresh_token.map(|t| t.refresh_token),
+        })
+    }
+
     async fn end_oauth2_session(
         &self,
         ctx: &Context<'_>,

crates/handlers/src/graphql/tests.rs

@@ -510,6 +510,7 @@ async fn test_oauth2_client_credentials(pool: PgPool) {
                 mutation {
                     addUser(input: {username: "alice"}) {
                         user {
+                            id
                             username
                         }
                     }
@@ -519,15 +520,42 @@ async fn test_oauth2_client_credentials(pool: PgPool) {
     let response = state.request(request).await;
     response.assert_status(StatusCode::OK);
     let response: GraphQLResponse = response.json();
-    assert!(response.errors.is_empty());
+    assert!(response.errors.is_empty(), "{:?}", response.errors);
+    let user_id = &response.data["addUser"]["user"]["id"];
+
     assert_eq!(
         response.data,
         serde_json::json!({
             "addUser": {
                 "user": {
+                    "id": user_id,
                     "username": "alice"
                 }
             }
         })
     );
+
+    // We should now be able to create an arbitrary access token for the user
+    let request = Request::post("/graphql")
+        .bearer(&access_token)
+        .json(serde_json::json!({
+            "query": r#"
+                mutation CreateSession($userId: String!, $scope: String!) {
+                    createOauth2Session(input: {userId: $userId, permanent: true, scope: $scope}) {
+                        accessToken
+                        refreshToken
+                    }
+                }
+            "#,
+            "variables": {
+                "userId": user_id,
+                "scope": "urn:matrix:org.matrix.msc2967.client:device:AABBCCDDEE urn:matrix:org.matrix.msc2967.client:api:* urn:synapse:admin:*"
+            },
+        }));
+    let response = state.request(request).await;
+    response.assert_status(StatusCode::OK);
+    let response: GraphQLResponse = response.json();
+    assert!(response.errors.is_empty(), "{:?}", response.errors);
+    assert!(response.data["createOauth2Session"]["refreshToken"].is_null());
+    assert!(response.data["createOauth2Session"]["accessToken"].is_string());
 }

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

@@ -14,7 +14,7 @@
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
-use mas_data_model::{BrowserSession, Client, Session, SessionState};
+use mas_data_model::{BrowserSession, Client, Session, SessionState, User};
 use mas_storage::{
     oauth2::{OAuth2SessionFilter, OAuth2SessionRepository},
     Clock, Page, Pagination,
@@ -133,24 +133,23 @@ impl<'c> OAuth2SessionRepository for PgOAuth2SessionRepository<'c> {
     }
 
     #[tracing::instrument(
-        name = "db.oauth2_session.add_from_browser_session",
+        name = "db.oauth2_session.add",
         skip_all,
         fields(
             db.statement,
-            %user_session.id,
-            user.id = %user_session.user.id,
             %client.id,
             session.id,
             session.scope = %scope,
         ),
         err,
     )]
-    async fn add_from_browser_session(
+    async fn add(
         &mut self,
         rng: &mut (dyn RngCore + Send),
         clock: &dyn Clock,
         client: &Client,
-        user_session: &BrowserSession,
+        user: Option<&User>,
+        user_session: Option<&BrowserSession>,
         scope: Scope,
     ) -> Result<Session, Self::Error> {
         let created_at = clock.now();
@@ -172,8 +171,8 @@ impl<'c> OAuth2SessionRepository for PgOAuth2SessionRepository<'c> {
                 VALUES ($1, $2, $3, $4, $5, $6)
             "#,
             Uuid::from(id),
-            Uuid::from(user_session.user.id),
-            Uuid::from(user_session.id),
+            user.map(|u| Uuid::from(u.id)),
+            user_session.map(|s| Uuid::from(s.id)),
             Uuid::from(client.id),
             &scope_list,
             created_at,
@@ -186,62 +185,8 @@ impl<'c> OAuth2SessionRepository for PgOAuth2SessionRepository<'c> {
             id,
             state: SessionState::Valid,
             created_at,
-            user_id: Some(user_session.user.id),
-            user_session_id: Some(user_session.id),
-            client_id: client.id,
-            scope,
-        })
-    }
-
-    #[tracing::instrument(
-        name = "db.oauth2_session.add_from_client_credentials",
-        skip_all,
-        fields(
-            db.statement,
-            %client.id,
-            session.id,
-            session.scope = %scope,
-        ),
-        err,
-    )]
-    async fn add_from_client_credentials(
-        &mut self,
-        rng: &mut (dyn RngCore + Send),
-        clock: &dyn Clock,
-        client: &Client,
-        scope: Scope,
-    ) -> Result<Session, Self::Error> {
-        let created_at = clock.now();
-        let id = Ulid::from_datetime_with_source(created_at.into(), rng);
-        tracing::Span::current().record("session.id", tracing::field::display(id));
-
-        let scope_list: Vec<String> = scope.iter().map(|s| s.as_str().to_owned()).collect();
-
-        sqlx::query!(
-            r#"
-                INSERT INTO oauth2_sessions
-                    ( oauth2_session_id
-                    , oauth2_client_id
-                    , scope_list
-                    , created_at
-                    )
-                VALUES ($1, $2, $3, $4)
-            "#,
-            Uuid::from(id),
-            Uuid::from(client.id),
-            &scope_list,
-            created_at,
-        )
-        .traced()
-        .execute(&mut *self.conn)
-        .await?;
-
-        Ok(Session {
-            id,
-            state: SessionState::Valid,
-            created_at,
-            user_id: None,
-            user_session_id: None,
+            user_id: user.map(|u| u.id),
+            user_session_id: user_session.map(|s| s.id),
             client_id: client.id,
             scope,
         })