xmtp_mls_core/
hmac_epoch.rs

1//! Report HMAC epoch changes observed while this client is open.
2
3use crate::context::CoreContext;
4use crate::worker::{BoxedWorker, NewWorker, Worker, WorkerFactory, WorkerKind, WorkerResult};
5use std::sync::{
6    Arc,
7    atomic::{AtomicI64, Ordering},
8};
9use xmtp_common::time::{self, Duration};
10use xmtp_events::{ClientEvent, EventWriter, HmacKeysUpdated};
11
12/// Current hmac epoch. HMAC keys change every 30 days
13pub fn hmac_epoch() -> i64 {
14    xmtp_push_types::hmac_epoch(xmtp_common::time::now_secs())
15}
16
17pub struct HmacEpochWorker<Context> {
18    context: Context,
19    last_epoch: Arc<AtomicI64>,
20    #[cfg(test)]
21    clock: Option<TestClock>,
22}
23
24#[cfg(test)]
25struct TestClock {
26    now: Arc<AtomicI64>,
27    ticks: futures::channel::mpsc::UnboundedReceiver<i64>,
28    waits: futures::channel::mpsc::UnboundedSender<Duration>,
29}
30
31impl<Context> HmacEpochWorker<Context> {
32    fn now(&self) -> i64 {
33        #[cfg(test)]
34        if let Some(clock) = &self.clock {
35            return clock.now.load(Ordering::Acquire);
36        }
37        time::now_secs()
38    }
39
40    async fn wait(&mut self, duration: Duration) {
41        #[cfg(test)]
42        if let Some(clock) = &mut self.clock {
43            use futures::StreamExt;
44            clock
45                .waits
46                .unbounded_send(duration)
47                .expect("test wait observer must stay open");
48            let now = clock.ticks.next().await.expect("test clock must stay open");
49            clock.now.store(now, Ordering::Release);
50            return;
51        }
52        time::sleep(duration).await;
53    }
54
55    fn wait_until_next_epoch(&self, now: i64) -> Duration {
56        let next =
57            (self.last_epoch.load(Ordering::Acquire) + 1) * xmtp_push_types::HMAC_EPOCH_SECONDS;
58        Duration::from_secs((next - now).max(1) as u64)
59    }
60
61    fn observe_epoch(&mut self, current: i64) -> bool {
62        loop {
63            let previous = self.last_epoch.load(Ordering::Acquire);
64            if current <= previous {
65                return false;
66            }
67            if self
68                .last_epoch
69                .compare_exchange(previous, current, Ordering::AcqRel, Ordering::Acquire)
70                .is_ok()
71            {
72                return true;
73            }
74        }
75    }
76}
77
78impl<Context: CoreContext> HmacEpochWorker<Context> {
79    // implements: EVENT-026
80    fn report_epoch(&mut self, observed: i64) {
81        if self.observe_epoch(observed) {
82            self.context.task_channels().mark_notification_changed();
83            self.context
84                .events()
85                .emit(Some(ClientEvent::HmacKeysUpdated(HmacKeysUpdated)), None);
86        }
87    }
88}
89
90struct Factory<Context> {
91    context: Context,
92    last_epoch: Arc<AtomicI64>,
93}
94
95impl<Context> Factory<Context> {
96    fn new(context: Context, startup_epoch: i64) -> Self {
97        Self {
98            context,
99            last_epoch: Arc::new(AtomicI64::new(startup_epoch)),
100        }
101    }
102}
103
104impl<Context: CoreContext + 'static> WorkerFactory for Factory<Context> {
105    fn kind(&self) -> WorkerKind {
106        WorkerKind::HmacEpoch
107    }
108
109    fn create(&self) -> BoxedWorker {
110        Box::new(HmacEpochWorker {
111            context: self.context.clone(),
112            last_epoch: self.last_epoch.clone(),
113            #[cfg(test)]
114            clock: None,
115        })
116    }
117}
118
119impl<Context: CoreContext + 'static> NewWorker<Context> for HmacEpochWorker<Context> {
120    fn factory(context: Context) -> impl WorkerFactory + 'static {
121        Factory::new(context, hmac_epoch())
122    }
123}
124
125#[xmtp_common::async_trait]
126impl<Context: CoreContext + 'static> Worker for HmacEpochWorker<Context> {
127    fn kind(&self) -> WorkerKind {
128        WorkerKind::HmacEpoch
129    }
130
131    async fn run_tasks(&mut self) -> WorkerResult<()> {
132        loop {
133            let now = self.now();
134            self.wait(self.wait_until_next_epoch(now)).await;
135            self.report_epoch(self.now() / xmtp_push_types::HMAC_EPOCH_SECONDS);
136        }
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use crate::test_context::real_context;
144    use xmtp_events::{EventFilter, EventKind};
145    use xmtp_proto::api::HasStats;
146
147    #[xmtp_common::test(unwrap_try = true)]
148    async fn reports_only_new_epochs_after_startup() {
149        let mut worker = HmacEpochWorker {
150            context: (),
151            last_epoch: Arc::new(AtomicI64::new(10)),
152            clock: None,
153        };
154        assert!(!worker.observe_epoch(10));
155        assert!(worker.observe_epoch(11));
156        assert!(!worker.observe_epoch(11));
157        assert!(worker.observe_epoch(12));
158    }
159
160    // verifies: EVENT-001
161    #[xmtp_common::test(unwrap_try = true)]
162    async fn epoch_timer_turn_emits_once_and_invalidates_notifications() {
163        let context = real_context().await;
164        let current = hmac_epoch();
165        let mut worker = HmacEpochWorker {
166            context: context.clone(),
167            last_epoch: Arc::new(AtomicI64::new(current)),
168            clock: None,
169        };
170        let events = context
171            .events()
172            .subscribe_app(EventFilter::new([EventKind::HmacKeysUpdated]))?;
173        let boundary = (current + 1) * xmtp_push_types::HMAC_EPOCH_SECONDS;
174        assert_eq!(
175            worker.wait_until_next_epoch(boundary - 1),
176            Duration::from_secs(1)
177        );
178        let previous_revision = context.task_channels().notification_revision();
179        worker.report_epoch(current);
180        assert!(events.drain().is_empty());
181        worker.report_epoch(current + 1);
182        assert!(matches!(
183            events.drain().as_slice(),
184            [xmtp_events::EventEnvelope {
185                client: Some(ClientEvent::HmacKeysUpdated(_)),
186                ..
187            }]
188        ));
189        assert_eq!(
190            context.task_channels().notification_revision(),
191            previous_revision + 1
192        );
193        worker.report_epoch(current + 1);
194        assert!(events.drain().is_empty());
195    }
196    // verifies: EVENT-026
197    #[xmtp_common::test(unwrap_try = true)]
198    async fn running_epoch_worker_skips_closed_epochs_and_uses_no_network() {
199        let context = real_context().await;
200        let events = context
201            .events()
202            .subscribe_app(EventFilter::new([EventKind::HmacKeysUpdated]))?;
203        let stats = context.api().api_client.mls_stats();
204        let identity = context.api().api_client.identity_stats();
205        let counts = || {
206            [
207                stats.publish.get_count(),
208                stats.create_upload.get_count(),
209                stats.query.get_count(),
210                stats.query_newest.get_count(),
211                stats.subscribe.get_count(),
212                stats.subscribe_static.get_count(),
213                identity.get_inbox_ids.get_count(),
214                identity.verify_smart_contract_wallet_signatures.get_count(),
215            ]
216        };
217        let before = counts();
218        let revision = context.task_channels().notification_revision();
219        let width = xmtp_push_types::HMAC_EPOCH_SECONDS;
220        let now = Arc::new(AtomicI64::new(10 * width));
221        let (ticks, receiver) = futures::channel::mpsc::unbounded();
222        let (waits, mut scheduled) = futures::channel::mpsc::unbounded();
223        let factory = Factory::new(context.clone(), 10);
224        let mut worker = HmacEpochWorker {
225            context: factory.context.clone(),
226            last_epoch: factory.last_epoch.clone(),
227            clock: Some(TestClock {
228                now: now.clone(),
229                ticks: receiver,
230                waits,
231            }),
232        };
233        let mut running = Box::pin(worker.run_tasks());
234        assert!(futures::poll!(&mut running).is_pending());
235        assert!(events.drain().is_empty());
236        assert_eq!(scheduled.try_recv()?, Duration::from_secs(width as u64));
237        let mut previous = 10;
238        for (epoch, expected) in [(10, 0), (11, 1), (11, 0), (14, 1), (13, 0), (14, 0)] {
239            ticks.unbounded_send(epoch * width)?;
240            assert!(futures::poll!(&mut running).is_pending());
241            let emitted = events.drain();
242            assert_eq!(emitted.len(), expected);
243            previous = previous.max(epoch);
244            assert_eq!(
245                scheduled.try_recv()?,
246                Duration::from_secs(((previous + 1 - epoch) * width) as u64)
247            );
248            assert!(
249                emitted
250                    .iter()
251                    .all(|event| matches!(event.client, Some(ClientEvent::HmacKeysUpdated(_))))
252            );
253        }
254        drop(running);
255        assert_eq!(
256            context.task_channels().notification_revision(),
257            revision + 2
258        );
259        now.store(18 * width, Ordering::Release);
260        let (ticks, receiver) = futures::channel::mpsc::unbounded();
261        let (waits, mut scheduled) = futures::channel::mpsc::unbounded();
262        let factory = Factory::new(context.clone(), 18);
263        let mut reopened = HmacEpochWorker {
264            context: factory.context.clone(),
265            last_epoch: factory.last_epoch.clone(),
266            clock: Some(TestClock {
267                now,
268                ticks: receiver,
269                waits,
270            }),
271        };
272        let mut running = Box::pin(reopened.run_tasks());
273        assert!(futures::poll!(&mut running).is_pending());
274        assert_eq!(scheduled.try_recv()?, Duration::from_secs(width as u64));
275        ticks.unbounded_send(18 * width)?;
276        assert!(futures::poll!(&mut running).is_pending());
277        assert!(events.drain().is_empty());
278        ticks.unbounded_send(19 * width)?;
279        assert!(futures::poll!(&mut running).is_pending());
280        assert_eq!(events.drain().len(), 1);
281        assert_eq!(
282            context.task_channels().notification_revision(),
283            revision + 3
284        );
285        assert_eq!(counts(), before);
286    }
287}
xmtp_mls_core/hmac_epoch.rs

xmtp_mls_core/
hmac_epoch.rs
