xmtp_mls/worker/device_sync/
preference_sync.rs

1use super::*;
2use xmtp_common::time::now_ns;
3use xmtp_db::consent_record::StoredConsentRecord;
4use xmtp_db::user_preferences::{HmacKey, StoredUserPreferences};
5use xmtp_proto::ConversionError;
6use xmtp_proto::xmtp::device_sync::content::HmacKeyUpdate as HmacKeyUpdateProto;
7use xmtp_proto::xmtp::device_sync::content::{
8    PreferenceUpdate as PreferenceUpdateProto, PreferenceUpdates,
9    device_sync_content::Content as ContentProto, preference_update::Update as UpdateProto,
10};
11
12#[derive(Clone, Debug, PartialEq)]
13pub enum PreferenceUpdate {
14    Consent(StoredConsentRecord),
15    Hmac { key: Vec<u8>, cycled_at_ns: i64 },
16}
17
18impl<Context> DeviceSyncClient<Context>
19where
20    Context: XmtpSharedContext,
21{
22    pub(crate) async fn sync_preferences(
23        &self,
24        updates: Vec<PreferenceUpdate>,
25    ) -> Result<Vec<PreferenceUpdate>, ClientError> {
26        self.send_device_sync_message(ContentProto::PreferenceUpdates(PreferenceUpdates {
27            updates: updates.clone().into_iter().map(From::from).collect(),
28        }))
29        .await?;
30
31        updates.iter().for_each(|update| match update {
32            PreferenceUpdate::Consent(_) => self.metrics.increment_metric(SyncMetric::ConsentSent),
33            PreferenceUpdate::Hmac { .. } => self.metrics.increment_metric(SyncMetric::HmacSent),
34        });
35
36        Ok(updates)
37    }
38
39    pub(crate) async fn cycle_hmac(&self) -> Result<(), ClientError> {
40        tracing::info!(
41            "[{}] Sending new HMAC key to sync group.",
42            self.context.installation_id()
43        );
44
45        self.sync_preferences(vec![PreferenceUpdate::Hmac {
46            key: HmacKey::random_key(),
47            cycled_at_ns: now_ns(),
48        }])
49        .await?;
50
51        Ok(())
52    }
53}
54
55pub(super) fn store_preference_updates(
56    updates: Vec<PreferenceUpdateProto>,
57    conn: &impl DbQuery,
58    handle: &WorkerMetrics<SyncMetric>,
59) -> Result<Vec<PreferenceUpdate>, StorageError> {
60    let mut changed = vec![];
61    for update in updates.into_iter().filter_map(|u| u.update) {
62        match update {
63            UpdateProto::Consent(consent_save) => {
64                tracing::info!(
65                    "Storing consent update from sync group. State: {:?}",
66                    consent_save.state
67                );
68
69                let consent_record: StoredConsentRecord = consent_save.try_into()?;
70                let updated = conn.insert_newer_consent_record(consent_record.clone())?;
71
72                if updated {
73                    changed.push(PreferenceUpdate::Consent(consent_record));
74                }
75
76                handle.increment_metric(SyncMetric::ConsentReceived);
77            }
78            UpdateProto::Hmac(HmacKeyUpdateProto { key, cycled_at_ns }) => {
79                tracing::info!("Storing new HMAC key from sync group");
80                StoredUserPreferences::store_hmac_key(conn, &key, Some(cycled_at_ns))?;
81                changed.push(PreferenceUpdate::Hmac { key, cycled_at_ns });
82                handle.increment_metric(SyncMetric::HmacReceived);
83            }
84        }
85    }
86
87    Ok(changed)
88}
89
90impl TryFrom<PreferenceUpdateProto> for PreferenceUpdate {
91    type Error = ConversionError;
92    fn try_from(update: PreferenceUpdateProto) -> Result<Self, Self::Error> {
93        let Some(update) = update.update else {
94            return Err(ConversionError::Unspecified("update"));
95        };
96        update.try_into()
97    }
98}
99impl TryFrom<UpdateProto> for PreferenceUpdate {
100    type Error = ConversionError;
101    fn try_from(update: UpdateProto) -> Result<Self, Self::Error> {
102        let update = match update {
103            UpdateProto::Consent(consent) => Self::Consent(consent.try_into()?),
104            UpdateProto::Hmac(HmacKeyUpdateProto { key, cycled_at_ns }) => {
105                Self::Hmac { key, cycled_at_ns }
106            }
107        };
108        Ok(update)
109    }
110}
111
112impl From<PreferenceUpdate> for PreferenceUpdateProto {
113    fn from(update: PreferenceUpdate) -> Self {
114        PreferenceUpdateProto {
115            update: Some(match update {
116                PreferenceUpdate::Consent(consent) => UpdateProto::Consent(consent.into()),
117                PreferenceUpdate::Hmac { key, cycled_at_ns } => {
118                    UpdateProto::Hmac(HmacKeyUpdateProto { key, cycled_at_ns })
119                }
120            }),
121        }
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use crate::{tester, worker::device_sync::worker::SyncMetric};
128    use xmtp_db::user_preferences::StoredUserPreferences;
129
130    #[rstest::rstest]
131    #[xmtp_common::test(unwrap_try = true)]
132    async fn test_hmac_sync() {
133        tester!(amal_a, sync_worker);
134        tester!(amal_b, from: amal_a);
135
136        amal_a.test_has_same_sync_group_as(&amal_b).await?;
137
138        amal_a
139            .worker()
140            .register_interest(SyncMetric::HmacSent, 1)
141            .wait()
142            .await?;
143
144        amal_a.sync_all_welcomes_and_device_sync_groups().await?;
145        amal_a
146            .worker()
147            .register_interest(SyncMetric::HmacReceived, 1)
148            .wait()
149            .await?;
150
151        // Wait for a to process the new hmac key
152        amal_b
153            .context
154            .device_sync_client()
155            .get_sync_group()
156            .await?
157            .sync()
158            .await?;
159        amal_b
160            .worker()
161            .register_interest(SyncMetric::HmacReceived, 1)
162            .wait()
163            .await?;
164
165        let pref_a = StoredUserPreferences::load(amal_a.context.db())?;
166        let pref_b = StoredUserPreferences::load(amal_b.context.db())?;
167
168        assert_eq!(pref_a.hmac_key, pref_b.hmac_key);
169
170        amal_a
171            .identity_updates()
172            .revoke_installations(vec![amal_b.context.installation_id().to_vec()])
173            .await?;
174
175        amal_a.sync_all_welcomes_and_device_sync_groups().await?;
176        amal_a
177            .worker()
178            .register_interest(SyncMetric::HmacReceived, 2)
179            .wait()
180            .await?;
181        let new_pref_a = StoredUserPreferences::load(amal_a.context.db())?;
182        assert_ne!(pref_a.hmac_key, new_pref_a.hmac_key);
183    }
184}
xmtp_mls/worker/device_sync/preference_sync.rs

xmtp_mls/worker/device_sync/
preference_sync.rs
