Revision control
Copy as Markdown
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
mod error;
pub use error::{ApiError, ApiResult, Error, Result};
use chrono::{DateTime, Duration, Utc};
use error_support::{error, handle_error};
use parking_lot::{Mutex, RwLock};
use uuid::Uuid;
uniffi::setup_scaffolding!("context_id");
mod callback;
pub use callback::{ContextIdCallback, DefaultContextIdCallback};
mod mars;
use mars::{MARSClient, SimpleMARSClient};
/// Top-level API for the context_id component
#[derive(uniffi::Object)]
pub struct ContextIDComponent {
inner: Mutex<ContextIDComponentInner>,
}
#[uniffi::export]
impl ContextIDComponent {
/// Construct a new [ContextIDComponent].
///
/// If no creation timestamp is provided, the current time will be used.
#[uniffi::constructor]
pub fn new(
init_context_id: &str,
creation_timestamp_s: i64,
running_in_test_automation: bool,
callback: Box<dyn ContextIdCallback>,
) -> Self {
Self {
inner: Mutex::new(ContextIDComponentInner::new(
init_context_id,
creation_timestamp_s,
running_in_test_automation,
callback,
Utc::now(),
Box::new(SimpleMARSClient::new()),
)),
}
}
/// Return the current context ID string.
#[handle_error(Error)]
pub fn request(&self, rotation_days_in_s: u8) -> ApiResult<String> {
let mut inner = self.inner.lock();
inner.request(rotation_days_in_s, Utc::now())
}
/// Regenerate the context ID.
#[handle_error(Error)]
pub fn force_rotation(&self) -> ApiResult<()> {
let mut inner = self.inner.lock();
inner.force_rotation(Utc::now());
Ok(())
}
/// Unset the callbacks set during construction, and use a default
/// no-op ContextIdCallback instead.
#[handle_error(Error)]
pub fn unset_callback(&self) -> ApiResult<()> {
let mut inner = self.inner.lock();
inner.unset_callback();
Ok(())
}
}
struct ContextIDComponentInner {
context_id: String,
creation_timestamp: DateTime<Utc>,
callback_handle: RwLock<Box<dyn ContextIdCallback>>,
mars_client: Box<dyn MARSClient>,
running_in_test_automation: bool,
}
impl ContextIDComponentInner {
pub fn new(
init_context_id: &str,
creation_timestamp_s: i64,
running_in_test_automation: bool,
callback: Box<dyn ContextIdCallback>,
now: DateTime<Utc>,
mars_client: Box<dyn MARSClient>,
) -> Self {
// Some historical context IDs are stored within opening and closing
// braces, and our endpoints have tolerated this, but ideally we'd
// send without the braces, so we strip any off here.
let (context_id, generated_context_id) = match init_context_id
.trim()
.trim_start_matches('{')
.trim_end_matches('}')
{
"" => (Uuid::new_v4().to_string(), true),
// If the passed in string isn't empty, but still not a valid UUID,
// just go ahead and generate a new UUID.
s => match Uuid::parse_str(s) {
Ok(_) => (s.to_string(), false),
Err(_) => (Uuid::new_v4().to_string(), true),
},
};
let (creation_timestamp, generated_creation_timestamp, force_rotation) =
match (generated_context_id, creation_timestamp_s) {
// We generated a new context ID then we need a fresh timestamp and no rotation needed
(true, _) => (now, true, false),
// Pre-existing context ID with a positive timestamp:
// try parsing it (any real UNIX-epoch timestamp is orders of magnitude within chrono’s 262_000-year range),
// and if it’s somehow out-of-range fall back to `now` and force rotation
(false, secs) if secs > 0 => DateTime::<Utc>::from_timestamp(secs, 0)
.map(|ts| (ts, false, false))
.unwrap_or((now, true, true)),
// Pre-existing context ID with zero timestamp then use current time and no rotation needed
(false, 0) => (now, true, false),
// Pre-existing context ID but INVALID timestamp then use current time but FORCE rotation
(false, _) => (now, true, true),
};
let mut instance = Self {
context_id,
creation_timestamp,
callback_handle: RwLock::new(callback),
mars_client,
running_in_test_automation,
};
if force_rotation {
// force_rotation will cause a implicit persist
instance.force_rotation(creation_timestamp)
} else if generated_context_id || generated_creation_timestamp {
// We only need to persist these if we just generated one.
instance.persist();
}
instance
}
pub fn request(&mut self, rotation_days: u8, now: DateTime<Utc>) -> Result<String> {
if rotation_days == 0 {
return Ok(self.context_id.clone());
}
let age = now - self.creation_timestamp;
if age >= Duration::days(rotation_days.into()) {
self.rotate_context_id(now);
}
Ok(self.context_id.clone())
}
fn rotate_context_id(&mut self, now: DateTime<Utc>) {
let original_context_id = self.context_id.clone();
self.context_id = Uuid::new_v4().to_string();
self.creation_timestamp = now;
self.persist();
// If we're running in test automation in the embedder, we don't want
// to be sending actual network requests to MARS.
if !self.running_in_test_automation {
let _ = self
.mars_client
.delete(original_context_id.clone())
.inspect_err(|e| error!("Failed to contact MARS: {}", e));
}
// In a perfect world, we'd call Glean ourselves here - however,
// our uniffi / Rust infrastructure doesn't yet support doing that,
// so we'll delegate to our embedder to send the Glean ping by
// calling a `rotated` callback method.
self.callback_handle
.read()
.rotated(original_context_id.clone());
}
pub fn force_rotation(&mut self, now: DateTime<Utc>) {
self.rotate_context_id(now);
}
pub fn persist(&self) {
self.callback_handle
.read()
.persist(self.context_id.clone(), self.creation_timestamp.timestamp());
}
pub fn unset_callback(&mut self) {
self.callback_handle = RwLock::new(Box::new(DefaultContextIdCallback));
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::callback::utils::SpyCallback;
use chrono::Utc;
use lazy_static::lazy_static;
use std::sync::{Arc, Mutex};
// 1745859061 ~= the timestamp for when this test was written (Apr 28, 2025)
const FAKE_NOW_TS: i64 = 1745859061;
const TEST_CONTEXT_ID: &str = "decafbad-0cd1-0cd2-0cd3-decafbad1000";
// 1706763600 ~= Jan 1st, 2024, which is long ago compared to FAKE_NOW.
const FAKE_LONG_AGO_TS: i64 = 1706763600;
lazy_static! {
static ref FAKE_NOW: DateTime<Utc> = DateTime::from_timestamp(FAKE_NOW_TS, 0).unwrap();
static ref FAKE_LONG_AGO: DateTime<Utc> =
DateTime::from_timestamp(FAKE_LONG_AGO_TS, 0).unwrap();
}
pub struct TestMARSClient {
delete_called: Arc<Mutex<bool>>,
}
impl TestMARSClient {
pub fn new(delete_called: Arc<Mutex<bool>>) -> Self {
Self { delete_called }
}
}
impl MARSClient for TestMARSClient {
fn delete(&self, _context_id: String) -> crate::Result<()> {
*self.delete_called.lock().unwrap() = true;
Ok(())
}
}
fn with_test_mars<F: FnOnce(Box<dyn MARSClient + Send + Sync>, Arc<Mutex<bool>>)>(test: F) {
let delete_called = Arc::new(Mutex::new(false));
let mars = Box::new(TestMARSClient::new(Arc::clone(&delete_called)));
test(mars, delete_called);
}
#[test]
fn test_creation_timestamp_with_some_value() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
let creation_timestamp = FAKE_NOW_TS;
let component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
creation_timestamp,
false,
spy.callback,
*FAKE_NOW,
mars,
);
// We should have left the context_id and creation_timestamp
// untouched if a creation_timestamp was passed.
assert_eq!(component.context_id, TEST_CONTEXT_ID);
assert_eq!(component.creation_timestamp.timestamp(), creation_timestamp);
assert!(!*delete_called.lock().unwrap());
// Neither persist nor rotate should have been called
assert_eq!(
*spy.persist_called.lock().unwrap(),
0,
"persist() should NOTE have been called"
);
assert_eq!(
*spy.rotated_called.lock().unwrap(),
0,
"rotated() should NOT have been called"
);
});
}
#[test]
fn test_creation_timestamp_with_zero_value() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
let component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
0,
false,
spy.callback,
*FAKE_NOW,
mars,
);
// If 0 was passed as the creation_timestamp, we'll interpret that
// as there having been no stored creation_timestamp. In that case,
// we'll use "now" as the creation_timestamp.
assert_eq!(component.context_id, TEST_CONTEXT_ID);
assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
assert!(!*delete_called.lock().unwrap());
// zero‐timestamp should trigger a persist, but NOT a rotation
assert_eq!(
*spy.persist_called.lock().unwrap(),
1,
"persist() should have been called once"
);
assert_eq!(
*spy.rotated_called.lock().unwrap(),
0,
"rotated() should NOT have been called"
);
});
}
#[test]
fn test_empty_initial_context_id() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
let component =
ContextIDComponentInner::new("", 0, false, spy.callback, *FAKE_NOW, mars);
// We expect a new UUID to have been generated for context_id.
assert!(Uuid::parse_str(&component.context_id).is_ok());
assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
assert_eq!(
*spy.persist_called.lock().unwrap(),
1,
"persist() should have been called once"
);
assert_eq!(
*spy.rotated_called.lock().unwrap(),
0,
"rotated() should NOT have been called"
);
assert!(!*delete_called.lock().unwrap());
});
}
#[test]
fn test_empty_initial_context_id_with_creation_date() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
let component =
ContextIDComponentInner::new("", 0, false, spy.callback, *FAKE_NOW, mars);
// We expect a new UUID to have been generated for context_id.
assert!(Uuid::parse_str(&component.context_id).is_ok());
assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
assert!(!*delete_called.lock().unwrap());
assert_eq!(
*spy.persist_called.lock().unwrap(),
1,
"persist() should have been called once"
);
assert_eq!(
*spy.rotated_called.lock().unwrap(),
0,
"rotated() should NOT have been called"
);
});
}
#[test]
fn test_invalid_context_id_with_no_creation_date() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
let component = ContextIDComponentInner::new(
"something-invalid",
0,
false,
spy.callback,
*FAKE_NOW,
mars,
);
// We expect a new UUID to have been generated for context_id.
assert!(Uuid::parse_str(&component.context_id).is_ok());
assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
// Also expect a persist to have been called, but not a rotation.
assert_eq!(
*spy.persist_called.lock().unwrap(),
1,
"persist() should have been called once"
);
assert_eq!(
*spy.rotated_called.lock().unwrap(),
0,
"rotated() should NOT have been called"
);
assert!(!*delete_called.lock().unwrap());
});
}
#[test]
fn test_invalid_context_id_with_creation_date() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
let component = ContextIDComponentInner::new(
"something-invalid",
FAKE_LONG_AGO_TS,
false,
spy.callback,
*FAKE_NOW,
mars,
);
// We expect a new UUID to have been generated for context_id.
assert!(Uuid::parse_str(&component.context_id).is_ok());
assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
// Also expect a persist to have been called, but not a rotation.
assert_eq!(
*spy.persist_called.lock().unwrap(),
1,
"persist() should have been called once"
);
assert_eq!(
*spy.rotated_called.lock().unwrap(),
0,
"rotated() should NOT have been called"
);
assert!(!*delete_called.lock().unwrap());
});
}
#[test]
fn test_context_id_with_invalid_creation_date() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
let component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
-1,
false,
spy.callback,
*FAKE_NOW,
mars,
);
// A new UUID should have been generated.
assert!(Uuid::parse_str(&component.context_id).is_ok());
assert_ne!(component.context_id, TEST_CONTEXT_ID);
// The creation timestamp must have been set to now.
assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
// Also expect a persist to have been called and a rotation.
assert_eq!(
*spy.persist_called.lock().unwrap(),
1,
"persist() should have been called once"
);
assert_eq!(
*spy.rotated_called.lock().unwrap(),
1,
"rotated() should have been called once"
);
// Since we forced a rotation, the MARS delete should have been called.
assert!(*delete_called.lock().unwrap());
});
}
#[test]
fn test_context_id_with_out_of_range_creation_date() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
// Way beyond chrono’s 262_000-year range
let huge_secs = i64::MAX;
let component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
huge_secs,
false,
spy.callback,
*FAKE_NOW,
mars,
);
// A new UUID should have been generated.
assert!(Uuid::parse_str(&component.context_id).is_ok());
assert_ne!(component.context_id, TEST_CONTEXT_ID);
// The creation timestamp must have been set to now.
assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
// Also expect a persist to have been called and a rotation.
assert_eq!(
*spy.persist_called.lock().unwrap(),
1,
"persist() should have been called once"
);
assert_eq!(
*spy.rotated_called.lock().unwrap(),
1,
"rotated() should have been called once"
);
// Since we forced a rotation on out-of-range, the MARS delete should have been called.
assert!(*delete_called.lock().unwrap());
});
}
#[test]
fn test_request_no_rotation() {
with_test_mars(|mars, delete_called| {
// Let's create a context_id with a creation date far in the past.
let mut component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
FAKE_LONG_AGO_TS,
false,
Box::new(DefaultContextIdCallback),
*FAKE_NOW,
mars,
);
// We expect neither the UUID nor creation_timestamp to have been changed.
assert_eq!(component.context_id, TEST_CONTEXT_ID);
assert_eq!(component.creation_timestamp, FAKE_LONG_AGO.clone());
// Now request the context_id, passing 0 for the rotation_days. We
// interpret this to mean "do not rotate".
assert_eq!(
component.request(0, *FAKE_NOW).unwrap(),
component.context_id
);
assert_eq!(component.creation_timestamp, FAKE_LONG_AGO.clone());
assert!(!*delete_called.lock().unwrap());
});
}
#[test]
fn test_request_with_rotation() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
// Let's create a context_id with a creation date far in the past.
let mut component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
FAKE_LONG_AGO_TS,
false,
spy.callback,
*FAKE_NOW,
mars,
);
// We expect neither the UUID nor creation_timestamp to have been changed.
assert_eq!(component.context_id, TEST_CONTEXT_ID);
assert_eq!(component.creation_timestamp, FAKE_LONG_AGO.clone());
// Now request the context_id, passing 2 for the rotation_days. Since
// the number of days since FAKE_LONG_AGO is greater than 2 days, we
// expect a new context_id to be generated, and the creation_timestamp
// to update to now.
assert!(Uuid::parse_str(&component.request(2, *FAKE_NOW).unwrap()).is_ok());
assert_ne!(component.context_id, TEST_CONTEXT_ID);
assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
assert!(*delete_called.lock().unwrap());
assert_eq!(
*spy.rotated_called.lock().unwrap(),
1,
"rotated() should have been called"
);
assert_eq!(
*spy.persist_called.lock().unwrap(),
1,
"persist() should have been called"
);
assert_eq!(
spy.last_rotated_id.lock().unwrap().as_deref().unwrap(),
TEST_CONTEXT_ID
);
});
}
#[test]
fn test_force_rotation() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
let mut component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
FAKE_LONG_AGO_TS,
false,
spy.callback,
*FAKE_NOW,
mars,
);
component.force_rotation(*FAKE_NOW);
assert!(Uuid::parse_str(&component.request(2, *FAKE_NOW).unwrap()).is_ok());
assert_ne!(component.context_id, TEST_CONTEXT_ID);
assert_eq!(component.creation_timestamp, FAKE_NOW.clone());
assert!(*delete_called.lock().unwrap());
assert_eq!(
*spy.rotated_called.lock().unwrap(),
1,
"rotated() should have been called"
);
assert_eq!(
spy.last_rotated_id.lock().unwrap().as_deref().unwrap(),
TEST_CONTEXT_ID
);
});
}
#[test]
fn test_accept_braces() {
with_test_mars(|mars, delete_called| {
// Some callers may store pre-existing context IDs with opening
// and closing curly braces. Our component should accept them, but
// return (and persist) UUIDs without such braces.
let wrapped_context_id = ["{", TEST_CONTEXT_ID, "}"].concat();
let mut component = ContextIDComponentInner::new(
&wrapped_context_id,
0,
false,
Box::new(DefaultContextIdCallback),
*FAKE_NOW,
mars,
);
// We expect to be storing TEST_CONTEXT_ID, and to return
// TEST_CONTEXT_ID without the braces.
assert_eq!(component.context_id, TEST_CONTEXT_ID);
assert!(Uuid::parse_str(&component.request(0, *FAKE_NOW).unwrap()).is_ok());
assert!(!*delete_called.lock().unwrap());
});
}
#[test]
fn test_persist_callback() {
with_test_mars(|mars, delete_called| {
let spy = SpyCallback::new();
let mut component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
FAKE_LONG_AGO_TS,
false,
spy.callback,
*FAKE_NOW,
mars,
);
component.force_rotation(*FAKE_NOW);
assert_eq!(
*spy.persist_called.lock().unwrap(),
1,
"persist() should have been called"
);
assert!(
Uuid::parse_str(spy.last_rotated_id.lock().unwrap().as_deref().unwrap()).is_ok(),
"persist() should have received a valid context_id string"
);
assert_eq!(
*spy.last_persist_ts.lock().unwrap(),
Some(FAKE_NOW_TS),
"persist() should have received the expected creation date"
);
assert!(*delete_called.lock().unwrap());
});
}
}