diff --git a/crates/proton-sdk/Cargo.toml b/crates/proton-sdk/Cargo.toml index 2c43f2d..472437d 100644 --- a/crates/proton-sdk/Cargo.toml +++ b/crates/proton-sdk/Cargo.toml @@ -21,6 +21,8 @@ base64.workspace = true hex.workspace = true hmac.workspace = true sha2.workspace = true +aes-gcm.workspace = true +hkdf.workspace = true bcrypt.workspace = true pgp.workspace = true tracing.workspace = true @@ -28,3 +30,8 @@ num-bigint.workspace = true num-traits.workspace = true getrandom.workspace = true rand.workspace = true + +[dev-dependencies] +# A one-shot loopback HTTP server backs the telemetry-instrumentation test; +# the `net` feature isn't needed by the library itself. +tokio = { workspace = true, features = ["net"] } diff --git a/crates/proton-sdk/src/cache.rs b/crates/proton-sdk/src/cache.rs new file mode 100644 index 0000000..9f3b308 --- /dev/null +++ b/crates/proton-sdk/src/cache.rs @@ -0,0 +1,436 @@ +//! Generic persistent-cache primitives. +//! +//! A string key/value store with tags, mirroring the C# `Proton.Sdk.Caching` +//! layer: the [`CacheRepository`] trait (C# `ICacheRepository`), an in-memory +//! implementation ([`InMemoryCacheRepository`], C# `InMemoryCacheRepository`), +//! and an at-rest-encryption wrapper ([`EncryptedCacheRepository`], C# +//! `EncryptedCacheRepository`). +//! +//! Higher layers (the Drive entity/secret caches) serialize typed values to +//! JSON strings and store them here; a consumer can supply an on-disk +//! implementation of [`CacheRepository`] (e.g. SQLite) without changing those +//! layers. + +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex}; + +use aes_gcm::aead::{Aead, KeyInit, Payload}; +use aes_gcm::{Aes256Gcm, Nonce}; +use async_trait::async_trait; +use base64::Engine; +use hkdf::Hkdf; +use sha2::Sha256; + +use crate::error::{ProtonError, Result}; + +/// A string key/value cache with secondary tag indexing. +/// +/// Mirrors C# `ICacheRepository`. Implementations must be cheap to share across +/// tasks (hence [`Send`] + [`Sync`]); the SDK holds them as +/// `Arc`. +#[async_trait] +pub trait CacheRepository: Send + Sync { + /// Store `value` under `key`, replacing any existing entry and its tags. + async fn set(&self, key: &str, value: &str, tags: &[String]) -> Result<()>; + + /// Fetch the value stored under `key`, or `None` if absent. + async fn get(&self, key: &str) -> Result>; + + /// Remove the entry stored under `key` (no-op if absent). + async fn remove(&self, key: &str) -> Result<()>; + + /// Remove every entry carrying `tag`. + async fn remove_by_tag(&self, tag: &str) -> Result<()>; + + /// Remove every entry. + async fn clear(&self) -> Result<()>; + + /// Return every `(key, value)` whose entry carries **all** of `tags` + /// (set intersection, matching C# `GetByTags`). An empty `tags` slice + /// returns nothing. + async fn get_by_tags(&self, tags: &[String]) -> Result>; +} + +/// Convenience: store a value with no tags. +pub async fn set_untagged(repo: &dyn CacheRepository, key: &str, value: &str) -> Result<()> { + repo.set(key, value, &[]).await +} + +/// Thread-safe in-memory [`CacheRepository`]. Mirrors C# +/// `InMemoryCacheRepository`. +#[derive(Default)] +pub struct InMemoryCacheRepository { + state: Mutex, +} + +#[derive(Default)] +struct InMemoryState { + entries: HashMap, + key_to_tags: HashMap>, + tag_to_keys: HashMap>, +} + +impl InMemoryCacheRepository { + /// Create an empty in-memory cache. + pub fn new() -> Self { + Self::default() + } + + /// Wrap a new in-memory cache in an [`Arc`] behind the trait object. + pub fn shared() -> Arc { + Arc::new(Self::new()) + } + + fn clear_tags_for_key(state: &mut InMemoryState, key: &str) { + if let Some(tags) = state.key_to_tags.remove(key) { + for tag in tags { + if let Some(keys) = state.tag_to_keys.get_mut(&tag) { + keys.remove(key); + if keys.is_empty() { + state.tag_to_keys.remove(&tag); + } + } + } + } + } +} + +#[async_trait] +impl CacheRepository for InMemoryCacheRepository { + async fn set(&self, key: &str, value: &str, tags: &[String]) -> Result<()> { + let mut state = self.state.lock().unwrap(); + Self::clear_tags_for_key(&mut state, key); + state.entries.insert(key.to_owned(), value.to_owned()); + let tag_set: HashSet = tags.iter().cloned().collect(); + for tag in &tag_set { + state + .tag_to_keys + .entry(tag.clone()) + .or_default() + .insert(key.to_owned()); + } + state.key_to_tags.insert(key.to_owned(), tag_set); + Ok(()) + } + + async fn get(&self, key: &str) -> Result> { + let state = self.state.lock().unwrap(); + Ok(state.entries.get(key).cloned()) + } + + async fn remove(&self, key: &str) -> Result<()> { + let mut state = self.state.lock().unwrap(); + state.entries.remove(key); + Self::clear_tags_for_key(&mut state, key); + Ok(()) + } + + async fn remove_by_tag(&self, tag: &str) -> Result<()> { + let mut state = self.state.lock().unwrap(); + let keys: Vec = state + .tag_to_keys + .get(tag) + .map(|keys| keys.iter().cloned().collect()) + .unwrap_or_default(); + for key in keys { + state.entries.remove(&key); + Self::clear_tags_for_key(&mut state, &key); + } + Ok(()) + } + + async fn clear(&self) -> Result<()> { + let mut state = self.state.lock().unwrap(); + state.entries.clear(); + state.key_to_tags.clear(); + state.tag_to_keys.clear(); + Ok(()) + } + + async fn get_by_tags(&self, tags: &[String]) -> Result> { + if tags.is_empty() { + return Ok(Vec::new()); + } + let state = self.state.lock().unwrap(); + let mut candidates: Option> = None; + for tag in tags { + match state.tag_to_keys.get(tag) { + Some(keys) => { + candidates = Some(match candidates { + Some(existing) => existing.intersection(keys).cloned().collect(), + None => keys.clone(), + }); + } + None => return Ok(Vec::new()), + } + if candidates.as_ref().is_some_and(|c| c.is_empty()) { + return Ok(Vec::new()); + } + } + let candidates = candidates.unwrap_or_default(); + Ok(candidates + .into_iter() + .filter_map(|key| state.entries.get(&key).map(|v| (key.clone(), v.clone()))) + .collect()) + } +} + +/// At-rest-encryption wrapper around any [`CacheRepository`]. Mirrors C# +/// `EncryptedCacheRepository`. +/// +/// Each value is encrypted independently: a random 16-byte salt feeds +/// HKDF-SHA256 (with the entry key mixed into the `info` parameter) to derive a +/// fresh AES-256-GCM key and 96-bit nonce; the stored payload is +/// `base64([salt(16)][ciphertext][tag(16)])`. Keys and tags are stored in the +/// clear (they drive lookup); only values are protected. +/// +/// A GCM authentication failure on read is treated as tampering or a changed +/// encryption key: the inner cache is cleared and the read reported as a miss +/// (matching the C# behavior). +pub struct EncryptedCacheRepository { + inner: Arc, + encryption_key: Vec, +} + +const SALT_LEN: usize = 16; +const KEY_LEN: usize = 32; +const NONCE_LEN: usize = 12; +const TAG_LEN: usize = 16; +const ENCRYPTION_CONTEXT: &[u8] = b"Drive.EncryptedCacheRepository"; + +impl EncryptedCacheRepository { + /// Wrap `inner`, encrypting values under `encryption_key`. + pub fn new(inner: Arc, encryption_key: impl Into>) -> Self { + Self { + inner, + encryption_key: encryption_key.into(), + } + } + + /// Wrap `inner` and box the result behind the trait object. + pub fn shared( + inner: Arc, + encryption_key: impl Into>, + ) -> Arc { + Arc::new(Self::new(inner, encryption_key)) + } + + /// Derive the per-entry AES key + nonce from the salt and entry key. + fn derive(&self, salt: &[u8], entry_key: &str) -> Result<([u8; KEY_LEN], [u8; NONCE_LEN])> { + let mut info = ENCRYPTION_CONTEXT.to_vec(); + info.extend_from_slice(entry_key.as_bytes()); + let hk = Hkdf::::new(Some(salt), &self.encryption_key); + let mut okm = [0u8; KEY_LEN + NONCE_LEN]; + hk.expand(&info, &mut okm) + .map_err(|e| ProtonError::invalid_operation(format!("cache HKDF expand: {e}")))?; + let mut key = [0u8; KEY_LEN]; + let mut nonce = [0u8; NONCE_LEN]; + key.copy_from_slice(&okm[..KEY_LEN]); + nonce.copy_from_slice(&okm[KEY_LEN..]); + Ok((key, nonce)) + } + + fn encrypt(&self, entry_key: &str, plaintext: &str) -> Result { + let mut salt = [0u8; SALT_LEN]; + getrandom::getrandom(&mut salt) + .map_err(|e| ProtonError::invalid_operation(format!("cache salt: {e}")))?; + let (key, nonce) = self.derive(&salt, entry_key)?; + let cipher = Aes256Gcm::new_from_slice(&key) + .map_err(|e| ProtonError::invalid_operation(format!("cache cipher: {e}")))?; + // aes-gcm appends the 16-byte tag to the ciphertext, so this yields + // `ciphertext || tag` — matching the C# `[salt][ciphertext][tag]` layout. + let sealed = cipher + .encrypt( + Nonce::from_slice(&nonce), + Payload { + msg: plaintext.as_bytes(), + aad: &[], + }, + ) + .map_err(|_| ProtonError::invalid_operation("cache encrypt failed"))?; + let mut out = Vec::with_capacity(SALT_LEN + sealed.len()); + out.extend_from_slice(&salt); + out.extend_from_slice(&sealed); + Ok(base64::engine::general_purpose::STANDARD.encode(out)) + } + + /// Decrypt a stored value. `Ok(None)` signals a GCM auth failure (tampered + /// or stale entry); the caller clears the cache and treats it as a miss. + fn decrypt(&self, entry_key: &str, encoded: &str) -> Result> { + let combined = base64::engine::general_purpose::STANDARD + .decode(encoded) + .map_err(|e| ProtonError::invalid_operation(format!("cache base64: {e}")))?; + if combined.len() < SALT_LEN + TAG_LEN { + return Err(ProtonError::invalid_operation("cache value too short")); + } + let (salt, sealed) = combined.split_at(SALT_LEN); + let (key, nonce) = self.derive(salt, entry_key)?; + let cipher = Aes256Gcm::new_from_slice(&key) + .map_err(|e| ProtonError::invalid_operation(format!("cache cipher: {e}")))?; + match cipher.decrypt( + Nonce::from_slice(&nonce), + Payload { + msg: sealed, + aad: &[], + }, + ) { + Ok(plaintext) => { + let text = String::from_utf8(plaintext) + .map_err(|e| ProtonError::invalid_operation(format!("cache utf8: {e}")))?; + Ok(Some(text)) + } + // Authentication failure: tampering or a changed key. Signal a miss. + Err(_) => Ok(None), + } + } +} + +#[async_trait] +impl CacheRepository for EncryptedCacheRepository { + async fn set(&self, key: &str, value: &str, tags: &[String]) -> Result<()> { + let encrypted = self.encrypt(key, value)?; + self.inner.set(key, &encrypted, tags).await + } + + async fn get(&self, key: &str) -> Result> { + let Some(encrypted) = self.inner.get(key).await? else { + return Ok(None); + }; + match self.decrypt(key, &encrypted)? { + Some(value) => Ok(Some(value)), + None => { + self.inner.clear().await?; + Ok(None) + } + } + } + + async fn remove(&self, key: &str) -> Result<()> { + self.inner.remove(key).await + } + + async fn remove_by_tag(&self, tag: &str) -> Result<()> { + self.inner.remove_by_tag(tag).await + } + + async fn clear(&self) -> Result<()> { + self.inner.clear().await + } + + async fn get_by_tags(&self, tags: &[String]) -> Result> { + let entries = self.inner.get_by_tags(tags).await?; + let mut out = Vec::with_capacity(entries.len()); + for (key, encrypted) in entries { + match self.decrypt(&key, &encrypted)? { + Some(value) => out.push((key, value)), + None => { + self.inner.clear().await?; + return Ok(Vec::new()); + } + } + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn tags(values: &[&str]) -> Vec { + values.iter().map(|s| s.to_string()).collect() + } + + #[tokio::test] + async fn in_memory_round_trips_and_overwrites() { + let cache = InMemoryCacheRepository::new(); + cache.set("k", "v1", &[]).await.unwrap(); + assert_eq!(cache.get("k").await.unwrap().as_deref(), Some("v1")); + cache.set("k", "v2", &[]).await.unwrap(); + assert_eq!(cache.get("k").await.unwrap().as_deref(), Some("v2")); + cache.remove("k").await.unwrap(); + assert_eq!(cache.get("k").await.unwrap(), None); + } + + #[tokio::test] + async fn in_memory_get_by_tags_intersects() { + let cache = InMemoryCacheRepository::new(); + cache.set("a", "1", &tags(&["x", "y"])).await.unwrap(); + cache.set("b", "2", &tags(&["x"])).await.unwrap(); + cache.set("c", "3", &tags(&["y"])).await.unwrap(); + + let mut both = cache.get_by_tags(&tags(&["x", "y"])).await.unwrap(); + both.sort(); + assert_eq!(both, vec![("a".to_string(), "1".to_string())]); + + let mut just_x = cache.get_by_tags(&tags(&["x"])).await.unwrap(); + just_x.sort(); + assert_eq!( + just_x, + vec![ + ("a".to_string(), "1".to_string()), + ("b".to_string(), "2".to_string()) + ] + ); + + assert!(cache.get_by_tags(&[]).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn in_memory_remove_by_tag_drops_only_tagged() { + let cache = InMemoryCacheRepository::new(); + cache.set("a", "1", &tags(&["x"])).await.unwrap(); + cache.set("b", "2", &tags(&["y"])).await.unwrap(); + cache.remove_by_tag("x").await.unwrap(); + assert_eq!(cache.get("a").await.unwrap(), None); + assert_eq!(cache.get("b").await.unwrap().as_deref(), Some("2")); + // The tag index is cleaned up too: re-querying yields nothing. + assert!(cache.get_by_tags(&tags(&["x"])).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn encrypted_round_trips_and_hides_plaintext() { + let inner = InMemoryCacheRepository::shared(); + let cache = EncryptedCacheRepository::new(inner.clone(), b"hunter2-master-key".to_vec()); + cache.set("share:1", "secret-value", &tags(&["t"])).await.unwrap(); + + // Stored ciphertext is not the plaintext. + let stored = inner.get("share:1").await.unwrap().unwrap(); + assert_ne!(stored, "secret-value"); + + // Round-trips through the wrapper. + assert_eq!( + cache.get("share:1").await.unwrap().as_deref(), + Some("secret-value") + ); + // Tags pass through to the inner store. + let by_tag = cache.get_by_tags(&tags(&["t"])).await.unwrap(); + assert_eq!(by_tag, vec![("share:1".to_string(), "secret-value".to_string())]); + } + + #[tokio::test] + async fn encrypted_wrong_key_is_a_miss_and_clears() { + let inner = InMemoryCacheRepository::shared(); + EncryptedCacheRepository::new(inner.clone(), b"key-one".to_vec()) + .set("k", "v", &[]) + .await + .unwrap(); + + // A different key fails the GCM tag check → treated as a miss, cache cleared. + let other = EncryptedCacheRepository::new(inner.clone(), b"key-two".to_vec()); + assert_eq!(other.get("k").await.unwrap(), None); + assert_eq!(inner.get("k").await.unwrap(), None); + } + + #[tokio::test] + async fn encrypted_salt_is_random_per_write() { + let inner = InMemoryCacheRepository::shared(); + let cache = EncryptedCacheRepository::new(inner.clone(), b"k".to_vec()); + cache.set("k", "same", &[]).await.unwrap(); + let first = inner.get("k").await.unwrap().unwrap(); + cache.set("k", "same", &[]).await.unwrap(); + let second = inner.get("k").await.unwrap().unwrap(); + // Random salt per write ⇒ identical plaintext yields different ciphertext. + assert_ne!(first, second); + } +} diff --git a/crates/proton-sdk/src/lib.rs b/crates/proton-sdk/src/lib.rs index a9dec5c..0fe0d2f 100644 --- a/crates/proton-sdk/src/lib.rs +++ b/crates/proton-sdk/src/lib.rs @@ -12,12 +12,14 @@ pub mod account; pub mod api; +pub mod cache; pub mod config; pub mod crypto; pub mod error; pub mod http; pub mod ids; pub mod session; +pub mod telemetry; pub use error::{ProtonApiError, ProtonError, Result}; pub use session::ProtonApiSession;