Add MasWriter support for compat sessions

crates/syn2mas/.sqlx/query-396c97dbfbc932c73301daa7376e36f4ef03e1224a0a89a921e5a3d398a5d35c.json

@@ -0,0 +1,22 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "\n                        INSERT INTO syn2mas__compat_sessions (\n                          compat_session_id, user_id,\n                          device_id, human_name,\n                          created_at, is_synapse_admin,\n                          last_active_at, last_active_ip,\n                          user_agent)\n                        SELECT * FROM UNNEST(\n                          $1::UUID[], $2::UUID[],\n                          $3::TEXT[], $4::TEXT[],\n                          $5::TIMESTAMP WITH TIME ZONE[], $6::BOOLEAN[],\n                          $7::TIMESTAMP WITH TIME ZONE[], $8::INET[],\n                          $9::TEXT[])\n                        ",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Left": [
+        "UuidArray",
+        "UuidArray",
+        "TextArray",
+        "TextArray",
+        "TimestamptzArray",
+        "BoolArray",
+        "TimestamptzArray",
+        "InetArray",
+        "TextArray"
+      ]
+    },
+    "nullable": []
+  },
+  "hash": "396c97dbfbc932c73301daa7376e36f4ef03e1224a0a89a921e5a3d398a5d35c"
+}

crates/syn2mas/src/mas_writer/mod.rs

@@ -7,7 +7,7 @@
 //!
 //! This module is responsible for writing new records to MAS' database.
 
-use std::fmt::Display;
+use std::{fmt::Display, net::IpAddr};
 
 use chrono::{DateTime, Utc};
 use futures_util::{future::BoxFuture, FutureExt, TryStreamExt};
@@ -230,6 +230,18 @@ pub struct MasNewUpstreamOauthLink {
     pub created_at: DateTime<Utc>,
 }
 
+pub struct MasNewCompatSession {
+    pub session_id: Uuid,
+    pub user_id: Uuid,
+    pub device_id: String,
+    pub human_name: Option<String>,
+    pub created_at: DateTime<Utc>,
+    pub is_synapse_admin: bool,
+    pub last_active_at: Option<DateTime<Utc>>,
+    pub last_active_ip: Option<IpAddr>,
+    pub user_agent: Option<String>,
+}
+
 /// 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.
@@ -761,6 +773,85 @@ impl<'conn> MasWriter<'conn> {
             })
         }).boxed()
     }
+
+    #[tracing::instrument(skip_all, level = Level::DEBUG)]
+    pub fn write_compat_sessions(
+        &mut self,
+        sessions: Vec<MasNewCompatSession>,
+    ) -> BoxFuture<'_, Result<(), Error>> {
+        self.writer_pool
+            .spawn_with_connection(move |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 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());
+                    let mut last_active_ats: Vec<Option<DateTime<Utc>>> =
+                        Vec::with_capacity(sessions.len());
+                    let mut last_active_ips: Vec<Option<IpAddr>> =
+                        Vec::with_capacity(sessions.len());
+                    let mut user_agents: Vec<Option<String>> = Vec::with_capacity(sessions.len());
+
+                    for MasNewCompatSession {
+                        session_id,
+                        user_id,
+                        device_id,
+                        human_name,
+                        created_at,
+                        is_synapse_admin,
+                        last_active_at,
+                        last_active_ip,
+                        user_agent,
+                    } in sessions
+                    {
+                        session_ids.push(session_id);
+                        user_ids.push(user_id);
+                        device_ids.push(device_id);
+                        human_names.push(human_name);
+                        created_ats.push(created_at);
+                        is_synapse_admins.push(is_synapse_admin);
+                        last_active_ats.push(last_active_at);
+                        last_active_ips.push(last_active_ip);
+                        user_agents.push(user_agent);
+                    }
+
+                    sqlx::query!(
+                        r#"
+                        INSERT INTO syn2mas__compat_sessions (
+                          compat_session_id, user_id,
+                          device_id, human_name,
+                          created_at, is_synapse_admin,
+                          last_active_at, last_active_ip,
+                          user_agent)
+                        SELECT * FROM UNNEST(
+                          $1::UUID[], $2::UUID[],
+                          $3::TEXT[], $4::TEXT[],
+                          $5::TIMESTAMP WITH TIME ZONE[], $6::BOOLEAN[],
+                          $7::TIMESTAMP WITH TIME ZONE[], $8::INET[],
+                          $9::TEXT[])
+                        "#,
+                        &session_ids[..],
+                        &user_ids[..],
+                        &device_ids[..],
+                        &human_names[..] as &[Option<String>],
+                        &created_ats[..],
+                        &is_synapse_admins[..],
+                        // We need to override the typing for arrays of optionals (sqlx limitation)
+                        &last_active_ats[..] as &[Option<DateTime<Utc>>],
+                        &last_active_ips[..] as &[Option<IpAddr>],
+                        &user_agents[..] as &[Option<String>],
+                    )
+                    .execute(&mut *conn)
+                    .await
+                    .into_database("writing compat sessions to MAS")?;
+
+                    Ok(())
+                })
+            })
+            .boxed()
+    }
 }
 
 // How many entries to buffer at once, before writing a batch of rows to the
