1use std::collections::{HashMap, HashSet};
15use std::num::NonZeroUsize;
16use std::sync::{Arc, Mutex, PoisonError};
17
18use lru::LruCache;
19
20use aes_gcm::aead::{Aead, KeyInit, Payload};
21use aes_gcm::{Aes256Gcm, Nonce};
22use async_trait::async_trait;
23use base64::Engine;
24use hkdf::Hkdf;
25use sha2::Sha256;
26
27use crate::error::{ProtonError, Result};
28
29#[async_trait]
35pub trait CacheRepository: Send + Sync {
36 async fn set(&self, key: &str, value: &str, tags: &[String]) -> Result<()>;
38
39 async fn get(&self, key: &str) -> Result<Option<String>>;
41
42 async fn remove(&self, key: &str) -> Result<()>;
44
45 async fn remove_by_tag(&self, tag: &str) -> Result<()>;
47
48 async fn clear(&self) -> Result<()>;
50
51 async fn get_by_tags(&self, tags: &[String]) -> Result<Vec<(String, String)>>;
55}
56
57pub async fn set_untagged(repo: &dyn CacheRepository, key: &str, value: &str) -> Result<()> {
59 repo.set(key, value, &[]).await
60}
61
62pub const DEFAULT_CACHE_CAPACITY: usize = 8192;
69
70pub struct InMemoryCacheRepository {
79 state: Mutex<InMemoryState>,
80}
81
82impl Default for InMemoryCacheRepository {
83 fn default() -> Self {
84 Self::with_capacity(DEFAULT_CACHE_CAPACITY)
85 }
86}
87
88struct InMemoryState {
89 entries: LruCache<String, String>,
95 key_to_tags: HashMap<String, HashSet<String>>,
96 tag_to_keys: HashMap<String, HashSet<String>>,
97}
98
99impl InMemoryCacheRepository {
100 pub fn new() -> Self {
103 Self::default()
104 }
105
106 pub fn with_capacity(capacity: usize) -> Self {
110 let capacity = NonZeroUsize::new(capacity).unwrap_or(NonZeroUsize::MIN);
111 Self {
112 state: Mutex::new(InMemoryState {
113 entries: LruCache::new(capacity),
114 key_to_tags: HashMap::new(),
115 tag_to_keys: HashMap::new(),
116 }),
117 }
118 }
119
120 pub fn shared() -> Arc<dyn CacheRepository> {
122 Arc::new(Self::new())
123 }
124
125 pub fn shared_with_capacity(capacity: usize) -> Arc<dyn CacheRepository> {
127 Arc::new(Self::with_capacity(capacity))
128 }
129
130 pub fn len(&self) -> usize {
133 self.state().entries.len()
134 }
135
136 pub fn is_empty(&self) -> bool {
138 self.len() == 0
139 }
140
141 pub fn tag_index_len(&self) -> usize {
144 self.state().key_to_tags.len()
145 }
146
147 fn state(&self) -> std::sync::MutexGuard<'_, InMemoryState> {
166 self.state.lock().unwrap_or_else(PoisonError::into_inner)
167 }
168
169 fn clear_tags_for_key(state: &mut InMemoryState, key: &str) {
170 if let Some(tags) = state.key_to_tags.remove(key) {
171 for tag in tags {
172 if let Some(keys) = state.tag_to_keys.get_mut(&tag) {
173 keys.remove(key);
174 if keys.is_empty() {
175 state.tag_to_keys.remove(&tag);
176 }
177 }
178 }
179 }
180 }
181}
182
183#[async_trait]
184impl CacheRepository for InMemoryCacheRepository {
185 async fn set(&self, key: &str, value: &str, tags: &[String]) -> Result<()> {
186 let mut state = self.state();
187 Self::clear_tags_for_key(&mut state, key);
188 if let Some((evicted, _)) = state.entries.push(key.to_owned(), value.to_owned())
192 && evicted != *key
193 {
194 Self::clear_tags_for_key(&mut state, &evicted);
195 }
196 let tag_set: HashSet<String> = tags.iter().cloned().collect();
197 for tag in &tag_set {
198 state
199 .tag_to_keys
200 .entry(tag.clone())
201 .or_default()
202 .insert(key.to_owned());
203 }
204 state.key_to_tags.insert(key.to_owned(), tag_set);
205 Ok(())
206 }
207
208 async fn get(&self, key: &str) -> Result<Option<String>> {
209 let mut state = self.state();
213 Ok(state.entries.get(key).cloned())
214 }
215
216 async fn remove(&self, key: &str) -> Result<()> {
217 let mut state = self.state();
218 state.entries.pop(key);
219 Self::clear_tags_for_key(&mut state, key);
220 Ok(())
221 }
222
223 async fn remove_by_tag(&self, tag: &str) -> Result<()> {
224 let mut state = self.state();
225 let keys: Vec<String> = state
226 .tag_to_keys
227 .get(tag)
228 .map(|keys| keys.iter().cloned().collect())
229 .unwrap_or_default();
230 for key in keys {
231 state.entries.pop(&key);
234 Self::clear_tags_for_key(&mut state, &key);
235 }
236 Ok(())
237 }
238
239 async fn clear(&self) -> Result<()> {
240 let mut state = self.state();
241 state.entries.clear();
242 state.key_to_tags.clear();
243 state.tag_to_keys.clear();
244 Ok(())
245 }
246
247 async fn get_by_tags(&self, tags: &[String]) -> Result<Vec<(String, String)>> {
248 if tags.is_empty() {
249 return Ok(Vec::new());
250 }
251 let state = self.state();
252 let mut candidates: Option<HashSet<String>> = None;
253 for tag in tags {
254 match state.tag_to_keys.get(tag) {
255 Some(keys) => {
256 candidates = Some(match candidates {
257 Some(existing) => existing.intersection(keys).cloned().collect(),
258 None => keys.clone(),
259 });
260 }
261 None => return Ok(Vec::new()),
262 }
263 if candidates.as_ref().is_some_and(|c| c.is_empty()) {
264 return Ok(Vec::new());
265 }
266 }
267 let candidates = candidates.unwrap_or_default();
268 Ok(candidates
272 .into_iter()
273 .filter_map(|key| state.entries.peek(&key).map(|v| (key.clone(), v.clone())))
274 .collect())
275 }
276}
277
278pub struct EncryptedCacheRepository {
291 inner: Arc<dyn CacheRepository>,
292 encryption_key: Vec<u8>,
293}
294
295const SALT_LEN: usize = 16;
296const KEY_LEN: usize = 32;
297const NONCE_LEN: usize = 12;
298const TAG_LEN: usize = 16;
299const ENCRYPTION_CONTEXT: &[u8] = b"Drive.EncryptedCacheRepository";
300
301impl EncryptedCacheRepository {
302 pub fn new(inner: Arc<dyn CacheRepository>, encryption_key: impl Into<Vec<u8>>) -> Self {
304 Self {
305 inner,
306 encryption_key: encryption_key.into(),
307 }
308 }
309
310 pub fn shared(
312 inner: Arc<dyn CacheRepository>,
313 encryption_key: impl Into<Vec<u8>>,
314 ) -> Arc<dyn CacheRepository> {
315 Arc::new(Self::new(inner, encryption_key))
316 }
317
318 fn derive(&self, salt: &[u8], entry_key: &str) -> Result<([u8; KEY_LEN], [u8; NONCE_LEN])> {
320 let mut info = ENCRYPTION_CONTEXT.to_vec();
321 info.extend_from_slice(entry_key.as_bytes());
322 let hk = Hkdf::<Sha256>::new(Some(salt), &self.encryption_key);
323 let mut okm = [0u8; KEY_LEN + NONCE_LEN];
324 hk.expand(&info, &mut okm)
325 .map_err(|e| ProtonError::invalid_operation(format!("cache HKDF expand: {e}")))?;
326 let mut key = [0u8; KEY_LEN];
327 let mut nonce = [0u8; NONCE_LEN];
328 key.copy_from_slice(&okm[..KEY_LEN]);
329 nonce.copy_from_slice(&okm[KEY_LEN..]);
330 Ok((key, nonce))
331 }
332
333 fn encrypt(&self, entry_key: &str, plaintext: &str) -> Result<String> {
334 let mut salt = [0u8; SALT_LEN];
335 getrandom::fill(&mut salt)
336 .map_err(|e| ProtonError::invalid_operation(format!("cache salt: {e}")))?;
337 let (key, nonce) = self.derive(&salt, entry_key)?;
338 let cipher = Aes256Gcm::new_from_slice(&key)
339 .map_err(|e| ProtonError::invalid_operation(format!("cache cipher: {e}")))?;
340 let sealed = cipher
343 .encrypt(
344 &Nonce::from(nonce),
345 Payload {
346 msg: plaintext.as_bytes(),
347 aad: &[],
348 },
349 )
350 .map_err(|_| ProtonError::invalid_operation("cache encrypt failed"))?;
351 let mut out = Vec::with_capacity(SALT_LEN + sealed.len());
352 out.extend_from_slice(&salt);
353 out.extend_from_slice(&sealed);
354 Ok(base64::engine::general_purpose::STANDARD.encode(out))
355 }
356
357 fn decrypt(&self, entry_key: &str, encoded: &str) -> Result<Option<String>> {
360 let combined = base64::engine::general_purpose::STANDARD
361 .decode(encoded)
362 .map_err(|e| ProtonError::invalid_operation(format!("cache base64: {e}")))?;
363 if combined.len() < SALT_LEN + TAG_LEN {
364 return Err(ProtonError::invalid_operation("cache value too short"));
365 }
366 let (salt, sealed) = combined.split_at(SALT_LEN);
367 let (key, nonce) = self.derive(salt, entry_key)?;
368 let cipher = Aes256Gcm::new_from_slice(&key)
369 .map_err(|e| ProtonError::invalid_operation(format!("cache cipher: {e}")))?;
370 match cipher.decrypt(
371 &Nonce::from(nonce),
372 Payload {
373 msg: sealed,
374 aad: &[],
375 },
376 ) {
377 Ok(plaintext) => {
378 let text = String::from_utf8(plaintext)
379 .map_err(|e| ProtonError::invalid_operation(format!("cache utf8: {e}")))?;
380 Ok(Some(text))
381 }
382 Err(_) => Ok(None),
384 }
385 }
386}
387
388#[async_trait]
389impl CacheRepository for EncryptedCacheRepository {
390 async fn set(&self, key: &str, value: &str, tags: &[String]) -> Result<()> {
391 let encrypted = self.encrypt(key, value)?;
392 self.inner.set(key, &encrypted, tags).await
393 }
394
395 async fn get(&self, key: &str) -> Result<Option<String>> {
396 let Some(encrypted) = self.inner.get(key).await? else {
397 return Ok(None);
398 };
399 match self.decrypt(key, &encrypted)? {
400 Some(value) => Ok(Some(value)),
401 None => {
402 self.inner.clear().await?;
403 Ok(None)
404 }
405 }
406 }
407
408 async fn remove(&self, key: &str) -> Result<()> {
409 self.inner.remove(key).await
410 }
411
412 async fn remove_by_tag(&self, tag: &str) -> Result<()> {
413 self.inner.remove_by_tag(tag).await
414 }
415
416 async fn clear(&self) -> Result<()> {
417 self.inner.clear().await
418 }
419
420 async fn get_by_tags(&self, tags: &[String]) -> Result<Vec<(String, String)>> {
421 let entries = self.inner.get_by_tags(tags).await?;
422 let mut out = Vec::with_capacity(entries.len());
423 for (key, encrypted) in entries {
424 match self.decrypt(&key, &encrypted)? {
425 Some(value) => out.push((key, value)),
426 None => {
427 self.inner.clear().await?;
428 return Ok(Vec::new());
429 }
430 }
431 }
432 Ok(out)
433 }
434}
435
436#[cfg(test)]
437mod tests {
438 use super::*;
439
440 fn tags(values: &[&str]) -> Vec<String> {
441 values.iter().map(|s| s.to_string()).collect()
442 }
443
444 #[tokio::test]
445 async fn in_memory_round_trips_and_overwrites() {
446 let cache = InMemoryCacheRepository::new();
447 cache.set("k", "v1", &[]).await.unwrap();
448 assert_eq!(cache.get("k").await.unwrap().as_deref(), Some("v1"));
449 cache.set("k", "v2", &[]).await.unwrap();
450 assert_eq!(cache.get("k").await.unwrap().as_deref(), Some("v2"));
451 cache.remove("k").await.unwrap();
452 assert_eq!(cache.get("k").await.unwrap(), None);
453 }
454
455 #[tokio::test]
456 async fn in_memory_get_by_tags_intersects() {
457 let cache = InMemoryCacheRepository::new();
458 cache.set("a", "1", &tags(&["x", "y"])).await.unwrap();
459 cache.set("b", "2", &tags(&["x"])).await.unwrap();
460 cache.set("c", "3", &tags(&["y"])).await.unwrap();
461
462 let mut both = cache.get_by_tags(&tags(&["x", "y"])).await.unwrap();
463 both.sort();
464 assert_eq!(both, vec![("a".to_string(), "1".to_string())]);
465
466 let mut just_x = cache.get_by_tags(&tags(&["x"])).await.unwrap();
467 just_x.sort();
468 assert_eq!(
469 just_x,
470 vec![
471 ("a".to_string(), "1".to_string()),
472 ("b".to_string(), "2".to_string())
473 ]
474 );
475
476 assert!(cache.get_by_tags(&[]).await.unwrap().is_empty());
477 }
478
479 #[tokio::test]
480 async fn in_memory_remove_by_tag_drops_only_tagged() {
481 let cache = InMemoryCacheRepository::new();
482 cache.set("a", "1", &tags(&["x"])).await.unwrap();
483 cache.set("b", "2", &tags(&["y"])).await.unwrap();
484 cache.remove_by_tag("x").await.unwrap();
485 assert_eq!(cache.get("a").await.unwrap(), None);
486 assert_eq!(cache.get("b").await.unwrap().as_deref(), Some("2"));
487 assert!(cache.get_by_tags(&tags(&["x"])).await.unwrap().is_empty());
489 }
490
491 #[tokio::test]
501 async fn in_memory_is_bounded() {
502 let cache = InMemoryCacheRepository::with_capacity(64);
503 for i in 0..10_000 {
504 cache
505 .set(&format!("node:{i}"), "{}", &tags(&["volume:1"]))
506 .await
507 .unwrap();
508 }
509 assert_eq!(
510 cache.len(),
511 64,
512 "entries bounded by the configured capacity"
513 );
514 assert_eq!(
515 cache.tag_index_len(),
516 64,
517 "the tag index is bounded with the entries, not left to grow"
518 );
519 assert!(cache.get("node:9999").await.unwrap().is_some());
521 assert!(cache.get("node:0").await.unwrap().is_none());
522 assert_eq!(
524 cache.get_by_tags(&tags(&["volume:1"])).await.unwrap().len(),
525 64
526 );
527 }
528
529 #[tokio::test]
533 async fn evicted_entries_can_still_be_invalidated() {
534 let cache = InMemoryCacheRepository::with_capacity(4);
535 for i in 0..64 {
536 cache
537 .set(&format!("k{i}"), "v", &tags(&["share:1"]))
538 .await
539 .unwrap();
540 }
541 cache.remove_by_tag("share:1").await.unwrap();
542 assert_eq!(cache.len(), 0);
543 assert_eq!(cache.tag_index_len(), 0);
544 cache.remove_by_tag("share:1").await.unwrap();
546 cache.remove_by_tag("share:never").await.unwrap();
547 }
548
549 #[tokio::test]
552 async fn reads_keep_an_entry_alive() {
553 let cache = InMemoryCacheRepository::with_capacity(2);
554 cache.set("a", "1", &[]).await.unwrap();
555 cache.set("b", "2", &[]).await.unwrap();
556 assert!(cache.get("a").await.unwrap().is_some()); cache.set("c", "3", &[]).await.unwrap();
558 assert!(cache.get("a").await.unwrap().is_some(), "recently read");
559 assert!(
560 cache.get("b").await.unwrap().is_none(),
561 "least recently used"
562 );
563 assert!(cache.get("c").await.unwrap().is_some(), "just written");
564 }
565
566 #[test]
575 fn a_panic_holding_the_lock_does_not_break_the_cache() {
576 let cache = Arc::new(InMemoryCacheRepository::with_capacity(8));
577
578 let poisoner = {
581 let cache = cache.clone();
582 std::thread::spawn(move || {
583 let _guard = cache.state.lock();
584 panic!("worker panicked mid-cache-update");
585 })
586 };
587 assert!(poisoner.join().is_err(), "the thread really did panic");
588
589 assert_eq!(cache.len(), 0);
591 let rt = tokio::runtime::Builder::new_current_thread()
592 .build()
593 .unwrap();
594 rt.block_on(async {
595 cache.set("k", "v", &[]).await.unwrap();
596 assert_eq!(cache.get("k").await.unwrap().as_deref(), Some("v"));
597 });
598 }
599
600 #[tokio::test]
601 async fn encrypted_round_trips_and_hides_plaintext() {
602 let inner = InMemoryCacheRepository::shared();
603 let cache = EncryptedCacheRepository::new(inner.clone(), b"hunter2-master-key".to_vec());
604 cache
605 .set("share:1", "secret-value", &tags(&["t"]))
606 .await
607 .unwrap();
608
609 let stored = inner.get("share:1").await.unwrap().unwrap();
611 assert_ne!(stored, "secret-value");
612
613 assert_eq!(
615 cache.get("share:1").await.unwrap().as_deref(),
616 Some("secret-value")
617 );
618 let by_tag = cache.get_by_tags(&tags(&["t"])).await.unwrap();
620 assert_eq!(
621 by_tag,
622 vec![("share:1".to_string(), "secret-value".to_string())]
623 );
624 }
625
626 #[tokio::test]
627 async fn encrypted_wrong_key_is_a_miss_and_clears() {
628 let inner = InMemoryCacheRepository::shared();
629 EncryptedCacheRepository::new(inner.clone(), b"key-one".to_vec())
630 .set("k", "v", &[])
631 .await
632 .unwrap();
633
634 let other = EncryptedCacheRepository::new(inner.clone(), b"key-two".to_vec());
636 assert_eq!(other.get("k").await.unwrap(), None);
637 assert_eq!(inner.get("k").await.unwrap(), None);
638 }
639
640 #[tokio::test]
641 async fn encrypted_salt_is_random_per_write() {
642 let inner = InMemoryCacheRepository::shared();
643 let cache = EncryptedCacheRepository::new(inner.clone(), b"k".to_vec());
644 cache.set("k", "same", &[]).await.unwrap();
645 let first = inner.get("k").await.unwrap().unwrap();
646 cache.set("k", "same", &[]).await.unwrap();
647 let second = inner.get("k").await.unwrap().unwrap();
648 assert_ne!(first, second);
650 }
651}