Fully sync the devices with the homeserver

crates/tasks/src/matrix.rs

@@ -12,13 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashSet;
+
 use anyhow::Context;
 use apalis_core::{context::JobContext, executor::TokioExecutor, monitor::Monitor};
+use mas_data_model::Device;
 use mas_matrix::ProvisionRequest;
 use mas_storage::{
-    job::{DeleteDeviceJob, JobWithSpanContext, ProvisionDeviceJob, ProvisionUserJob},
+    compat::CompatSessionFilter,
+    job::{
+        DeleteDeviceJob, JobRepositoryExt as _, JobWithSpanContext, ProvisionDeviceJob,
+        ProvisionUserJob, SyncDevicesJob,
+    },
+    oauth2::OAuth2SessionFilter,
     user::{UserEmailRepository, UserRepository},
-    RepositoryAccess,
+    Pagination, RepositoryAccess,
 };
 use tracing::info;
 
@@ -56,9 +64,6 @@ async fn provision_user(
         .filter(|email| email.confirmed_at.is_some())
         .map(|email| email.email)
         .collect();
-
-    repo.cancel().await?;
-
     let mut request = ProvisionRequest::new(mxid.clone(), user.sub.clone()).set_emails(emails);
 
     if let Some(display_name) = job.display_name_to_set() {
@@ -73,6 +78,12 @@ async fn provision_user(
         info!(%user.id, %mxid, "User updated");
     }
 
+    // Schedule a device sync job
+    let sync_device_job = SyncDevicesJob::new(&user);
+    repo.job().schedule_job(sync_device_job).await?;
+
+    repo.save().await?;
+
     Ok(())
 }
 
@@ -144,6 +155,84 @@ async fn delete_device(
     Ok(())
 }
 
+/// Job to sync the list of devices of a user with the homeserver.
+#[tracing::instrument(
+    name = "job.sync_devices",
+    fields(user.id = %job.user_id()),
+    skip_all,
+    err(Debug),
+)]
+async fn sync_devices(
+    job: JobWithSpanContext<SyncDevicesJob>,
+    ctx: JobContext,
+) -> Result<(), anyhow::Error> {
+    let state = ctx.state();
+    let matrix = state.matrix_connection();
+    let mut repo = state.repository().await?;
+
+    let user = repo
+        .user()
+        .lookup(job.user_id())
+        .await?
+        .context("User not found")?;
+
+    let mut devices = HashSet::new();
+
+    // Cycle through all the compat sessions of the user, and grab the devices
+    let mut cursor = Pagination::first(100);
+    loop {
+        let page = repo
+            .compat_session()
+            .list(
+                CompatSessionFilter::new().for_user(&user).active_only(),
+                cursor,
+            )
+            .await?;
+
+        for (compat_session, _) in page.edges {
+            devices.insert(compat_session.device.as_str().to_owned());
+            cursor = cursor.after(compat_session.id);
+        }
+
+        if !page.has_next_page {
+            break;
+        }
+    }
+
+    // Cycle though all the oauth2 sessions of the user, and grab the devices
+    let mut cursor = Pagination::first(100);
+    loop {
+        let page = repo
+            .oauth2_session()
+            .list(
+                OAuth2SessionFilter::new().for_user(&user).active_only(),
+                cursor,
+            )
+            .await?;
+
+        for oauth2_session in page.edges {
+            for scope in &*oauth2_session.scope {
+                if let Some(device) = Device::from_scope_token(scope) {
+                    devices.insert(device.as_str().to_owned());
+                }
+            }
+
+            cursor = cursor.after(oauth2_session.id);
+        }
+
+        if !page.has_next_page {
+            break;
+        }
+    }
+
+    // We now have a complete list of devices, so we can sync them with the
+    // homeserver
+    let mxid = matrix.mxid(&user.username);
+    matrix.sync_devices(&mxid, devices).await?;
+
+    Ok(())
+}
+
 pub(crate) fn register(
     suffix: &str,
     monitor: Monitor<TokioExecutor>,
@@ -156,9 +245,12 @@ pub(crate) fn register(
         crate::build!(ProvisionDeviceJob => provision_device, suffix, state, storage_factory);
     let delete_device_worker =
         crate::build!(DeleteDeviceJob => delete_device, suffix, state, storage_factory);
+    let sync_devices_worker =
+        crate::build!(SyncDevicesJob => sync_devices, suffix, state, storage_factory);
 
     monitor
         .register(provision_user_worker)
         .register(provision_device_worker)
         .register(delete_device_worker)
+        .register(sync_devices_worker)
 }