@@ -839,8 +930,8 @@ mod test {
 
     use crate::{
         mas_writer::{
-            MasNewEmailThreepid, MasNewUnsupportedThreepid, MasNewUpstreamOauthLink, MasNewUser,
-            MasNewUserPassword,
+            MasNewCompatSession, MasNewEmailThreepid, MasNewUnsupportedThreepid,
+            MasNewUpstreamOauthLink, MasNewUser, MasNewUserPassword,
         },
         LockedMasDatabase, MasWriter,
     };
@@ -1105,4 +1196,41 @@ mod test {
 
         assert_db_snapshot!(&mut conn);
     }
+
+    /// Tests writing a single user, with a device (compat session).
+    #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
+    async fn test_write_user_with_device(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: "ADEVICE".to_owned(),
+                human_name: Some("alice's pinephone".to_owned()),
+                is_synapse_admin: true,
+                last_active_at: Some(DateTime::default()),
+                last_active_ip: Some("203.0.113.1".parse().unwrap()),
+                user_agent: Some("Browser/5.0".to_owned()),
+            }])
+            .await
+            .expect("failed to write link");
+
+        writer.finish().await.expect("failed to finish MasWriter");
+
+        assert_db_snapshot!(&mut conn);
+    }
 }

crates/syn2mas/src/mas_writer/snapshots/syn2mas__mas_writer__test__write_user_with_device.snap

@@ -0,0 +1,23 @@
+---
+source: crates/syn2mas/src/mas_writer/mod.rs
+expression: db_snapshot
+---
+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: "alice's pinephone"
+    is_synapse_admin: "true"
+    last_active_at: "1970-01-01 00:00:00+00"
+    last_active_ip: 203.0.113.1/32
+    user_agent: Browser/5.0
+    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

crates/syn2mas/src/mas_writer/syn2mas_revert_temporary_tables.sql

@@ -13,3 +13,6 @@ ALTER TABLE syn2mas__user_passwords RENAME TO user_passwords;
 ALTER TABLE syn2mas__user_emails RENAME TO user_emails;
 ALTER TABLE syn2mas__user_unsupported_third_party_ids RENAME TO user_unsupported_third_party_ids;
 ALTER TABLE syn2mas__upstream_oauth_links RENAME TO upstream_oauth_links;
+ALTER TABLE syn2mas__compat_sessions RENAME TO compat_sessions;
+ALTER TABLE syn2mas__compat_access_tokens RENAME TO compat_access_tokens;
+ALTER TABLE syn2mas__compat_refresh_tokens RENAME TO compat_refresh_tokens;

crates/syn2mas/src/mas_writer/syn2mas_temporary_tables.sql

@@ -42,3 +42,6 @@ ALTER TABLE user_passwords RENAME TO syn2mas__user_passwords;
 ALTER TABLE user_emails RENAME TO syn2mas__user_emails;
 ALTER TABLE user_unsupported_third_party_ids RENAME TO syn2mas__user_unsupported_third_party_ids;
 ALTER TABLE upstream_oauth_links RENAME TO syn2mas__upstream_oauth_links;
+ALTER TABLE compat_sessions RENAME TO syn2mas__compat_sessions;
+ALTER TABLE compat_access_tokens RENAME TO syn2mas__compat_access_tokens;
+ALTER TABLE compat_refresh_tokens RENAME TO syn2mas__compat_refresh_tokens;

crates/syn2mas/src/migration.rs

