xmtp_mls/worker/device_sync/

mod.rs

use crate::{
    client::ClientError,
    context::XmtpSharedContext,
    groups::{
        GroupError, MlsGroup, PreconfiguredPolicies, send_message_opts, summary::SyncSummary,
    },
    mls_store::{MlsStore, MlsStoreError},
    subscriptions::{SubscribeError, SyncWorkerEvent},
    worker::{NeedsDbReconnect, metrics::WorkerMetrics},
};
use owo_colors::OwoColorize;
use prost::Message;
use std::{collections::HashMap, sync::Arc};
use thiserror::Error;
use tokio::sync::broadcast::error::RecvError;
use tracing::instrument;
use worker::SyncMetric;
use xmtp_archive::ArchiveError;
use xmtp_common::ErrorCode;
use xmtp_common::{NS_IN_DAY, RetryableError, time::now_ns};
use xmtp_content_types::encoded_content_to_bytes;
use xmtp_db::tasks::NewTask;
use xmtp_db::{NotFound, StorageError, consent_record::ConsentState, group::GroupQueryArgs};
use xmtp_db::{XmtpDb, group::ConversationType, prelude::*};
use xmtp_id::{InboxIdRef, associations::DeserializationError};
use xmtp_mls_common::group::GroupMetadataOptions;
use xmtp_proto::types::{GroupId, InstallationId};
use xmtp_proto::xmtp::{
    device_sync::content::{
        DeviceSyncContent as DeviceSyncContentProto, device_sync_content::Content as ContentProto,
    },
    mls::{
        database::{
            AddMissingInstallations as AddMissingInstallationsProto, Task as TaskProto,
            task::Task as TaskKindProto,
        },
        message_contents::{
            ContentTypeId, EncodedContent, PlaintextEnvelope,
            plaintext_envelope::{Content, V1},
        },
    },
};

pub mod archive;
pub mod preference_sync;
pub mod worker;

pub use xmtp_archive::archive_options::{ArchiveOptions, BackupElementSelection};

#[cfg(test)]
mod tests;

