Add MasWriter support for compat refresh tokens + some migration progress

Author: reivilibre

Cargo.lock

@@ -223,6 +223,7 @@ impl Options {
                 }
                 let mut writer = MasWriter::new(mas_connection, writer_mas_connections).await?;
 
+                let clock = SystemClock::default();
                 // TODO is this rng ok?
                 #[allow(clippy::disallowed_methods)]
                 let mut rng = thread_rng();
@@ -233,6 +234,7 @@ impl Options {
                     &mut reader,
                     &mut writer,
                     &mas_matrix.homeserver,
+                    &clock,
                     &mut rng,
                     &provider_id_mappings,
                 )

crates/cli/src/commands/syn2mas.rs

@@ -28,6 +28,7 @@ uuid = "1.10.0"
 ulid = { workspace = true, features = ["uuid"] }
 
 mas-config.workspace = true
+mas-storage.workspace = true
 
 [dev-dependencies]
 mas-storage-pg.workspace = true

crates/syn2mas/.sqlx/query-88975196c4c174d464b33aa015ce5d8cac3836701fc24922f4f0e8b98d330796.json

@@ -233,7 +233,7 @@ pub struct MasNewUpstreamOauthLink {
 pub struct MasNewCompatSession {
     pub session_id: Uuid,
     pub user_id: Uuid,
-    pub device_id: String,
+    pub device_id: Option<String>,
     pub human_name: Option<String>,
     pub created_at: DateTime<Utc>,
     pub is_synapse_admin: bool,
@@ -250,6 +250,14 @@ pub struct MasNewCompatAccessToken {
     pub expires_at: Option<DateTime<Utc>>,
 }
 
+pub struct MasNewCompatRefreshToken {
+    pub refresh_token_id: Uuid,
+    pub session_id: Uuid,
+    pub access_token_id: Uuid,
+    pub refresh_token: String,
+    pub created_at: DateTime<Utc>,
+}
+
 /// The 'version' of the password hashing scheme used for passwords when they
 /// are migrated from Synapse to MAS.
 /// This is version 1, as in the previous syn2mas script.
@@ -795,7 +803,7 @@ impl<'conn> MasWriter<'conn> {
                 Box::pin(async move {
                     let mut session_ids: Vec<Uuid> = Vec::with_capacity(sessions.len());
                     let mut user_ids: Vec<Uuid> = Vec::with_capacity(sessions.len());
-                    let mut device_ids: Vec<String> = Vec::with_capacity(sessions.len());
+                    let mut device_ids: Vec<Option<String>> = Vec::with_capacity(sessions.len());
                     let mut human_names: Vec<Option<String>> = Vec::with_capacity(sessions.len());
                     let mut created_ats: Vec<DateTime<Utc>> = Vec::with_capacity(sessions.len());
                     let mut is_synapse_admins: Vec<bool> = Vec::with_capacity(sessions.len());
@@ -845,7 +853,7 @@ impl<'conn> MasWriter<'conn> {
                         "#,
                         &session_ids[..],
                         &user_ids[..],
-                        &device_ids[..],
+                        &device_ids[..] as &[Option<String>],
                         &human_names[..] as &[Option<String>],
                         &created_ats[..],
                         &is_synapse_admins[..],
@@ -925,6 +933,66 @@ impl<'conn> MasWriter<'conn> {
             })
             .boxed()
     }
+
+    #[tracing::instrument(skip_all, level = Level::DEBUG)]
+    pub fn write_compat_refresh_tokens(
+        &mut self,
+        tokens: Vec<MasNewCompatRefreshToken>,
+    ) -> BoxFuture<'_, Result<(), Error>> {
+        self.writer_pool
+            .spawn_with_connection(move |conn| {
+                Box::pin(async move {
+                    let mut refresh_token_ids: Vec<Uuid> = Vec::with_capacity(tokens.len());
+                    let mut session_ids: Vec<Uuid> = Vec::with_capacity(tokens.len());
+                    let mut access_token_ids: Vec<Uuid> = Vec::with_capacity(tokens.len());
+                    let mut refresh_tokens: Vec<String> = Vec::with_capacity(tokens.len());
+                    let mut created_ats: Vec<DateTime<Utc>> = Vec::with_capacity(tokens.len());
+
+                    for MasNewCompatRefreshToken {
+                        refresh_token_id,
+                        session_id,
+                        access_token_id,
+                        refresh_token,
+                        created_at,
+                    } in tokens
+                    {
+                        refresh_token_ids.push(refresh_token_id);
+                        session_ids.push(session_id);
+                        access_token_ids.push(access_token_id);
+                        refresh_tokens.push(refresh_token);
+                        created_ats.push(created_at);
+                    }
+
+                    sqlx::query!(
+                        r#"
+                        INSERT INTO syn2mas__compat_refresh_tokens (
+                          compat_refresh_token_id,
+                          compat_session_id,
+                          compat_access_token_id,
+                          refresh_token,
+                          created_at)
+                        SELECT * FROM UNNEST(
+                          $1::UUID[],
+                          $2::UUID[],
+                          $3::UUID[],
+                          $4::TEXT[],
+                          $5::TIMESTAMP WITH TIME ZONE[])
+                        "#,
+                        &refresh_token_ids[..],
+                        &session_ids[..],
+                        &access_token_ids[..],
+                        &refresh_tokens[..],
+                        &created_ats[..],
+                    )
+                    .execute(&mut *conn)
+                    .await
+                    .into_database("writing compat refresh tokens to MAS")?;
+
+                    Ok(())
+                })
+            })
+            .boxed()
+    }
 }
 
 // How many entries to buffer at once, before writing a batch of rows to the
@@ -1003,8 +1071,9 @@ mod test {
 
     use crate::{
         mas_writer::{
-            MasNewCompatAccessToken, MasNewCompatSession, MasNewEmailThreepid,
-            MasNewUnsupportedThreepid, MasNewUpstreamOauthLink, MasNewUser, MasNewUserPassword,
+            MasNewCompatAccessToken, MasNewCompatRefreshToken, MasNewCompatSession,
+            MasNewEmailThreepid, MasNewUnsupportedThreepid, MasNewUpstreamOauthLink, MasNewUser,
+            MasNewUserPassword,
         },
         LockedMasDatabase, MasWriter,
     };
@@ -1292,7 +1361,7 @@ mod test {
                 user_id: Uuid::from_u128(1u128),
                 session_id: Uuid::from_u128(5u128),
                 created_at: DateTime::default(),
-                device_id: "ADEVICE".to_owned(),
+                device_id: Some("ADEVICE".to_owned()),
                 human_name: Some("alice's pinephone".to_owned()),
                 is_synapse_admin: true,
                 last_active_at: Some(DateTime::default()),
@@ -1329,7 +1398,7 @@ mod test {
                 user_id: Uuid::from_u128(1u128),
                 session_id: Uuid::from_u128(5u128),
                 created_at: DateTime::default(),
-                device_id: "ADEVICE".to_owned(),
+                device_id: Some("ADEVICE".to_owned()),
                 human_name: None,
                 is_synapse_admin: false,
                 last_active_at: None,
@@ -1354,4 +1423,63 @@ mod test {
 
         assert_db_snapshot!(&mut conn);
     }
+
+    /// Tests writing a single user, with a device, an access token and a refresh token.
+    #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
+    async fn test_write_user_with_refresh_token(pool: PgPool) {
+        let mut conn = pool.acquire().await.unwrap();
+        let mut writer = make_mas_writer(&pool, &mut conn).await;
+
+        writer
+            .write_users(vec![MasNewUser {
+                user_id: Uuid::from_u128(1u128),
+                username: "alice".to_owned(),
+                created_at: DateTime::default(),
+                locked_at: None,
+                can_request_admin: false,
+            }])
+            .await
+            .expect("failed to write user");
+
+        writer
+            .write_compat_sessions(vec![MasNewCompatSession {
+                user_id: Uuid::from_u128(1u128),
+                session_id: Uuid::from_u128(5u128),
+                created_at: DateTime::default(),
+                device_id: Some("ADEVICE".to_owned()),
+                human_name: None,
+                is_synapse_admin: false,
+                last_active_at: None,
+                last_active_ip: None,
+                user_agent: None,
+            }])
+            .await
+            .expect("failed to write compat session");
+
+        writer
+            .write_compat_access_tokens(vec![MasNewCompatAccessToken {
+                token_id: Uuid::from_u128(6u128),
+                session_id: Uuid::from_u128(5u128),
+                access_token: "syt_zxcvzxcvzxcvzxcv_zxcv".to_owned(),
+                created_at: DateTime::default(),
+                expires_at: None,
+            }])
+            .await
+            .expect("failed to write access token");
+
+        writer
+            .write_compat_refresh_tokens(vec![MasNewCompatRefreshToken {
+                refresh_token_id: Uuid::from_u128(7u128),
+                session_id: Uuid::from_u128(5u128),
+                access_token_id: Uuid::from_u128(6u128),
+                refresh_token: "syr_zxcvzxcvzxcvzxcv_zxcv".to_owned(),
+                created_at: DateTime::default(),
+            }])
+            .await
+            .expect("failed to write refresh token");
+
+        writer.finish().await.expect("failed to finish MasWriter");
+
+        assert_db_snapshot!(&mut conn);
+    }
 }

crates/syn2mas/Cargo.toml

@@ -0,0 +1,36 @@
+---
+source: crates/syn2mas/src/mas_writer/mod.rs
+expression: db_snapshot
+---
+compat_access_tokens:
+  - access_token: syt_zxcvzxcvzxcvzxcv_zxcv
+    compat_access_token_id: 00000000-0000-0000-0000-000000000006
+    compat_session_id: 00000000-0000-0000-0000-000000000005
+    created_at: "1970-01-01 00:00:00+00"
+    expires_at: ~
+compat_refresh_tokens:
+  - compat_access_token_id: 00000000-0000-0000-0000-000000000006
+    compat_refresh_token_id: 00000000-0000-0000-0000-000000000007
+    compat_session_id: 00000000-0000-0000-0000-000000000005
+    consumed_at: ~
+    created_at: "1970-01-01 00:00:00+00"
+    refresh_token: syr_zxcvzxcvzxcvzxcv_zxcv
+compat_sessions:
+  - compat_session_id: 00000000-0000-0000-0000-000000000005
+    created_at: "1970-01-01 00:00:00+00"
+    device_id: ADEVICE
+    finished_at: ~
+    human_name: ~
+    is_synapse_admin: "false"
+    last_active_at: ~
+    last_active_ip: ~
+    user_agent: ~
+    user_id: 00000000-0000-0000-0000-000000000001
+    user_session_id: ~
+users:
+  - can_request_admin: "false"
+    created_at: "1970-01-01 00:00:00+00"
+    locked_at: ~
+    primary_user_email_id: ~
+    user_id: 00000000-0000-0000-0000-000000000001
+    username: alice