@@ -11,7 +11,10 @@
 //! This module does not implement any of the safety checks that should be run
 //! *before* the migration.
 
-use std::{collections::HashMap, pin::pin};
+use std::{
+    collections::{HashMap, HashSet},
+    pin::pin,
+};
 
 use chrono::{DateTime, Utc};
 use compact_str::CompactString;
@@ -25,11 +28,12 @@ use uuid::Uuid;
 
 use crate::{
     mas_writer::{
-        self, MasNewEmailThreepid, MasNewUnsupportedThreepid, MasNewUpstreamOauthLink, MasNewUser,
-        MasNewUserPassword, MasWriteBuffer, MasWriter,
+        self, MasNewCompatSession, MasNewEmailThreepid, MasNewUnsupportedThreepid,
+        MasNewUpstreamOauthLink, MasNewUser, MasNewUserPassword, MasWriteBuffer, MasWriter,
     },
     synapse_reader::{
-        self, ExtractLocalpartError, FullUserId, SynapseExternalId, SynapseThreepid, SynapseUser,
+        self, ExtractLocalpartError, FullUserId, SynapseDevice, SynapseExternalId, SynapseThreepid,
+        SynapseUser,
     },
     SynapseReader,
 };
@@ -66,12 +70,9 @@ pub enum Error {
 struct UsersMigrated {
     /// Lookup table from user localpart to that user's UUID in MAS.
     user_localparts_to_uuid: HashMap<CompactString, Uuid>,
-}
 
-struct DevicesMigrated {
-    /// Lookup table from `(user_id, device_id)` pairs to the
-    /// UUID of the `compat_session` in MAS
-    devices_to_uuid: HashMap<(CompactString, CompactString), Uuid>,
+    /// Set of user UUIDs that correspond to Synapse admins
+    synapse_admins: HashSet<Uuid>,
 }
 
 /// Performs a migration from Synapse's database to MAS' database.
@@ -127,18 +128,14 @@ pub async fn migrate(
     )
     .await?;
 
-    let migrated_devices = migrate_devices(
-        synapse,
-        mas,
-        counts
-            .devices
-            .try_into()
-            .expect("More than usize::MAX devices — unable to handle this many!"),
-        server_name,
-        rng,
-        &migrated_users.user_localparts_to_uuid,
-    )
-    .await?;
+    // `(MAS user_id, device_id)` mapped to `compat_session` ULID
+    let mut devices_to_compat_sessions: HashMap<(Uuid, CompactString), Uuid> =
+        HashMap::with_capacity(
+            counts
+                .devices
+                .try_into()
+                .expect("More than usize::MAX devices — unable to handle this many!"),
+        );
 
     migrate_access_and_refresh_tokens(
         synapse,
@@ -146,7 +143,18 @@ pub async fn migrate(
         server_name,
         rng,
         &migrated_users.user_localparts_to_uuid,
-        &migrated_devices.devices_to_uuid,
+        &mut devices_to_compat_sessions,
+    )
+    .await?;
+
+    migrate_devices(
+        synapse,
+        mas,
+        server_name,
+        rng,
+        &migrated_users.user_localparts_to_uuid,
+        &mut devices_to_compat_sessions,
+        &migrated_users.synapse_admins,
     )
     .await?;
 
@@ -166,11 +174,19 @@ async fn migrate_users(
     let mut users_stream = pin!(synapse.read_users());
     // TODO is 1:1 capacity enough for a hashmap?
     let mut user_localparts_to_uuid = HashMap::with_capacity(user_count_hint);
+    let mut synapse_admins = HashSet::new();
 
     while let Some(user_res) = users_stream.next().await {
         let user = user_res.into_synapse("reading user")?;
         let (mas_user, mas_password_opt) = transform_user(&user, server_name, rng)?;
 
+        if bool::from(user.admin) {
+            // Note down the fact that this user is a Synapse admin,
+            // because we will grant their existing devices the Synapse admin
+            // flag
+            synapse_admins.insert(mas_user.user_id);
+        }
+
         user_localparts_to_uuid.insert(CompactString::new(&mas_user.username), mas_user.user_id);
 
         user_buffer
@@ -194,6 +210,7 @@ async fn migrate_users(
 
     Ok(UsersMigrated {
         user_localparts_to_uuid,
+        synapse_admins,
     })
 }
 
@@ -341,21 +358,100 @@ async fn migrate_external_ids(
 
     Ok(())
 }
+
+/// Migrate devices from Synapse to MAS (as compat sessions).
+///
+/// In order to get the right session creation timestamps, the access tokens
+/// must counterintuitively be migrated first, with the ULIDs passed in as
+/// `devices`.
+///
+/// This is because only access tokens store a timestamp that in any way
+/// resembles a creation timestamp.
 #[tracing::instrument(skip_all, level = Level::INFO)]
 async fn migrate_devices(
     synapse: &mut SynapseReader<'_>,
     mas: &mut MasWriter<'_>,
-    device_count_hint: usize,
     server_name: &str,
     rng: &mut impl RngCore,
     user_localparts_to_uuid: &HashMap<CompactString, Uuid>,
-) -> Result<DevicesMigrated, Error> {
-    // TODO is 1:1 enough capacity for a HashMap?
-    let mut devices_to_uuid = HashMap::with_capacity(device_count_hint);
+    devices: &mut HashMap<(Uuid, CompactString), Uuid>,
+    synapse_admins: &HashSet<Uuid>,
+) -> Result<(), Error> {
+    let mut devices_stream = pin!(synapse.read_devices());
+    let mut write_buffer = MasWriteBuffer::new(MasWriter::write_compat_sessions);
 
-    todo!();
+    while let Some(device_res) = devices_stream.next().await {
+        let SynapseDevice {
+            user_id: synapse_user_id,
+            device_id,
+            display_name,
+            last_seen,
+            ip,
+            user_agent,
+        } = device_res.into_synapse("reading Synapse device")?;
 
-    Ok(DevicesMigrated { devices_to_uuid })
+        let username = synapse_user_id
+            .extract_localpart(server_name)
+            .into_extract_localpart(synapse_user_id.clone())?
+            .to_owned();
+        let Some(user_id) = user_localparts_to_uuid.get(username.as_str()).copied() else {
+            return Err(Error::MissingUserFromDependentTable {
+                table: "devices".to_owned(),
+                user: synapse_user_id,
+            });
+        };
+
+        let session_id = *devices
+            .entry((user_id, CompactString::new(&device_id)))
+            .or_insert_with(||
+                // We don't have a creation time for this device (as it has no access token),
+                // so use now as a least-evil fallback.
+                Ulid::with_source(rng).into());
+        let created_at = Ulid::from(session_id).datetime().into();
+
+        // As we're using a real IP type in the MAS database, it is possible
+        // that we encounter invalid IP addresses in the Synapse database.
+        // In that case, we should ignore them, but still log a warning.
+        let last_active_ip = ip.and_then(|ip| {
+            ip.parse()
+                .map_err(|e| {
+                    tracing::warn!(
+                        error = &e as &dyn std::error::Error,
+                        mxid = %synapse_user_id,
+                        %device_id,
+                        %ip,
+                        "Failed to parse device IP, ignoring"
+                    );
+                })
+                .ok()
+        });
+
+        // TODO skip access tokens for deactivated users
+        write_buffer
+            .write(
+                mas,
+                MasNewCompatSession {
+                    session_id,
+                    user_id,
+                    device_id,
+                    human_name: display_name,
+                    created_at,
+                    is_synapse_admin: synapse_admins.contains(&user_id),
+                    last_active_at: last_seen.map(DateTime::from),
+                    last_active_ip,
+                    user_agent,
+                },
+            )
+            .await
+            .into_mas("writing compat sessions")?;
+    }
+
+    write_buffer
+        .finish(mas)
+        .await
+        .into_mas("writing compat sessions")?;
+
+    Ok(())
 }
 
 #[tracing::instrument(skip_all, level = Level::INFO)]
@@ -365,8 +461,11 @@ async fn migrate_access_and_refresh_tokens(
     server_name: &str,
     rng: &mut impl RngCore,
     user_localparts_to_uuid: &HashMap<CompactString, Uuid>,
-    devices: &HashMap<(CompactString, CompactString), Uuid>,
+    devices: &mut HashMap<(Uuid, CompactString), Uuid>,
 ) -> Result<(), Error> {
+    let mut access_token_stream = pin!(synapse.read_access_tokens());
+    // let mut write_buffer =
+    // MasWriteBuffer::new(MasWriter::write_compat_access_token);
     todo!();
 
     Ok(())