#[derive(Debug, Error, ErrorCode)]
pub enum DeviceSyncError {
    /// I/O error.
    ///
    /// File system or network I/O failed. May be retryable.
    #[error("IO error: {0}")]
    IO(#[from] std::io::Error),
    /// Serialization error.
    ///
    /// JSON serialization/deserialization failed. Retryable.
    #[error("Serialization/Deserialization Error {0}")]
    Serde(#[from] serde_json::Error),
    #[error(transparent)]
    #[error_code(inherit)]
    ProtoConversion(#[from] xmtp_proto::ConversionError),
    /// AES-GCM encryption error.
    ///
    /// Encryption/decryption of sync payload failed. Retryable.
    #[error("AES-GCM encryption error")]
    AesGcm(#[from] aes_gcm::Error),
    #[error("storage error: {0}")]
    #[error_code(inherit)]
    Storage(#[from] StorageError),
    /// Type conversion error.
    ///
    /// Internal type conversion failed. Retryable.
    #[error("type conversion error")]
    Conversion,
    /// UTF-8 error.
    ///
    /// String is not valid UTF-8. Retryable.
    #[error("utf-8 error: {0}")]
    UTF8(#[from] std::str::Utf8Error),
    #[error("client error: {0}")]
    #[error_code(inherit)]
    Client(#[from] ClientError),
    #[error("group error: {0}")]
    #[error_code(inherit)]
    Group(#[from] GroupError),
    /// Invalid payload.
    ///
    /// Sync message payload is malformed. Retryable.
    #[error("invalid history message payload")]
    InvalidPayload,
    /// Unspecified sync kind.
    ///
    /// Device sync kind not specified. Not retryable.
    #[error("unspecified device sync kind")]
    UnspecifiedDeviceSyncKind,
    #[error(transparent)]
    #[error_code(inherit)]
    Subscribe(#[from] SubscribeError),
    /// Bincode error.
    ///
    /// Binary serialization failed. Retryable.
    #[error(transparent)]
    Bincode(#[from] bincode::Error),
    /// Archive error.
    ///
    /// Sync archive operation failed. Retryable.
    #[error(transparent)]
    Archive(#[from] ArchiveError),
    /// Decode error.
    ///
    /// Protobuf decoding failed. Retryable.
    #[error(transparent)]
    Decode(#[from] prost::DecodeError),
    #[error(transparent)]
    #[error_code(inherit)]
    Deserialization(#[from] DeserializationError),
    /// Missing sync group.
    ///
    /// Sync group not found. Not retryable.
    #[error("Missing sync group")]
    MissingSyncGroup,
    #[error(transparent)]
    #[error_code(inherit)]
    Db(#[from] xmtp_db::ConnectionError),
    /// Sync summary.
    ///
    /// Sync completed with errors. May be retryable.
    #[error("{}", _0.to_string())]
    Sync(Box<SyncSummary>),
    /// MLS store error.
    ///
    /// OpenMLS key store operation failed. Retryable.
    #[error(transparent)]
    MlsStore(#[from] MlsStoreError),
    /// Receive error.
    ///
    /// Channel receive failed. Retryable.
    #[error(transparent)]
    Recv(#[from] RecvError),
    /// Missing field.
    ///
    /// Required field not present. Retryable.
    #[error("Missing Field: {0:?} {1}")]
    MissingField(MissingField, String),
}

#[derive(Debug)]
pub enum MissingField {
    Conversation(ConversationField),
}
#[derive(Debug)]
pub enum ConversationField {
    DmId,
}

impl From<SyncSummary> for DeviceSyncError {
    fn from(value: SyncSummary) -> Self {
        DeviceSyncError::Sync(Box::new(value))
    }
}

impl NeedsDbReconnect for DeviceSyncError {
    fn needs_db_reconnect(&self) -> bool {
        match self {
            Self::Client(s) => s.db_needs_connection(),
            Self::Storage(s) => s.db_needs_connection(),
            // A dropped pool can hide in these wrapped errors; forward so the
            // worker stops instead of hot-looping (was `_ => false`).
            Self::Db(c) => c.db_needs_connection(),
            Self::Group(e) => e.needs_db_reconnect(),
            Self::MlsStore(e) => e.needs_db_reconnect(),
            Self::Subscribe(e) => e.needs_db_reconnect(),
            _ => false,
        }
    }
}

impl RetryableError for DeviceSyncError {
    fn is_retryable(&self) -> bool {
        !matches!(
            self,
            Self::MissingSyncGroup | Self::UnspecifiedDeviceSyncKind
        )
    }
}

impl From<NotFound> for DeviceSyncError {
    fn from(value: NotFound) -> Self {
        DeviceSyncError::Storage(StorageError::NotFound(value))
    }
}

#[derive(Clone)]
pub struct DeviceSyncClient<Context> {
    pub(crate) context: Context,
    pub(crate) mls_store: MlsStore<Context>,
    pub(crate) metrics: Arc<WorkerMetrics<SyncMetric>>,
}

impl<Context: XmtpSharedContext> DeviceSyncClient<Context> {
    pub fn new(context: Context, metrics: Arc<WorkerMetrics<SyncMetric>>) -> Self {
        Self {
            context: context.clone(),
            mls_store: MlsStore::new(context),
            metrics,
        }
    }
}

impl<Context> DeviceSyncClient<Context>
where
    Context: XmtpSharedContext,
{
    pub fn inbox_id(&self) -> InboxIdRef<'_> {
        self.context.identity().inbox_id()
    }

    pub fn installation_id(&self) -> InstallationId {
        self.context.installation_id()
    }

    pub fn db(&self) -> <Context::Db as XmtpDb>::DbQuery {
        self.context.db()
    }

    /// Blocks until the sync worker notifies that it is initialized and running.
    pub async fn wait_for_sync_worker_init(&self) -> Result<(), xmtp_common::time::Expired> {
        self.metrics.wait_for_init().await
    }

    /// Sends a device sync message.
    /// If the `group_id` is `None`, the message will be sent
    /// to the primary sync group ID.
    #[cfg_attr(any(test, feature = "test-utils"), tracing::instrument(level = "info", fields(inbox_id = self.context.inbox_id()), skip(self)))]
    #[cfg_attr(
        not(any(test, feature = "test-utils")),
        tracing::instrument(level = "trace", skip(self))
    )]
    async fn send_device_sync_message(
        &self,
        content: ContentProto,
    ) -> Result<Vec<u8>, ClientError> {
        let content = DeviceSyncContentProto {
            content: Some(content),
        };

        let sync_group = self.get_sync_group().await?;

        let msg = format!(
            "[{}] Sending sync message to group {:?}",
            self.context.installation_id(),
            xmtp_common::fmt::debug_hex(sync_group.group_id)
        );
        tracing::info!("{}", msg.yellow());

        let mut content_bytes = vec![];
        content
            .encode(&mut content_bytes)
            .map_err(|err| ClientError::Generic(err.to_string()))?;

        let encoded_content = EncodedContent {
            r#type: Some(ContentTypeId {
                authority_id: "xmtp.org".to_string(),
                type_id: "application/x-protobuf".to_string(),
                version_major: 1,
                version_minor: 0,
            }),
            parameters: HashMap::new(),
            fallback: None,
            compression: None,
            content: content_bytes,
        };
        let content_bytes = encoded_content_to_bytes(encoded_content);

        let message_id = sync_group.prepare_message(
            &content_bytes,
            send_message_opts::SendMessageOpts {
                should_push: false,
                idempotency_key: None,
            },
            |key| PlaintextEnvelope {
                content: Some(Content::V1(V1 {
                    content: content_bytes.clone(),
                    idempotency_key: key.to_string(),
                })),
            },
        )?;

        sync_group.sync_until_last_intent_resolved().await?;

        // Notify our own worker of our own message so it can process it.
        let _ = self
            .context
            .worker_events()
            .send(SyncWorkerEvent::NewSyncGroupMsg);

        Ok(message_id)
    }

    #[instrument(level = "trace", skip_all)]
    pub async fn get_sync_group(&self) -> Result<MlsGroup<Context>, GroupError> {
        let db = self.context.db();
        let sync_group = match db.primary_sync_group()? {
            Some(sync_group) => self.mls_store.group(&sync_group.id)?,
            None => {
                let sync_group = MlsGroup::create_and_insert(
                    self.context.clone(),
                    ConversationType::Sync,
                    PreconfiguredPolicies::default().to_policy_set(),
                    GroupMetadataOptions::default(),
                    None,
                )?;
                tracing::info!(
                    "[{}] Creating sync group: {}",
                    hex::encode(self.context.installation_id()),
                    hex::encode(sync_group.group_id)
                );
                if let Err(inline_err) = sync_group.add_missing_installations().await {
                    // The group row is already persisted, so this add is never
                    // re-attempted (later calls take the `Some` branch) — arm
                    // the durable reconcile task so the TaskRunner heals it,
                    // then surface the original error. Armed only on failure:
                    // enqueue-first duplicated the reconcile (and its identity
                    // fetch) on every sync-group creation.
                    self.schedule_add_missing_installations_task(sync_group.group_id)
                        .map_err(Box::new)?;
                    return Err(inline_err);
                }
                sync_group.sync_with_conn().await?;

                self.metrics.increment_metric(SyncMetric::SyncGroupCreated);

                sync_group
            }
        };

        Ok(sync_group)
    }

    /// This should be triggered when a new sync group appears,
    /// indicating the presence of a new installation.
    ///
    /// Schedules one durable AddMissingInstallations task per eligible group
    /// on the TaskRunner (deduped by payload hash) so the membership add is
    /// retried with backoff instead of being lost if a single inline attempt
    /// fails (e.g. identity propagation lag → MissingSequenceId).
    #[cfg_attr(
        any(test, feature = "test-utils"),
        tracing::instrument(level = "info", skip_all)
    )]
    pub fn schedule_add_installations_to_groups(&self) -> Result<usize, DeviceSyncError> {
        let groups = self.mls_store.find_groups(GroupQueryArgs {
            last_activity_after_ns: Some(now_ns() - NS_IN_DAY * 90),
            consent_states: Some(vec![ConsentState::Allowed, ConsentState::Unknown]),
            ..Default::default()
        })?;

        for group in &groups {
            self.schedule_add_missing_installations_task(group.group_id)?;
        }
        Ok(groups.len())
    }

    /// Durably enqueue one AddMissingInstallations task for `group_id`
    /// (deduped by payload hash against any pending row for the same group)
    /// and wake the TaskRunner.
    pub(crate) fn schedule_add_missing_installations_task(
        &self,
        group_id: GroupId,
    ) -> Result<(), DeviceSyncError> {
        let task = NewTask::builder()
            .originating_message_sequence_id(0)
            .build(TaskProto {
                task: Some(TaskKindProto::AddMissingInstallations(
                    AddMissingInstallationsProto {
                        group_id: group_id.to_vec(),
                    },
                )),
            })?;
        self.context.db().create_or_ignore_task(task)?;
        self.context.task_channels().wake();
        Ok(())
    }
}

/// Decode only device-sync content that this client supports.
///
/// Old archive transfer fields decode as an unset oneof after their field
/// numbers became reserved. Ignore them so older installations do not make
/// the sync worker fail or retry the message.
fn decode_supported_content(bytes: &[u8]) -> Option<ContentProto> {
    DeviceSyncContentProto::decode(bytes).ok()?.content
}