context_id/
lib.rs

Help
1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
4
5mod error;
6
7pub use error::{ApiError, ApiResult, Error, Result};
8
9use chrono::{DateTime, Duration, Utc};
10use error_support::{error, handle_error};
11use parking_lot::{Mutex, RwLock};
12use uuid::Uuid;
13
14uniffi::setup_scaffolding!("context_id");
15
16mod callback;
17pub use callback::{ContextIdCallback, DefaultContextIdCallback};
18
19mod mars;
20use mars::{MARSClient, SimpleMARSClient};
21
22/// Top-level API for the context_id component
23#[derive(uniffi::Object)]
24pub struct ContextIDComponent {
25    inner: Mutex<ContextIDComponentInner>,
26}
27
28#[uniffi::export]
29impl ContextIDComponent {
30    /// Construct a new [ContextIDComponent].
31    ///
32    /// If no creation timestamp is provided, the current time will be used.
33    #[uniffi::constructor]
34
35    pub fn new(
36        init_context_id: &str,
37        creation_timestamp_s: i64,
38        running_in_test_automation: bool,
39        callback: Box<dyn ContextIdCallback>,
40    ) -> Self {
41        Self {
42            inner: Mutex::new(ContextIDComponentInner::new(
43                init_context_id,
44                creation_timestamp_s,
45                running_in_test_automation,
46                callback,
47                Utc::now(),
48                Box::new(SimpleMARSClient::new()),
49            )),
50        }
51    }
52
53    /// Return the current context ID string.
54    #[handle_error(Error)]
55    pub fn request(&self, rotation_days_in_s: u8) -> ApiResult<String> {
56        let mut inner = self.inner.lock();
57        inner.request(rotation_days_in_s, Utc::now())
58    }
59
60    /// Regenerate the context ID.
61    #[handle_error(Error)]
62    pub fn force_rotation(&self) -> ApiResult<()> {
63        let mut inner = self.inner.lock();
64        inner.force_rotation(Utc::now());
65        Ok(())
66    }
67
68    /// Unset the callbacks set during construction, and use a default
69    /// no-op ContextIdCallback instead.
70    #[handle_error(Error)]
71    pub fn unset_callback(&self) -> ApiResult<()> {
72        let mut inner = self.inner.lock();
73        inner.unset_callback();
74        Ok(())
75    }
76}
77
78struct ContextIDComponentInner {
79    context_id: String,
80    creation_timestamp: DateTime<Utc>,
81    callback_handle: RwLock<Box<dyn ContextIdCallback>>,
82    mars_client: Box<dyn MARSClient>,
83    running_in_test_automation: bool,
84}
85
86impl ContextIDComponentInner {
87    pub fn new(
88        init_context_id: &str,
89        creation_timestamp_s: i64,
90        running_in_test_automation: bool,
91        callback: Box<dyn ContextIdCallback>,
92        now: DateTime<Utc>,
93        mars_client: Box<dyn MARSClient>,
94    ) -> Self {
95        // Some historical context IDs are stored within opening and closing
96        // braces, and our endpoints have tolerated this, but ideally we'd
97        // send without the braces, so we strip any off here.
98        let (context_id, generated_context_id) = match init_context_id
99            .trim()
100            .trim_start_matches('{')
101            .trim_end_matches('}')
102        {
103            "" => (Uuid::new_v4().to_string(), true),
104            // If the passed in string isn't empty, but still not a valid UUID,
105            // just go ahead and generate a new UUID.
106            s => match Uuid::parse_str(s) {
107                Ok(_) => (s.to_string(), false),
108                Err(_) => (Uuid::new_v4().to_string(), true),
109            },
110        };
111
112        let (creation_timestamp, generated_creation_timestamp, force_rotation) =
113            match (generated_context_id, creation_timestamp_s) {
114                // We generated a new context ID then we need a fresh timestamp and no rotation needed
115                (true, _) => (now, true, false),
116                // Pre-existing context ID with a positive timestamp:
117                // try parsing it (any real UNIX-epoch timestamp is orders of magnitude within chrono’s 262_000-year range),
118                // and if it’s somehow out-of-range fall back to `now`  and force rotation
119                // See: https://docs.rs/chrono/latest/chrono/naive/struct.NaiveDateTime.html#panics
120                (false, secs) if secs > 0 => DateTime::<Utc>::from_timestamp(secs, 0)
121                    .map(|ts| (ts, false, false))
122                    .unwrap_or((now, true, true)),
123                // Pre-existing context ID with zero timestamp then use current time and no rotation needed
124                (false, 0) => (now, true, false),
125                // Pre-existing context ID but INVALID timestamp then use current time but FORCE rotation
126                (false, _) => (now, true, true),
127            };
128
129        let mut instance = Self {
130            context_id,
131            creation_timestamp,
132            callback_handle: RwLock::new(callback),
133            mars_client,
134            running_in_test_automation,
135        };
136
137        if force_rotation {
138            // force_rotation will cause a implicit persist
139            instance.force_rotation(creation_timestamp)
140        } else if generated_context_id || generated_creation_timestamp {
141            // We only need to persist these if we just generated one.
142            instance.persist();
143        }
144
145        instance
146    }
147
148    pub fn request(&mut self, rotation_days: u8, now: DateTime<Utc>) -> Result<String> {
149        if rotation_days == 0 {
150            return Ok(self.context_id.clone());
151        }
152
153        let age = now - self.creation_timestamp;
154        if age >= Duration::days(rotation_days.into()) {
155            self.rotate_context_id(now);
156        }
157
158        Ok(self.context_id.clone())
159    }
160
161    fn rotate_context_id(&mut self, now: DateTime<Utc>) {
162        let original_context_id = self.context_id.clone();
163
164        self.context_id = Uuid::new_v4().to_string();
165        self.creation_timestamp = now;
166        self.persist();
167
168        // If we're running in test automation in the embedder, we don't want
169        // to be sending actual network requests to MARS.
170        if !self.running_in_test_automation {
171            let _ = self
172                .mars_client
173                .delete(original_context_id.clone())
174                .inspect_err(|e| error!("Failed to contact MARS: {}", e));
175        }
176
177        // In a perfect world, we'd call Glean ourselves here - however,
178        // our uniffi / Rust infrastructure doesn't yet support doing that,
179        // so we'll delegate to our embedder to send the Glean ping by
180        // calling a `rotated` callback method.
181        self.callback_handle
182            .read()
183            .rotated(original_context_id.clone());
184    }
185
186    pub fn force_rotation(&mut self, now: DateTime<Utc>) {
187        self.rotate_context_id(now);
188    }
189
190    pub fn persist(&self) {
191        self.callback_handle
192            .read()
193            .persist(self.context_id.clone(), self.creation_timestamp.timestamp());
194    }
195
196    pub fn unset_callback(&mut self) {
197        self.callback_handle = RwLock::new(Box::new(DefaultContextIdCallback));
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204    use crate::callback::utils::SpyCallback;
205    use chrono::Utc;
206    use lazy_static::lazy_static;
207
208    use std::sync::{Arc, Mutex};
209
210    // 1745859061 ~= the timestamp for when this test was written (Apr 28, 2025)
211    const FAKE_NOW_TS: i64 = 1745859061;
212    const TEST_CONTEXT_ID: &str = "decafbad-0cd1-0cd2-0cd3-decafbad1000";
213    // 1706763600 ~= Jan 1st, 2024, which is long ago compared to FAKE_NOW.
214    const FAKE_LONG_AGO_TS: i64 = 1706763600;
215
216    lazy_static! {
217        static ref FAKE_NOW: DateTime<Utc> = DateTime::from_timestamp(FAKE_NOW_TS, 0).unwrap();
218        static ref FAKE_LONG_AGO: DateTime<Utc> =
219            DateTime::from_timestamp(FAKE_LONG_AGO_TS, 0).unwrap();
220    }
221
222    pub struct TestMARSClient {
223        delete_called: Arc<Mutex<bool>>,
224    }
225
226    impl TestMARSClient {
227        pub fn new(delete_called: Arc<Mutex<bool>>) -> Self {
228            Self { delete_called }
229        }
230    }
231
232    impl MARSClient for TestMARSClient {
233        fn delete(&self, _context_id: String) -> crate::Result<()> {
234            *self.delete_called.lock().unwrap() = true;
235            Ok(())
236        }
237    }
238
239    fn with_test_mars<F: FnOnce(Box<dyn MARSClient + Send + Sync>, Arc<Mutex<bool>>)>(test: F) {
240        let delete_called = Arc::new(Mutex::new(false));
241        let mars = Box::new(TestMARSClient::new(Arc::clone(&delete_called)));
242        test(mars, delete_called);
243    }
244
245    #[test]
246    fn test_creation_timestamp_with_some_value() {
247        with_test_mars(|mars, delete_called| {
248            let spy = SpyCallback::new();
249            let creation_timestamp = FAKE_NOW_TS;
250            let component = ContextIDComponentInner::new(
251                TEST_CONTEXT_ID,
252                creation_timestamp,
253                false,
254                spy.callback,
255                *FAKE_NOW,
256                mars,
257            );
258
259            // We should have left the context_id and creation_timestamp
260            // untouched if a creation_timestamp was passed.
261            assert_eq!(component.context_id, TEST_CONTEXT_ID);
262            assert_eq!(component.creation_timestamp.timestamp(), creation_timestamp);
263            assert!(!*delete_called.lock().unwrap());
264            // Neither persist nor rotate should have been called
265            assert_eq!(
266                *spy.persist_called.lock().unwrap(),
267                0,
268                "persist() should NOTE have been called"
269            );
270            assert_eq!(
271                *spy.rotated_called.lock().unwrap(),
272                0,
273                "rotated() should NOT have been called"
274            );
275        });
276    }
277
278    #[test]
279    fn test_creation_timestamp_with_zero_value() {
280        with_test_mars(|mars, delete_called| {
281            let spy = SpyCallback::new();
282            let component = ContextIDComponentInner::new(
283                TEST_CONTEXT_ID,
284                0,
285                false,
286                spy.callback,
287                *FAKE_NOW,
288                mars,
289            );
290
291            // If 0 was passed as the creation_timestamp, we'll interpret that
292            // as there having been no stored creation_timestamp. In that case,
293            // we'll use "now" as the creation_timestamp.
294            assert_eq!(component.context_id, TEST_CONTEXT_ID);
295            assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
296            assert!(!*delete_called.lock().unwrap());
297            // zero‐timestamp should trigger a persist, but NOT a rotation
298            assert_eq!(
299                *spy.persist_called.lock().unwrap(),
300                1,
301                "persist() should have been called once"
302            );
303            assert_eq!(
304                *spy.rotated_called.lock().unwrap(),
305                0,
306                "rotated() should NOT have been called"
307            );
308        });
309    }
310
311    #[test]
312    fn test_empty_initial_context_id() {
313        with_test_mars(|mars, delete_called| {
314            let spy = SpyCallback::new();
315            let component =
316                ContextIDComponentInner::new("", 0, false, spy.callback, *FAKE_NOW, mars);
317
318            // We expect a new UUID to have been generated for context_id.
319            assert!(Uuid::parse_str(&component.context_id).is_ok());
320            assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
321            assert_eq!(
322                *spy.persist_called.lock().unwrap(),
323                1,
324                "persist() should have been called once"
325            );
326            assert_eq!(
327                *spy.rotated_called.lock().unwrap(),
328                0,
329                "rotated() should NOT have been called"
330            );
331            assert!(!*delete_called.lock().unwrap());
332        });
333    }
334
335    #[test]
336    fn test_empty_initial_context_id_with_creation_date() {
337        with_test_mars(|mars, delete_called| {
338            let spy = SpyCallback::new();
339            let component =
340                ContextIDComponentInner::new("", 0, false, spy.callback, *FAKE_NOW, mars);
341
342            // We expect a new UUID to have been generated for context_id.
343            assert!(Uuid::parse_str(&component.context_id).is_ok());
344            assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
345            assert!(!*delete_called.lock().unwrap());
346            assert_eq!(
347                *spy.persist_called.lock().unwrap(),
348                1,
349                "persist() should have been called once"
350            );
351            assert_eq!(
352                *spy.rotated_called.lock().unwrap(),
353                0,
354                "rotated() should NOT have been called"
355            );
356        });
357    }
358
359    #[test]
360    fn test_invalid_context_id_with_no_creation_date() {
361        with_test_mars(|mars, delete_called| {
362            let spy = SpyCallback::new();
363            let component = ContextIDComponentInner::new(
364                "something-invalid",
365                0,
366                false,
367                spy.callback,
368                *FAKE_NOW,
369                mars,
370            );
371
372            // We expect a new UUID to have been generated for context_id.
373            assert!(Uuid::parse_str(&component.context_id).is_ok());
374            assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
375            // Also expect a persist to have been called, but not a rotation.
376            assert_eq!(
377                *spy.persist_called.lock().unwrap(),
378                1,
379                "persist() should have been called once"
380            );
381            assert_eq!(
382                *spy.rotated_called.lock().unwrap(),
383                0,
384                "rotated() should NOT have been called"
385            );
386            assert!(!*delete_called.lock().unwrap());
387        });
388    }
389
390    #[test]
391    fn test_invalid_context_id_with_creation_date() {
392        with_test_mars(|mars, delete_called| {
393            let spy = SpyCallback::new();
394            let component = ContextIDComponentInner::new(
395                "something-invalid",
396                FAKE_LONG_AGO_TS,
397                false,
398                spy.callback,
399                *FAKE_NOW,
400                mars,
401            );
402
403            // We expect a new UUID to have been generated for context_id.
404            assert!(Uuid::parse_str(&component.context_id).is_ok());
405            assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
406            // Also expect a persist to have been called, but not a rotation.
407            assert_eq!(
408                *spy.persist_called.lock().unwrap(),
409                1,
410                "persist() should have been called once"
411            );
412            assert_eq!(
413                *spy.rotated_called.lock().unwrap(),
414                0,
415                "rotated() should NOT have been called"
416            );
417            assert!(!*delete_called.lock().unwrap());
418        });
419    }
420
421    #[test]
422    fn test_context_id_with_invalid_creation_date() {
423        with_test_mars(|mars, delete_called| {
424            let spy = SpyCallback::new();
425            let component = ContextIDComponentInner::new(
426                TEST_CONTEXT_ID,
427                -1,
428                false,
429                spy.callback,
430                *FAKE_NOW,
431                mars,
432            );
433
434            // A new UUID should have been generated.
435            assert!(Uuid::parse_str(&component.context_id).is_ok());
436            assert_ne!(component.context_id, TEST_CONTEXT_ID);
437
438            // The creation timestamp must have been set to now.
439            assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
440
441            // Also expect a persist to have been called and a rotation.
442            assert_eq!(
443                *spy.persist_called.lock().unwrap(),
444                1,
445                "persist() should have been called once"
446            );
447            assert_eq!(
448                *spy.rotated_called.lock().unwrap(),
449                1,
450                "rotated() should have been called once"
451            );
452            // Since we forced a rotation, the MARS delete should have been called.
453            assert!(*delete_called.lock().unwrap());
454        });
455    }
456
457    #[test]
458    fn test_context_id_with_out_of_range_creation_date() {
459        with_test_mars(|mars, delete_called| {
460            let spy = SpyCallback::new();
461            // Way beyond chrono’s 262_000-year range
462            let huge_secs = i64::MAX;
463            let component = ContextIDComponentInner::new(
464                TEST_CONTEXT_ID,
465                huge_secs,
466                false,
467                spy.callback,
468                *FAKE_NOW,
469                mars,
470            );
471
472            // A new UUID should have been generated.
473            assert!(Uuid::parse_str(&component.context_id).is_ok());
474            assert_ne!(component.context_id, TEST_CONTEXT_ID);
475
476            // The creation timestamp must have been set to now.
477            assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
478
479            // Also expect a persist to have been called and a rotation.
480            assert_eq!(
481                *spy.persist_called.lock().unwrap(),
482                1,
483                "persist() should have been called once"
484            );
485            assert_eq!(
486                *spy.rotated_called.lock().unwrap(),
487                1,
488                "rotated() should have been called once"
489            );
490            // Since we forced a rotation on out-of-range, the MARS delete should have been called.
491            assert!(*delete_called.lock().unwrap());
492        });
493    }
494
495    #[test]
496    fn test_request_no_rotation() {
497        with_test_mars(|mars, delete_called| {
498            // Let's create a context_id with a creation date far in the past.
499            let mut component = ContextIDComponentInner::new(
500                TEST_CONTEXT_ID,
501                FAKE_LONG_AGO_TS,
502                false,
503                Box::new(DefaultContextIdCallback),
504                *FAKE_NOW,
505                mars,
506            );
507
508            // We expect neither the UUID nor creation_timestamp to have been changed.
509            assert_eq!(component.context_id, TEST_CONTEXT_ID);
510            assert_eq!(component.creation_timestamp, FAKE_LONG_AGO.clone());
511
512            // Now request the context_id, passing 0 for the rotation_days. We
513            // interpret this to mean "do not rotate".
514            assert_eq!(
515                component.request(0, *FAKE_NOW).unwrap(),
516                component.context_id
517            );
518            assert_eq!(component.creation_timestamp, FAKE_LONG_AGO.clone());
519            assert!(!*delete_called.lock().unwrap());
520        });
521    }
522
523    #[test]
524    fn test_request_with_rotation() {
525        with_test_mars(|mars, delete_called| {
526            let spy = SpyCallback::new();
527            // Let's create a context_id with a creation date far in the past.
528            let mut component = ContextIDComponentInner::new(
529                TEST_CONTEXT_ID,
530                FAKE_LONG_AGO_TS,
531                false,
532                spy.callback,
533                *FAKE_NOW,
534                mars,
535            );
536
537            // We expect neither the UUID nor creation_timestamp to have been changed.
538            assert_eq!(component.context_id, TEST_CONTEXT_ID);
539            assert_eq!(component.creation_timestamp, FAKE_LONG_AGO.clone());
540
541            // Now request the context_id, passing 2 for the rotation_days. Since
542            // the number of days since FAKE_LONG_AGO is greater than 2 days, we
543            // expect a new context_id to be generated, and the creation_timestamp
544            // to update to now.
545            assert!(Uuid::parse_str(&component.request(2, *FAKE_NOW).unwrap()).is_ok());
546            assert_ne!(component.context_id, TEST_CONTEXT_ID);
547            assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
548            assert!(*delete_called.lock().unwrap());
549
550            assert_eq!(
551                *spy.rotated_called.lock().unwrap(),
552                1,
553                "rotated() should have been called"
554            );
555            assert_eq!(
556                *spy.persist_called.lock().unwrap(),
557                1,
558                "persist() should have been called"
559            );
560            assert_eq!(
561                spy.last_rotated_id.lock().unwrap().as_deref().unwrap(),
562                TEST_CONTEXT_ID
563            );
564        });
565    }
566
567    #[test]
568    fn test_force_rotation() {
569        with_test_mars(|mars, delete_called| {
570            let spy = SpyCallback::new();
571
572            let mut component = ContextIDComponentInner::new(
573                TEST_CONTEXT_ID,
574                FAKE_LONG_AGO_TS,
575                false,
576                spy.callback,
577                *FAKE_NOW,
578                mars,
579            );
580
581            component.force_rotation(*FAKE_NOW);
582            assert!(Uuid::parse_str(&component.request(2, *FAKE_NOW).unwrap()).is_ok());
583            assert_ne!(component.context_id, TEST_CONTEXT_ID);
584            assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
585            assert!(*delete_called.lock().unwrap());
586            assert_eq!(
587                *spy.rotated_called.lock().unwrap(),
588                1,
589                "rotated() should have been called"
590            );
591            assert_eq!(
592                spy.last_rotated_id.lock().unwrap().as_deref().unwrap(),
593                TEST_CONTEXT_ID
594            );
595        });
596    }
597
598    #[test]
599    fn test_accept_braces() {
600        with_test_mars(|mars, delete_called| {
601            // Some callers may store pre-existing context IDs with opening
602            // and closing curly braces. Our component should accept them, but
603            // return (and persist) UUIDs without such braces.
604            let wrapped_context_id = ["{", TEST_CONTEXT_ID, "}"].concat();
605            let mut component = ContextIDComponentInner::new(
606                &wrapped_context_id,
607                0,
608                false,
609                Box::new(DefaultContextIdCallback),
610                *FAKE_NOW,
611                mars,
612            );
613
614            // We expect to be storing TEST_CONTEXT_ID, and to return
615            // TEST_CONTEXT_ID without the braces.
616            assert_eq!(component.context_id, TEST_CONTEXT_ID);
617            assert!(Uuid::parse_str(&component.request(0, *FAKE_NOW).unwrap()).is_ok());
618            assert!(!*delete_called.lock().unwrap());
619        });
620    }
621
622    #[test]
623    fn test_persist_callback() {
624        with_test_mars(|mars, delete_called| {
625            let spy = SpyCallback::new();
626
627            let mut component = ContextIDComponentInner::new(
628                TEST_CONTEXT_ID,
629                FAKE_LONG_AGO_TS,
630                false,
631                spy.callback,
632                *FAKE_NOW,
633                mars,
634            );
635
636            component.force_rotation(*FAKE_NOW);
637
638            assert_eq!(
639                *spy.persist_called.lock().unwrap(),
640                1,
641                "persist() should have been called"
642            );
643
644            assert!(
645                Uuid::parse_str(spy.last_rotated_id.lock().unwrap().as_deref().unwrap()).is_ok(),
646                "persist() should have received a valid context_id string"
647            );
648
649            assert_eq!(
650                *spy.last_persist_ts.lock().unwrap(),
651                Some(FAKE_NOW_TS),
652                "persist() should have received the expected creation date"
653            );
654            assert!(*delete_called.lock().unwrap());
655        });
656    }
657}
context_id/lib.rs

context_id/
lib.rs
