586 lines
20 KiB
Rust
586 lines
20 KiB
Rust
/* 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
|
|
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
|
|
|
|
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());
|
|
});
|
|
}
|
|
}
|