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]
#[handle_error(Error)]
pub fn new(
init_context_id: &str,
creation_timestamp_s: i64,
running_in_test_automation: bool,
callback: Box<dyn ContextIdCallback>,
) -> ApiResult<Self> {
Ok(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>,
) -> Result<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) = if generated_context_id {
// If we had to generate a UUID, also force a new timestamp.
(now, true)
} else {
match creation_timestamp_s {
secs if secs > 0 => (
DateTime::<Utc>::from_timestamp(secs, 0)
.ok_or(Error::InvalidTimestamp { timestamp: secs })?,
false,
),
_ => (now, true),
}
};
let instance = Self {
context_id,
creation_timestamp,
callback_handle: RwLock::new(callback),
mars_client,
running_in_test_automation,
};
// We only need to persist these if we just generated one.
if generated_context_id || generated_creation_timestamp {
instance.persist();
}
Ok(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())
}
pub 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 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 creation_timestamp = FAKE_NOW_TS;
let component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
creation_timestamp,
false,
Box::new(DefaultContextIdCallback),
*FAKE_NOW,
mars,
)
.unwrap();
// 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());
});
}
#[test]
fn test_creation_timestamp_with_zero_value() {
with_test_mars(|mars, delete_called| {
let component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
0,
false,
Box::new(DefaultContextIdCallback),
*FAKE_NOW,
mars,
)
.unwrap();
// 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());
});
}
#[test]
fn test_empty_initial_context_id() {
with_test_mars(|mars, delete_called| {
let component = ContextIDComponentInner::new(
"",
0,
false,
Box::new(DefaultContextIdCallback),
*FAKE_NOW,
mars,
)
.unwrap();
// 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());
});
}
#[test]
fn test_empty_initial_context_id_with_creation_date() {
with_test_mars(|mars, delete_called| {
let component = ContextIDComponentInner::new(
"",
0,
false,
Box::new(DefaultContextIdCallback),
*FAKE_NOW,
mars,
)
.unwrap();
// 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());
});
}
#[test]
fn test_invalid_context_id_with_no_creation_date() {
with_test_mars(|mars, delete_called| {
let component = ContextIDComponentInner::new(
"something-invalid",
0,
false,
Box::new(DefaultContextIdCallback),
*FAKE_NOW,
mars,
)
.unwrap();
// 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());
});
}
#[test]
fn test_invalid_context_id_with_creation_date() {
with_test_mars(|mars, delete_called| {
let component = ContextIDComponentInner::new(
"something-invalid",
FAKE_LONG_AGO_TS,
false,
Box::new(DefaultContextIdCallback),
*FAKE_NOW,
mars,
)
.unwrap();
// 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());
});
}
#[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,
)
.unwrap();
// 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| {
struct FakeContextIdCallback {
rotated_called: Arc<Mutex<bool>>,
original_context_id: Arc<Mutex<Option<String>>>,
}
impl ContextIdCallback for FakeContextIdCallback {
fn persist(&self, _context_id: String, _creation_date: i64) {}
fn rotated(&self, original_context_id: String) {
*self.rotated_called.lock().unwrap() = true;
*self.original_context_id.lock().unwrap() = Some(original_context_id);
}
}
let rotated_called_flag = Arc::new(Mutex::new(false));
let original_context_id = Arc::new(Mutex::new(None));
let callback = FakeContextIdCallback {
rotated_called: Arc::clone(&rotated_called_flag),
original_context_id: Arc::clone(&original_context_id),
};
// 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(callback),
*FAKE_NOW,
mars,
)
.unwrap();
// 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!(
*rotated_called_flag.lock().unwrap(),
"rotated() should have been called"
);
assert_eq!(
original_context_id.lock().unwrap().as_deref().unwrap(),
TEST_CONTEXT_ID
);
});
}
#[test]
fn test_force_rotation() {
with_test_mars(|mars, delete_called| {
struct FakeContextIdCallback {
rotated_called: Arc<Mutex<bool>>,
original_context_id: Arc<Mutex<Option<String>>>,
}
impl ContextIdCallback for FakeContextIdCallback {
fn persist(&self, _context_id: String, _creation_date: i64) {}
fn rotated(&self, original_context_id: String) {
*self.rotated_called.lock().unwrap() = true;
*self.original_context_id.lock().unwrap() = Some(original_context_id);
}
}
let rotated_called_flag = Arc::new(Mutex::new(false));
let original_context_id = Arc::new(Mutex::new(None));
let callback = FakeContextIdCallback {
rotated_called: Arc::clone(&rotated_called_flag),
original_context_id: Arc::clone(&original_context_id),
};
let mut component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
FAKE_LONG_AGO_TS,
false,
Box::new(callback),
*FAKE_NOW,
mars,
)
.unwrap();
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!(
*rotated_called_flag.lock().unwrap(),
"rotated() should have been called"
);
assert_eq!(
original_context_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,
)
.unwrap();
// 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| {
struct FakeContextIdCallback {
persist_called: Arc<Mutex<bool>>,
context_id: Arc<Mutex<Option<String>>>,
creation_timestamp: Arc<Mutex<Option<i64>>>,
}
impl ContextIdCallback for FakeContextIdCallback {
fn persist(&self, context_id: String, creation_date: i64) {
*self.persist_called.lock().unwrap() = true;
*self.context_id.lock().unwrap() = Some(context_id);
*self.creation_timestamp.lock().unwrap() = Some(creation_date);
}
fn rotated(&self, _original_context_id: String) {}
}
let persist_called_flag = Arc::new(Mutex::new(false));
let context_id = Arc::new(Mutex::new(None));
let creation_timestamp = Arc::new(Mutex::new(None));
let callback = FakeContextIdCallback {
persist_called: Arc::clone(&persist_called_flag),
context_id: Arc::clone(&context_id),
creation_timestamp: Arc::clone(&creation_timestamp),
};
let mut component = ContextIDComponentInner::new(
TEST_CONTEXT_ID,
FAKE_LONG_AGO_TS,
false,
Box::new(callback),
*FAKE_NOW,
mars,
)
.unwrap();
component.force_rotation(*FAKE_NOW);
assert!(
*persist_called_flag.lock().unwrap(),
"persist() should have been called"
);
assert!(
Uuid::parse_str(context_id.lock().unwrap().as_deref().unwrap()).is_ok(),
"persist() should have received a valid context_id string"
);
assert_eq!(
*creation_timestamp.lock().unwrap(),
Some(FAKE_NOW_TS),
"persist() should have received the expected creation date"
);
assert!(*delete_called.lock().unwrap());
});
}
}