diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 022d68a72..619bc4b6d 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -5,28 +5,25 @@ use anyhow::Context as _; impl Client { /// Build a [`SignalProtocolStoreAdapter`] from the current device state and signal cache. - pub(crate) async fn signal_adapter( + pub(crate) fn signal_adapter( &self, ) -> crate::store::signal_adapter::SignalProtocolStoreAdapter { - let device_store = self.persistence_manager.get_device_arc().await; - self.signal_adapter_from(device_store) + self.signal_adapter_from(self.persistence_manager.clone()) } /// Build a standalone [`SenderKeyAdapter`] from the current device state and /// signal cache, avoiding the full five-store adapter on the SKDM path. - pub(crate) async fn sender_key_adapter( - &self, - ) -> crate::store::signal_adapter::SenderKeyAdapter { + pub(crate) fn sender_key_adapter(&self) -> crate::store::signal_adapter::SenderKeyAdapter { crate::store::signal_adapter::SenderKeyAdapter::new( - self.persistence_manager.get_device_arc().await, + self.persistence_manager.clone(), self.signal_cache.clone(), ) } - /// Build a [`SignalProtocolStoreAdapter`] from a pre-fetched device arc. + /// Build a [`SignalProtocolStoreAdapter`] from a pre-fetched persistence handle. pub(crate) fn signal_adapter_from( &self, - device_store: Arc>, + device_store: Arc, ) -> crate::store::signal_adapter::SignalProtocolStoreAdapter { crate::store::signal_adapter::SignalProtocolStoreAdapter::new( device_store, diff --git a/src/client/device_registry.rs b/src/client/device_registry.rs index ebd52310b..cbd28d840 100644 --- a/src/client/device_registry.rs +++ b/src/client/device_registry.rs @@ -3318,7 +3318,7 @@ mod tests { "rotation must wait for the in-flight advance" ); - let mut sender_key_store = client.sender_key_adapter().await; + let mut sender_key_store = client.sender_key_adapter(); group_encrypt( &mut sender_key_store, &name, diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 820d8a4ed..3e1cfbc2e 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -485,7 +485,7 @@ impl Client { } } - let mut adapter = self.signal_adapter().await; + let mut adapter = self.signal_adapter(); let mut rng = rand::make_rng::(); let mut success_count = 0; @@ -596,7 +596,7 @@ impl Client { /// pkmsg too, not as plain msg. #[cfg(feature = "voip-runtime")] pub(crate) async fn would_emit_pkmsg(&self, jid: &Jid) -> Result { - let device_store = self.persistence_manager.get_device_arc().await; + let device_store = self.persistence_manager.clone(); let mut adapter = self.signal_adapter_from(device_store); let signal_addr = jid.to_protocol_address(); wacore::send::pkmsg_would_be_emitted(&mut adapter.session_store, &signal_addr).await diff --git a/src/features/signal.rs b/src/features/signal.rs index 86a2b0c93..43a231aab 100644 --- a/src/features/signal.rs +++ b/src/features/signal.rs @@ -182,7 +182,7 @@ impl<'a> Signal<'a> { let address = jid.to_protocol_address(); let lock = self.client.session_lock_for(address.as_str()).await; let _guard = lock.lock().await; - let mut adapter = self.client.signal_adapter().await; + let mut adapter = self.client.signal_adapter(); Ok(message_encrypt( plaintext, &address, @@ -200,7 +200,7 @@ impl<'a> Signal<'a> { let address = jid.to_protocol_address(); let lock = self.client.session_lock_for(address.as_str()).await; let _guard = lock.lock().await; - let mut adapter = self.client.signal_adapter().await; + let mut adapter = self.client.signal_adapter(); let mut rng = rand::make_rng::(); let decrypted = message_decrypt( parsed, @@ -242,7 +242,7 @@ impl<'a> Signal<'a> { bundle: &PreKeyBundle, ) -> Result { let resolved = self.client.resolve_encryption_jid(jid).await; - let mut adapter = self.client.signal_adapter().await; + let mut adapter = self.client.signal_adapter(); let mut rng = rand::make_rng::(); let identity_change = self .client @@ -276,7 +276,7 @@ impl<'a> Signal<'a> { let distribution = decode_sender_key_distribution(distribution)?; let sender_address = sender_jid.to_non_ad().to_protocol_address(); let sender_key_name = make_sender_key_name(group_jid, &sender_address); - let mut store = self.client.sender_key_adapter().await; + let mut store = self.client.sender_key_adapter(); let chain_lock = store.sender_key_lock(&sender_key_name).await; let chain_guard = chain_lock.lock().await; @@ -294,7 +294,7 @@ impl<'a> Signal<'a> { ) -> Result, SignalError> { let sender_address = sender_jid.to_non_ad().to_protocol_address(); let sender_key_name = make_sender_key_name(group_jid, &sender_address); - let mut store = self.client.sender_key_adapter().await; + let mut store = self.client.sender_key_adapter(); let chain_lock = store.sender_key_lock(&sender_key_name).await; let chain_guard = chain_lock.lock().await; let distribution = wacore::send::create_sender_key_distribution_message_for_group( @@ -512,7 +512,7 @@ impl<'a> Signal<'a> { .await? .is_some(); - let mut store = self.client.sender_key_adapter().await; + let mut store = self.client.sender_key_adapter(); let mut rng = rand::make_rng::(); let pending_distribution = self @@ -574,7 +574,7 @@ impl<'a> Signal<'a> { let sender_key_name = make_sender_key_name(group_jid, &sender_jid.to_non_ad().to_protocol_address()); - let mut store = self.client.sender_key_adapter().await; + let mut store = self.client.sender_key_adapter(); let chain_lock = store.sender_key_lock(&sender_key_name).await; let _chain_guard = chain_lock.lock().await; @@ -652,7 +652,7 @@ impl<'a> Signal<'a> { let _session_guards = self.client.session_guards_for(&lock_jids).await; let plaintext = MessageUtils::encode_and_pad(message); - let mut adapter = self.client.signal_adapter().await; + let mut adapter = self.client.signal_adapter(); let mediatype = wacore::send::media_type_from_message(message); let hide_decrypt_fail = wacore::send::should_hide_decrypt_fail(message); diff --git a/src/message/receive.rs b/src/message/receive.rs index f014027b7..c43a1022e 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -638,7 +638,7 @@ impl Client { // Started after the lock so the histogram excludes lock/queue wait. let _t = wacore::telemetry::timer(wacore::telemetry::DECRYPT_DURATION); - let mut adapter = self.signal_adapter().await; + let mut adapter = self.signal_adapter(); let mut rng = rand::make_rng::(); let mut outcome = SessionBatchOutcome::default(); // Buffer plaintexts to handle after the ratchet lock drops (see the drain @@ -1229,7 +1229,7 @@ impl Client { if payloads.is_empty() { return Ok(()); } - let mut adapter = self.signal_adapter().await; + let mut adapter = self.signal_adapter(); // Always use bare sender for sender key operations. Real WA delivers // skmsg with bare participant but pkmsg (SKDM) with device-qualified diff --git a/src/message/tests.rs b/src/message/tests.rs index c01797dbe..12d2b6e18 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -778,7 +778,7 @@ async fn bobs_prekey_bundle_with_spk_id(client: &Arc, spk_id: u32) -> (P // Read/write prekeys through the same trait surface production uses // (see signal_adapter.rs). Avoids reaching past `PersistenceManager` // to mutate device storage directly. - let mut adapter = client.signal_adapter().await; + let mut adapter = client.signal_adapter(); let spk_record = adapter .signed_pre_key_store .get_signed_pre_key(1.into()) diff --git a/src/retry.rs b/src/retry.rs index f619fbfae..dfaf21567 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -839,7 +839,7 @@ impl Client { let signal_address = encryption_jid.to_protocol_address(); let session_mutex = self.session_lock_for(signal_address.as_str()).await; let session_guard = session_mutex.lock().await; - let mut store_adapter = self.signal_adapter().await; + let mut store_adapter = self.signal_adapter(); let device_snapshot = self.persistence_manager.get_device_snapshot(); let edit = wacore::types::message::EditAttribute::infer_from_message(&message); @@ -930,7 +930,7 @@ impl Client { let encoded = pre_encoded .filter(|_| can_reuse_encoding) .or(encoded_fallback.as_deref()); - let device_store = self.persistence_manager.get_device_arc().await; + let device_store = self.persistence_manager.clone(); let mut store_adapter = self.signal_adapter_from(device_store); let mut stores = store_adapter.as_signal_stores(); let edit = wacore::types::message::EditAttribute::infer_from_message(&message); @@ -1420,7 +1420,7 @@ impl Client { identity_key.into(), )?; - let mut adapter = self.signal_adapter().await; + let mut adapter = self.signal_adapter(); let mut rng = rand::make_rng::(); self.install_prekey_bundle_cached(requester_jid, &bundle, &mut adapter, &mut rng) .await?; diff --git a/src/send/mod.rs b/src/send/mod.rs index 791066835..1c93bd292 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -1137,7 +1137,7 @@ impl Client { self.add_recent_message(&to, &request_id, &message, shared_content.clone()) .await; - let device_store_arc = self.persistence_manager.get_device_arc().await; + let device_store_arc = self.persistence_manager.clone(); let to_str = to.to_string(); let distribution_guard = self.group_distribution_lock(&to).await; @@ -1909,7 +1909,7 @@ impl Client { let session_mutex = self.session_lock_for(signal_addr.as_str()).await; let _session_guard = session_mutex.lock().await; - let mut store_adapter = self.signal_adapter().await; + let mut store_adapter = self.signal_adapter(); let device_snapshot = self.persistence_manager.get_device_snapshot(); wacore::send::prepare_peer_stanza( @@ -1981,7 +1981,7 @@ impl Client { .await; } - let device_store_arc = self.persistence_manager.get_device_arc().await; + let device_store_arc = self.persistence_manager.clone(); let to_str = to.to_string(); let (own_sending_jid, _) = match group_info.addressing_mode { @@ -2427,7 +2427,7 @@ impl Client { let lock_jids = self.build_session_lock_keys(dm_devices.devices()).await; let _session_guards = self.session_guards_for(&lock_jids).await; - let mut store_adapter = self.signal_adapter().await; + let mut store_adapter = self.signal_adapter(); let mut stores = store_adapter.as_signal_stores(); @@ -4907,7 +4907,7 @@ mod tests { .expect("prekey bundle task") .expect("prekey bundle"); - let mut adapter = client.signal_adapter().await; + let mut adapter = client.signal_adapter(); let mut rng = rand::make_rng::(); process_prekey_bundle( &peer.to_protocol_address(), @@ -5124,7 +5124,7 @@ mod tests { .expect("prekey bundle task") .expect("prekey bundle"); { - let mut adapter = client.signal_adapter().await; + let mut adapter = client.signal_adapter(); let mut rng = rand::make_rng::(); process_prekey_bundle( &lid_addr.to_protocol_address(), diff --git a/src/store/persistence_manager.rs b/src/store/persistence_manager.rs index 5bc658f91..5c0b0fefd 100644 --- a/src/store/persistence_manager.rs +++ b/src/store/persistence_manager.rs @@ -74,7 +74,7 @@ impl PersistenceManager { }) } - /// Handle for store adapters that need `&mut Device` trait access. + /// Handle for callers that need `&mut Device` trait access directly. /// For plain reads, prefer [`get_device_snapshot`](Self::get_device_snapshot). pub async fn get_device_arc(&self) -> Arc> { self.device.clone() diff --git a/src/store/signal_adapter.rs b/src/store/signal_adapter.rs index 90d820254..5d832da62 100644 --- a/src/store/signal_adapter.rs +++ b/src/store/signal_adapter.rs @@ -1,6 +1,6 @@ use crate::store::Device; +use crate::store::persistence_manager::PersistenceManager; use crate::store::signal_cache::SignalStoreCache; -use async_lock::RwLock; use async_trait::async_trait; use std::sync::Arc; use wacore::libsignal::protocol::{ @@ -33,12 +33,25 @@ type BoxFut<'a, T> = std::pin::Pin + 'a>>; use std::future::{Future, ready}; +/// Snapshots per call instead of holding the device lock: `get_device_snapshot` +/// is a refcount bump, so a backend round-trip no longer keeps a read guard +/// that every later Signal read would queue behind once a write arrives. +/// +/// Per call, not per adapter: `signed_pre_key_id` is promoted by rotation and +/// its staged backend row is then dropped, so an adapter pinned to a +/// pre-rotation snapshot would fail to load the very key a peer just used. #[derive(Clone)] struct SharedDevice { - device: Arc>, + persistence_manager: Arc, cache: Arc, } +impl SharedDevice { + fn device(&self) -> Arc { + self.persistence_manager.get_device_snapshot() + } +} + #[derive(Clone)] pub struct SessionAdapter(SharedDevice); #[derive(Clone)] @@ -55,8 +68,11 @@ impl SenderKeyAdapter { /// Build a standalone sender-key store without constructing the full /// five-store [`SignalProtocolStoreAdapter`]. Used on the SKDM-processing /// path, which only needs the sender-key store. - pub fn new(device: Arc>, cache: Arc) -> Self { - Self(SharedDevice { device, cache }) + pub fn new(persistence_manager: Arc, cache: Arc) -> Self { + Self(SharedDevice { + persistence_manager, + cache, + }) } } @@ -70,8 +86,11 @@ pub struct SignalProtocolStoreAdapter { } impl SignalProtocolStoreAdapter { - pub fn new(device: Arc>, cache: Arc) -> Self { - let shared = SharedDevice { device, cache }; + pub fn new(persistence_manager: Arc, cache: Arc) -> Self { + let shared = SharedDevice { + persistence_manager, + cache, + }; Self { session_store: SessionAdapter(shared.clone()), identity_store: IdentityAdapter(shared.clone()), @@ -99,7 +118,7 @@ impl SessionStore for SessionAdapter { &self, address: &ProtocolAddress, ) -> Result, SignalProtocolError> { - let device = self.0.device.read().await; + let device = self.0.device(); self.0 .cache .peek_session(address, &*device.backend) @@ -112,7 +131,7 @@ impl SessionStore for SessionAdapter { &self, address: &ProtocolAddress, ) -> Result<(Option, Option), SignalProtocolError> { - let device = self.0.device.read().await; + let device = self.0.device(); self.0 .cache .checkout_session(address, &*device.backend) @@ -184,7 +203,7 @@ impl SessionStore for SessionAdapter { return Box::pin(ready(answer)); } Box::pin(async move { - let device = self.0.device.read().await; + let device = self.0.device(); self.0 .cache .has_session(address, &*device.backend) @@ -247,15 +266,15 @@ impl IdentityAdapter { #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl IdentityKeyStore for IdentityAdapter { async fn get_identity_key_pair(&self) -> Result { - let device = self.0.device.read().await; - IdentityKeyStore::get_identity_key_pair(&*device) + let device = self.0.device(); + IdentityKeyStore::get_identity_key_pair(device.as_ref()) .await .map_err(signal_err("get_identity_key_pair")) } async fn get_local_registration_id(&self) -> Result { - let device = self.0.device.read().await; - IdentityKeyStore::get_local_registration_id(&*device) + let device = self.0.device(); + IdentityKeyStore::get_local_registration_id(device.as_ref()) .await .map_err(signal_err("get_local_registration_id")) } @@ -345,7 +364,7 @@ impl IdentityKeyStore for IdentityAdapter { return Box::pin(ready(parse_cached_identity(cached))); } Box::pin(async move { - let device = self.0.device.read().await; + let device = self.0.device(); let data = self .0 .cache @@ -376,8 +395,8 @@ fn parse_cached_identity( #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl PreKeyStore for PreKeyAdapter { async fn get_pre_key(&self, prekey_id: PreKeyId) -> Result { - let device = self.0.device.read().await; - WacorePreKeyStore::load_prekey(&*device, prekey_id.into()) + let device = self.0.device(); + WacorePreKeyStore::load_prekey(device.as_ref(), prekey_id.into()) .await .map_err(signal_err("backend"))? .ok_or(SignalProtocolError::InvalidPreKeyId) @@ -388,9 +407,9 @@ impl PreKeyStore for PreKeyAdapter { prekey_id: PreKeyId, record: &PreKeyRecord, ) -> Result<(), SignalProtocolError> { - let device = self.0.device.read().await; + let device = self.0.device(); let structure = wacore_record::prekey_record_to_structure(record)?; - WacorePreKeyStore::store_prekey(&*device, prekey_id.into(), structure, false) + WacorePreKeyStore::store_prekey(device.as_ref(), prekey_id.into(), structure, false) .await .map_err(signal_err("backend")) } @@ -399,7 +418,7 @@ impl PreKeyStore for PreKeyAdapter { // through here: message_decrypt reports the consumed prekey and the receive // path buffers it via buffer_consumed_prekey so the durable delete is // atomic with the session flush (matching WAWebSignalProtocolStoreUnifiedApi). - let device = self.0.device.read().await; + let device = self.0.device(); device .backend .remove_prekey(prekey_id.into()) @@ -429,10 +448,20 @@ impl SignedPreKeyStore for SignedPreKeyAdapter { &self, signed_prekey_id: SignedPreKeyId, ) -> Result { - let device = self.0.device.read().await; - WacoreSignedPreKeyStore::load_signed_prekey(&*device, signed_prekey_id.into()) + let id = signed_prekey_id.into(); + let mut record = WacoreSignedPreKeyStore::load_signed_prekey(self.0.device().as_ref(), id) .await - .map_err(signal_err("backend"))? + .map_err(signal_err("backend"))?; + if record.is_none() { + // Rotation promotes an id into the device field and only then drops + // its staged row, so a snapshot taken just before the promotion + // resolves it in neither place. Re-read before calling it unknown: + // once the row is gone the field definitely holds it. + record = WacoreSignedPreKeyStore::load_signed_prekey(self.0.device().as_ref(), id) + .await + .map_err(signal_err("backend"))?; + } + record .ok_or(SignalProtocolError::InvalidSignedPreKeyId) .and_then(wacore_record::signed_prekey_structure_to_record) } @@ -463,7 +492,7 @@ impl wacore::libsignal::protocol::SenderKeyStore for SenderKeyAdapter { ) -> wacore::libsignal::protocol::error::Result< Option, > { - let device = self.0.device.read().await; + let device = self.0.device(); // group_decrypt mutates the loaded record (catch-up + ratchet) and stores // it back, so the trait needs an owned copy. The cache keeps its `Arc`, so // this clones the inner record (unchanged from the prior behavior). @@ -490,11 +519,123 @@ impl wacore::libsignal::protocol::SenderKeyStore for SenderKeyAdapter { #[cfg(test)] mod tests { use super::*; - use crate::store::Device; use wacore::store::in_memory::InMemoryBackend; const PREKEY_ID: u32 = 7777; + async fn test_persistence_manager( + backend: Arc, + ) -> Arc { + Arc::new( + PersistenceManager::new(backend) + .await + .expect("in-memory persistence manager"), + ) + } + + /// Rotation promotes the new id into the device field and only then drops + /// its staged backend row, so a snapshot taken before the promotion + /// resolves that id in neither place. This is why the adapter re-reads the + /// snapshot before reporting an id unknown: a fresh one resolves what a + /// stale one cannot, and a decrypt that raced the promotion would otherwise + /// reject a valid pre-key message. + #[tokio::test] + async fn a_promoted_signed_pre_key_is_only_visible_to_a_fresh_snapshot() { + use wacore::libsignal::protocol::KeyPair; + use wacore::store::commands::DeviceCommand; + + let backend: Arc = Arc::new(InMemoryBackend::new()); + let pm = test_persistence_manager(backend).await; + + let stale = pm.get_device_snapshot(); + let promoted_id = stale.signed_pre_key_id + 1; + + // Nothing staged for the new id, and it is not the snapshot's current + // one: the pre-promotion view cannot resolve it. + assert!( + WacoreSignedPreKeyStore::load_signed_prekey(stale.as_ref(), promoted_id) + .await + .expect("load") + .is_none(), + "the new id must be unresolvable before promotion" + ); + + let key_pair = KeyPair::generate(&mut rand::rng()); + pm.process_command(DeviceCommand::SetSignedPreKey { + key_pair, + id: promoted_id, + signature: [0u8; 64], + rotation_ms: 0, + }) + .await; + + // The stale snapshot still cannot resolve it, which is exactly the + // failure a pinned snapshot would produce. + assert!( + WacoreSignedPreKeyStore::load_signed_prekey(stale.as_ref(), promoted_id) + .await + .expect("load") + .is_none(), + "a stale snapshot must not resolve the promoted id" + ); + + let adapter = SignalProtocolStoreAdapter::new(pm, Arc::new(SignalStoreCache::new())); + assert!( + adapter + .signed_pre_key_store + .get_signed_pre_key(promoted_id.into()) + .await + .is_ok(), + "the adapter must resolve the promoted id from a fresh snapshot" + ); + } + + /// The window the re-read closes is intra-call: the promotion can land + /// after the first snapshot is taken and before its backend lookup + /// resolves. Parks that lookup, promotes, then releases, so the retry is + /// the only thing that can produce an answer. + #[tokio::test] + async fn a_promotion_racing_a_lookup_is_resolved_by_the_retry() { + use wacore::libsignal::protocol::KeyPair; + use wacore::store::commands::DeviceCommand; + + let backend = Arc::new(InMemoryBackend::new()); + let pm = test_persistence_manager(backend.clone()).await; + let promoted_id = pm.get_device_snapshot().signed_pre_key_id + 1; + + // Built before the promotion, so its first snapshot is the stale one. + let adapter = + SignalProtocolStoreAdapter::new(pm.clone(), Arc::new(SignalStoreCache::new())); + + let gate = Arc::new(async_lock::Barrier::new(2)); + backend.gate_next_signed_prekey_read(gate.clone()); + + let lookup = tokio::spawn(async move { + adapter + .signed_pre_key_store + .get_signed_pre_key(promoted_id.into()) + .await + }); + + // The first load is now parked inside the backend, holding a snapshot + // that predates everything below. + gate.wait().await; + let key_pair = KeyPair::generate(&mut rand::rng()); + pm.process_command(DeviceCommand::SetSignedPreKey { + key_pair, + id: promoted_id, + signature: [0u8; 64], + rotation_ms: 0, + }) + .await; + gate.wait().await; + + assert!( + lookup.await.expect("lookup task").is_ok(), + "the retry must resolve an id promoted mid-lookup" + ); + } + /// The inbound decrypt path consumes a one-time prekey and buffers it via /// `buffer_consumed_prekey`. It must NOT delete the prekey from the backend /// synchronously: the promoted session is still volatile at that point, so an @@ -508,9 +649,9 @@ mod tests { .await .unwrap(); - let device = Arc::new(RwLock::new(Device::new(backend.clone()))); + let device = test_persistence_manager(backend.clone()).await; let cache = Arc::new(SignalStoreCache::new()); - let adapter = SignalProtocolStoreAdapter::new(device, cache.clone()); + let adapter = SignalProtocolStoreAdapter::new(device.clone(), cache.clone()); let addr = ProtocolAddress::new("bob", 1.into()); // The real path stores the promoted session before buffering the prekey. @@ -544,9 +685,9 @@ mod tests { .await .unwrap(); - let device = Arc::new(RwLock::new(Device::new(backend.clone()))); + let device = test_persistence_manager(backend.clone()).await; let cache = Arc::new(SignalStoreCache::new()); - let mut adapter = SignalProtocolStoreAdapter::new(device, cache.clone()); + let mut adapter = SignalProtocolStoreAdapter::new(device.clone(), cache.clone()); adapter .pre_key_store @@ -560,10 +701,10 @@ mod tests { ); } - fn test_adapter() -> SignalProtocolStoreAdapter { + async fn test_adapter() -> SignalProtocolStoreAdapter { let backend: Arc = Arc::new(InMemoryBackend::new()); - let device = Arc::new(RwLock::new(Device::new(backend))); - SignalProtocolStoreAdapter::new(device, Arc::new(SignalStoreCache::new())) + let device = test_persistence_manager(backend).await; + SignalProtocolStoreAdapter::new(device.clone(), Arc::new(SignalStoreCache::new())) } struct BlockingIdentityStore { @@ -650,9 +791,9 @@ mod tests { seed_durable, ); - let device = Arc::new(RwLock::new(Device::new(backend.clone()))); + let device = test_persistence_manager(backend.clone()).await; let mut session_store = - SignalProtocolStoreAdapter::new(device, cache.clone()).session_store; + SignalProtocolStoreAdapter::new(device.clone(), cache.clone()).session_store; let entered = Arc::new(async_lock::Barrier::new(2)); let mut identity_store = BlockingIdentityStore { pair: identity_pair, @@ -715,8 +856,8 @@ mod tests { let backend: Arc = Arc::new(InMemoryBackend::new()); let cache = Arc::new(SignalStoreCache::new()); - let device = Arc::new(RwLock::new(Device::new(backend))); - let mut adapter = SignalProtocolStoreAdapter::new(device, cache.clone()); + let device = test_persistence_manager(backend).await; + let mut adapter = SignalProtocolStoreAdapter::new(device.clone(), cache.clone()); let address = ProtocolAddress::new("15550005555", 1.into()); cache .put_session(&address, SessionRecord::new_fresh()) @@ -742,7 +883,7 @@ mod tests { #[tokio::test] async fn session_store_fast_paths_round_trip() { use wacore::libsignal::protocol::SessionStore as _; - let mut adapter = test_adapter(); + let mut adapter = test_adapter().await; let addr = ProtocolAddress::new("15550002222", 1.into()); // Cold cache: goes through the async fallback (backend consult). @@ -782,7 +923,7 @@ mod tests { #[tokio::test] async fn identity_fast_paths_keep_change_semantics() { use wacore::libsignal::protocol::{IdentityKeyPair, IdentityKeyStore as _}; - let mut adapter = test_adapter(); + let mut adapter = test_adapter().await; let addr = ProtocolAddress::new("15550003333", 1.into()); let mut rng = rand::make_rng::(); @@ -839,15 +980,18 @@ mod tests { #[cfg(test)] mod hook_alloc_tests { use super::*; - use crate::store::Device; use crate::test_alloc::min_allocs; use wacore::libsignal::protocol::{Direction, IdentityKeyPair}; use wacore::store::in_memory::InMemoryBackend; - fn adapter_for_test() -> SignalProtocolStoreAdapter { + async fn adapter_for_test() -> SignalProtocolStoreAdapter { let backend: Arc = Arc::new(InMemoryBackend::new()); - let device = Arc::new(RwLock::new(Device::new(backend))); - SignalProtocolStoreAdapter::new(device, Arc::new(SignalStoreCache::new())) + let device = Arc::new( + PersistenceManager::new(backend) + .await + .expect("in-memory persistence manager"), + ); + SignalProtocolStoreAdapter::new(device.clone(), Arc::new(SignalStoreCache::new())) } fn some_identity() -> IdentityKey { @@ -858,9 +1002,9 @@ mod hook_alloc_tests { /// `#[async_trait]` still boxes that future once per encrypt and once per /// decrypt. The hook exists to skip the box, and the only way to know it /// does is to count. - #[test] - fn the_trusted_identity_hook_answers_without_allocating() { - let adapter = adapter_for_test(); + #[tokio::test] + async fn the_trusted_identity_hook_answers_without_allocating() { + let adapter = adapter_for_test().await; let address = ProtocolAddress::new("bob@s.whatsapp.net", 1.into()); let identity = some_identity(); @@ -885,9 +1029,9 @@ mod hook_alloc_tests { /// Bad path: with nothing cached for the address, the hook must decline so /// the caller takes the async path that can read the backend. Answering /// from an empty cache would report every identity as new. - #[test] - fn the_save_identity_hook_declines_when_nothing_is_cached() { - let mut adapter = adapter_for_test(); + #[tokio::test] + async fn the_save_identity_hook_declines_when_nothing_is_cached() { + let mut adapter = adapter_for_test().await; let address = ProtocolAddress::new("never-seen@s.whatsapp.net", 1.into()); let identity = some_identity(); @@ -905,7 +1049,7 @@ mod hook_alloc_tests { /// correctly, which is what lets the caller skip the box. #[tokio::test] async fn the_save_identity_hook_answers_once_the_entry_is_cached() { - let mut adapter = adapter_for_test(); + let mut adapter = adapter_for_test().await; let address = ProtocolAddress::new("bob@s.whatsapp.net", 1.into()); let first = some_identity(); let second = some_identity(); @@ -936,9 +1080,9 @@ mod hook_alloc_tests { /// The session hook has the same contract: decline when the cache cannot /// answer, rather than reporting "no session" and forcing a needless /// session rebuild. - #[test] - fn the_session_hook_declines_when_the_cache_cannot_answer() { - let adapter = adapter_for_test(); + #[tokio::test] + async fn the_session_hook_declines_when_the_cache_cannot_answer() { + let adapter = adapter_for_test().await; let address = ProtocolAddress::new("never-seen@s.whatsapp.net", 1.into()); assert!( diff --git a/src/test_utils.rs b/src/test_utils.rs index 68d01c61a..fb8b20cac 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -233,7 +233,7 @@ pub async fn seed_peer_session(client: &Arc, peer: &Jid) { .expect("bundle task") .expect("bundle"); - let mut adapter = client.signal_adapter().await; + let mut adapter = client.signal_adapter(); let mut rng = rand::make_rng::(); process_prekey_bundle( &peer.to_protocol_address(), diff --git a/src/voip/facade.rs b/src/voip/facade.rs index 196cefa6d..9e60724ed 100644 --- a/src/voip/facade.rs +++ b/src/voip/facade.rs @@ -1471,7 +1471,7 @@ async fn fanout_group_epoch_for_generation( } } let plan = wacore::send::SessionPlan::assume_ready(recipients.len()); - let mut adapter = client.signal_adapter().await; + let mut adapter = client.signal_adapter(); let mut stores = adapter.as_signal_stores(); let encrypted = wacore::send::encrypt_for_devices_with_sessions_raw( &*client.runtime, @@ -1654,7 +1654,7 @@ async fn place_call( // encrypt against the existing sessions directly: a device whose session is somehow still // missing fails its encrypt and is skipped, exactly as the old per-device loop did. let plan = wacore::send::SessionPlan::assume_ready(devices.len()); - let mut adapter = client.signal_adapter().await; + let mut adapter = client.signal_adapter(); let mut stores = adapter.as_signal_stores(); let raw = wacore::send::encrypt_for_devices_with_sessions_raw( &*client.runtime, diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index c7f93479f..6ac3e259b 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -141,6 +141,11 @@ pub struct InMemoryBackend { /// flush). #[cfg(any(test, feature = "test-util"))] fail_sender_key_writes: AtomicBool, + /// Parks the next `load_signed_prekey` on a barrier. Test hook: lets a + /// harness promote a rotated key while a lookup is already in flight, + /// which is the window a caller's re-read exists to close. + #[cfg(any(test, feature = "test-util"))] + signed_prekey_read_gate: std::sync::Mutex>>, } impl InMemoryBackend { @@ -157,6 +162,8 @@ impl InMemoryBackend { fail_session_writes: AtomicBool::new(false), #[cfg(any(test, feature = "test-util"))] fail_sender_key_writes: AtomicBool::new(false), + #[cfg(any(test, feature = "test-util"))] + signed_prekey_read_gate: std::sync::Mutex::new(None), } } @@ -186,6 +193,16 @@ impl InMemoryBackend { self.fail_sender_key_writes.store(fail, Ordering::Relaxed); } + /// Park the next `load_signed_prekey` on `gate`, which it waits on twice: + /// once to signal it has arrived, once to be released. + #[cfg(any(test, feature = "test-util"))] + pub fn gate_next_signed_prekey_read(&self, gate: Arc) { + *self + .signed_prekey_read_gate + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(gate); + } + /// Lets recovery tests remove only the state needed to trigger a key request. #[cfg(any(test, feature = "test-util"))] pub async fn remove_sync_key_for_test(&self, key_id: &[u8]) -> bool { @@ -362,6 +379,20 @@ impl SignalStore for InMemoryBackend { } async fn load_signed_prekey(&self, id: u32) -> Result>> { + #[cfg(any(test, feature = "test-util"))] + { + // Taken, not borrowed, so the guard never crosses the await and + // only the first read after arming is gated. + let gate = self + .signed_prekey_read_gate + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + if let Some(gate) = gate { + gate.wait().await; + gate.wait().await; + } + } Ok(self.state.lock().await.signed_prekeys.get(&id).cloned()) } diff --git a/wacore/src/store/signal_cache.rs b/wacore/src/store/signal_cache.rs index d0980235e..df660b01a 100644 --- a/wacore/src/store/signal_cache.rs +++ b/wacore/src/store/signal_cache.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::num::NonZeroU64; use std::sync::{Arc, Mutex as SyncMutex, MutexGuard as SyncMutexGuard}; @@ -31,11 +31,12 @@ fn new_store_incarnation() -> StoreIncarnation { /// that grows the map, including read-populate (cache-miss) inserts, so the cache /// stays bounded even under unique-key read floods; the early-out keeps it cheap. fn evict_clean_entries( - cache: &mut HashMap, Option>, + cache: &mut UserIndexedCache>, dirty: &HashSet>, deleted: Option<&HashSet>>, max_entries: usize, ) { + compact_users_if_needed(cache, max_entries); if cache.len() <= high_watermark(max_entries) { return; } @@ -62,9 +63,35 @@ fn evict_clean_entries( } } +/// Rebuild the user superset once eviction has let it drift a whole watermark +/// past the entries backing it. A rebuild leaves it no larger than the live key +/// count, so the next one is that many inserts away. +fn compact_users_if_needed(cache: &mut UserIndexedCache, max_entries: usize) { + // Both conditions matter. The watermark keeps the rebuild rare; requiring + // the index to exceed the live key count keeps it self-limiting, since a + // rebuild always lands at or below that count. Without it, a store holding + // more distinct users than the watermark — dirty entries that eviction + // cannot trim, as a run of failing flushes produces — would rebuild on + // every single update, under the global mutex. + if cache.users_len() > high_watermark(max_entries) && cache.users_len() > cache.len() { + cache.compact_users(); + } +} + /// Default max entries per store before clean entry eviction triggers. const DEFAULT_MAX_CACHE_ENTRIES: usize = 2_000; +/// Removals retained for cold readers to consult. A reader that spans more +/// than this many removals is told its key went, which costs a re-read; the +/// window is what keeps the bookkeeping fixed-size and free of any per-reader +/// state that a cancelled read could strand. +const RECENT_REMOVALS: usize = 64; + +/// Unlocked cold sender-key reads to try before falling back to reading under +/// the lock. Losing twice means a removal landed in both windows, which needs a +/// flush plus an eviction each time; the fallback keeps that bounded. +const SENDER_KEY_UNLOCKED_READ_ATTEMPTS: usize = 2; + /// Slack above `max_entries` the cache may grow to before an eviction scan /// fires, expressed as a divisor of `max_entries` (1/8th here). Trimming back /// to `max_entries` then amortizes the O(n) scan over this many inserts. A @@ -84,6 +111,191 @@ fn protocol_address_matches_user(address: &str, user: &str) -> bool { .is_some_and(|suffix| suffix.starts_with('@') || suffix.starts_with(':')) } +/// The user half of a protocol address, matching the prefix +/// [`protocol_address_matches_user`] tests. +fn user_of_protocol_address(address: &str) -> &str { + match address.find(['@', ':']) { + Some(end) => &address[..end], + None => address, + } +} + +/// A cache map that can answer "is any address here owned by this user?" +/// without scanning every key. +/// +/// The user set is a deliberate superset: removals leave it untouched, so its +/// only error is a `true` for a user whose last entry is gone, costing one +/// migration pass that finds nothing. It can never answer "no state" for a +/// user that has some, which is the direction that would silently skip a +/// migration. `insert` is the only way in, so an entry cannot reach the map +/// without registering its user. +struct UserIndexedCache { + map: HashMap, V>, + users: HashSet>, + /// Counts removal events. A cold reader that saw a slot absent, released + /// the lock, and finds it absent again cannot otherwise tell "never + /// written" from "written, flushed, and evicted": a clean removal keeps the + /// incarnation, so its pre-write bytes would be trusted as an exact reload. + removal_seq: u64, + /// The last `RECENT_REMOVALS` removals, newest last, so a reader can ask + /// about its own key rather than about cache-wide churn. A fixed window + /// needs no per-reader registration, so a cancelled read leaves nothing + /// behind; a reader older than the window is told "removed" instead. + recent_removals: VecDeque<(u64, Arc)>, + /// Sequence of the last removal that could not name its keys (`clear`, + /// `retain`), which invalidates every reader older than it. + opaque_removal_seq: u64, +} + +impl UserIndexedCache { + fn new() -> Self { + Self { + map: HashMap::new(), + users: HashSet::new(), + removal_seq: 0, + recent_removals: VecDeque::new(), + opaque_removal_seq: 0, + } + } + + fn insert(&mut self, key: Arc, value: V) -> Option { + let user = user_of_protocol_address(&key); + if !self.users.contains(user) { + self.users.insert(Arc::from(user)); + } + self.map.insert(key, value) + } + + /// Stamp to take before a cold read releases the lock. + fn removal_seq(&self) -> u64 { + self.removal_seq + } + + /// Whether `key` was removed after `since`. Conservative in two places: a + /// removal that could not name its keys, and a reader older than the + /// retained window. Both answer "removed", which costs a re-read rather + /// than admitting bytes that predate a write. + fn removed_since(&self, key: &str, since: u64) -> bool { + if self.opaque_removal_seq > since { + return true; + } + if self.removal_seq.saturating_sub(since) > RECENT_REMOVALS as u64 { + return true; + } + self.recent_removals + .iter() + .any(|(seq, removed)| *seq > since && removed.as_ref() == key) + } + + fn note_removal(&mut self, key: Arc) { + self.removal_seq += 1; + if self.recent_removals.len() == RECENT_REMOVALS { + self.recent_removals.pop_front(); + } + self.recent_removals.push_back((self.removal_seq, key)); + } + + fn note_opaque_removal(&mut self) { + self.removal_seq += 1; + self.opaque_removal_seq = self.removal_seq; + } + + fn has_user(&self, user: &str) -> bool { + // Normalize the query the way the keys were derived. A matching + // address begins with `user`, so a separator inside `user` is also the + // address's first one and both collapse to the same key: an addressed + // `19995551006:5` and a bare `19995551006` are one entry here. + self.users.contains(user_of_protocol_address(user)) + } + + /// Drop users no longer backed by an entry. Bounds the superset's drift + /// after eviction; callers gate it on a watermark so it stays amortized. + fn compact_users(&mut self) { + self.users.clear(); + let users: Vec> = self + .map + .keys() + .map(|key| Arc::from(user_of_protocol_address(key))) + .collect(); + self.users.extend(users); + } + + fn users_len(&self) -> usize { + self.users.len() + } + + /// Retained bytes of the bookkeeping beside the primary map: the user + /// index, plus the removal window, whose keys outlive the entries they + /// name and so are owned solely here. Keeps `memory_stats` from reporting + /// only the map. + fn overhead_bytes(&self) -> usize { + self.users.iter().map(|user| user.len()).sum::() + + self + .recent_removals + .iter() + .map(|(_, key)| key.len() + size_of::()) + .sum::() + } + + fn get(&self, key: &str) -> Option<&V> { + self.map.get(key) + } + + fn get_mut(&mut self, key: &str) -> Option<&mut V> { + self.map.get_mut(key) + } + + fn get_key_value(&self, key: &str) -> Option<(&Arc, &V)> { + self.map.get_key_value(key) + } + + fn contains_key(&self, key: &str) -> bool { + self.map.contains_key(key) + } + + fn remove(&mut self, key: &str) -> Option { + let (key, value) = self.map.remove_entry(key)?; + // Reuses the map's own `Arc`, so recording a removal allocates + // nothing. + self.note_removal(key); + Some(value) + } + + fn retain(&mut self, keep: impl FnMut(&Arc, &mut V) -> bool) { + let before = self.map.len(); + self.map.retain(keep); + if self.map.len() != before { + // `retain` does not report which keys went. + self.note_opaque_removal(); + } + } + + fn clear(&mut self) { + if !self.map.is_empty() { + self.note_opaque_removal(); + } + self.map.clear(); + self.users.clear(); + } + + fn len(&self) -> usize { + self.map.len() + } + + fn is_empty(&self) -> bool { + self.map.is_empty() + } + + fn iter(&self) -> impl Iterator, &V)> { + self.map.iter() + } + + #[cfg(test)] + fn values(&self) -> impl Iterator { + self.map.values() + } +} + /// In-memory write-back cache for Signal protocol state. /// Keys use `Arc` for O(1) clone. Sessions cached as objects (serialized on flush). /// Capacity-bounded: every path that grows a store (writes and read-populate @@ -150,7 +362,7 @@ struct SessionStoreState { incarnation: StoreIncarnation, checkout_generation: u64, next_checkout_token: u64, - cache: HashMap, SessionEntry>, + cache: UserIndexedCache, dirty: HashSet>, deleted: HashSet>, /// Sessions whose raised counter reservation has not reached the backend @@ -168,7 +380,7 @@ impl SessionStoreState { incarnation, checkout_generation: 0, next_checkout_token: 1, - cache: HashMap::new(), + cache: UserIndexedCache::new(), dirty: HashSet::new(), deleted: HashSet::new(), reservation_pending: HashSet::new(), @@ -258,6 +470,10 @@ impl SessionStoreState { fn clear_clean_entries(&mut self) { self.cache .retain(|_, entry| matches!(entry, SessionEntry::CheckedOut { .. })); + // Teardown drops nearly every entry at once and is not a hot path, so + // settle the user superset here instead of letting it drift until the + // watermark rebuild. + self.cache.compact_users(); } fn discard(&mut self, incarnation: StoreIncarnation, generation: u64) { @@ -267,6 +483,7 @@ impl SessionStoreState { } fn evict_if_needed(&mut self, max_entries: usize) { + compact_users_if_needed(&mut self.cache, max_entries); if self.cache.len() <= high_watermark(max_entries) { return; } @@ -304,7 +521,7 @@ struct SenderKeyStoreState { // `Arc`-wrapped so a warm `get_sender_key` (the per-send peek reads and the // per-decrypt load) bumps a refcount instead of deep-cloning the record's // `VecDeque` with up to `MAX_MESSAGE_KEYS` message keys each. - cache: HashMap, Option>>, + cache: UserIndexedCache>>, dirty: HashSet>, /// Chains whose outbound iteration lease was raised but not yet persisted; /// the send path must flush before the wire while any entry is here. @@ -323,7 +540,7 @@ impl SenderKeyStoreState { fn new(incarnation: StoreIncarnation) -> Self { Self { incarnation, - cache: HashMap::new(), + cache: UserIndexedCache::new(), dirty: HashSet::new(), wire_gate_pending: HashSet::new(), pending_distributions: HashMap::new(), @@ -375,7 +592,7 @@ impl SenderKeyStoreState { struct ByteStoreState { /// Cached entries. `None` value = known-absent (negative cache). - cache: HashMap, Option>>, + cache: UserIndexedCache>>, dirty: HashSet>, deleted: HashSet>, } @@ -383,7 +600,7 @@ struct ByteStoreState { impl ByteStoreState { fn new() -> Self { Self { - cache: HashMap::new(), + cache: UserIndexedCache::new(), dirty: HashSet::new(), deleted: HashSet::new(), } @@ -637,6 +854,16 @@ impl SignalStoreCache { drop(self.lock_sessions().await); } + /// Model a flush followed by a capacity eviction or `clear_after_flush`: + /// the chain becomes clean and then leaves the cache outright. + #[cfg(test)] + async fn drop_clean_sender_key_for_test(&self, cache_key: &str) { + let mut state = self.sender_keys.lock().await; + state.dirty.remove(cache_key); + state.wire_gate_pending.remove(cache_key); + state.cache.remove(cache_key); + } + /// Whether any session or identity is known for `user` (across device ids), /// checking the in-memory cache first, then the durable backend. Lets a /// caller skip a per-device migration scan for a user we've never had Signal @@ -644,25 +871,11 @@ impl SignalStoreCache { /// (even a stale/checked-out marker), so it never reports "none" when state /// might exist. pub async fn has_state_for_user(&self, user: &str, backend: &dyn SignalStore) -> Result { - { - let state = self.lock_sessions().await; - if state - .cache - .keys() - .any(|address| protocol_address_matches_user(address, user)) - { - return Ok(true); - } + if self.lock_sessions().await.cache.has_user(user) { + return Ok(true); } - { - let state = self.identities.lock().await; - if state - .cache - .keys() - .any(|address| protocol_address_matches_user(address, user)) - { - return Ok(true); - } + if self.identities.lock().await.cache.has_user(user) { + return Ok(true); } Ok(backend.has_signal_state_for_user(user).await?) } @@ -1014,6 +1227,50 @@ impl SignalStoreCache { backend: &dyn SignalStore, ) -> Result>> { let key = name.cache_key(); + for _ in 0..SENDER_KEY_UNLOCKED_READ_ATTEMPTS { + let (incarnation, since) = { + let state = self.sender_keys.lock().await; + if let Some(cached) = state.cache.get(key) { + return Ok(cached.clone()); + } + (state.incarnation, state.cache.removal_seq()) + }; + + // Decoding stays outside the lock too: a cold chain can carry + // MAX_MESSAGE_KEYS skipped keys, and parsing them is the expensive + // half of the miss. Errors are held rather than raised, so a + // concurrent write still wins the re-check below: an unreadable row + // must not fail an operation the cache can already answer. + let decoded: Result>> = + match backend.get_sender_key(key).await { + Ok(Some(bytes)) => SenderKeyRecord::deserialize_for_store(&bytes, &incarnation) + .map(|record| Some(Arc::new(record))) + .map_err(anyhow::Error::from), + Ok(None) => Ok(None), + Err(error) => Err(anyhow::Error::from(error)), + }; + + let mut state = self.sender_keys.lock().await; + // A put or delete that landed while we awaited describes the chain + // more recently than the bytes we read, so it wins; we would + // otherwise resurrect a deleted chain or undo a fresher iteration. + if let Some(cached) = state.cache.get(key) { + return Ok(cached.clone()); + } + // Still absent, but that is only trustworthy if this key was not + // removed and no lossy clear reset the incarnation meanwhile. + // Either could mean a newer record was written and then dropped, + // leaving our older bytes to be adopted as an exact reload and let + // the chain resume an iteration that has already been published. + if state.incarnation == incarnation && !state.cache.removed_since(key, since) { + let record = decoded?; + state.cache.insert(Arc::from(key), record.clone()); + state.evict_if_needed(self.max_entries); + return Ok(record); + } + } + + // Repeatedly raced. Read under the lock, which cannot be raced at all. let mut state = self.sender_keys.lock().await; if let Some(cached) = state.cache.get(key) { return Ok(cached.clone()); @@ -1432,6 +1689,7 @@ impl SignalStoreCache { } }) .collect(); + keys_len += s.cache.overhead_bytes(); (s.cache.len() as u64, keys_len, recs) }; let session_bytes: usize = session_keys_len @@ -1447,7 +1705,8 @@ impl SignalStoreCache { .cache .iter() .map(|(k, v)| k.len() + v.as_ref().map_or(0, |b| b.len())) - .sum(); + .sum::() + + i.cache.overhead_bytes(); CollectionStats::new(i.cache.len() as u64, bytes as u64) }; @@ -1474,7 +1733,7 @@ impl SignalStoreCache { .fold((0usize, 0usize), |(count, bytes), key| { (count + 1, bytes + key.len()) }); - keys_len += pending_only_key_bytes; + keys_len += pending_only_key_bytes + sk.cache.overhead_bytes(); ( (sk.cache.len() + pending_only_count) as u64, keys_len, @@ -1654,6 +1913,747 @@ mod sender_key_lock_tests { } } + /// Sender-key backend whose read parks until every expected reader has + /// arrived, so a test can prove that N cold readers reach it concurrently + /// rather than queueing on the global cache mutex. + struct GatedSenderKeyLookup { + arrived: async_lock::Barrier, + release: async_lock::Barrier, + hits: std::sync::atomic::AtomicUsize, + gated_reads: usize, + payload: SyncMutex>>, + } + + impl GatedSenderKeyLookup { + fn new(readers: usize, payload: Option>) -> Self { + Self::with_rounds(readers, readers, payload) + } + + /// `readers` sizes the rendezvous (plus the driving test task); + /// `gated_reads` is how many backend calls park on it, which is a + /// different number when one reader is made to read repeatedly. + fn with_rounds(readers: usize, gated_reads: usize, payload: Option>) -> Self { + Self { + arrived: async_lock::Barrier::new(readers + 1), + release: async_lock::Barrier::new(readers + 1), + hits: std::sync::atomic::AtomicUsize::new(0), + gated_reads, + payload: SyncMutex::new(payload), + } + } + + fn hits(&self) -> usize { + self.hits.load(Ordering::Relaxed) + } + + /// Model the flush that makes a concurrent write durable, so a reader + /// forced to read again sees what a real backend would now hold. + fn set_payload(&self, payload: Option>) { + *self.payload.lock().unwrap_or_else(|p| p.into_inner()) = payload; + } + } + + #[async_trait::async_trait] + impl SignalStore for GatedSenderKeyLookup { + async fn get_sender_key(&self, _: &str) -> StoreResult>> { + // Only the first round is gated. A reader whose install loses the + // epoch check reads again, and that retry must not wait on a + // rendezvous the test has already passed through. + let hit = self.hits.fetch_add(1, Ordering::Relaxed); + // Sampled before parking: these bytes belong to the moment the read + // started, so a write landing while we are gated is invisible here + // and visible to the next call, as a real backend behaves. + let sampled = self + .payload + .lock() + .unwrap_or_else(|p| p.into_inner()) + .clone(); + if hit < self.gated_reads { + self.arrived.wait().await; + self.release.wait().await; + } + Ok(sampled) + } + + async fn put_identity(&self, _: &str, _: [u8; 32]) -> StoreResult<()> { + unreachable!() + } + async fn load_identity(&self, _: &str) -> StoreResult> { + unreachable!() + } + async fn delete_identity(&self, _: &str) -> StoreResult<()> { + unreachable!() + } + async fn get_session(&self, _: &str) -> StoreResult> { + unreachable!() + } + async fn put_session(&self, _: &str, _: &[u8]) -> StoreResult<()> { + unreachable!() + } + async fn delete_session(&self, _: &str) -> StoreResult<()> { + unreachable!() + } + async fn store_prekey(&self, _: u32, _: &[u8], _: bool) -> StoreResult<()> { + unreachable!() + } + async fn load_prekey(&self, _: u32) -> StoreResult> { + unreachable!() + } + async fn mark_prekeys_uploaded(&self, _: &[u32]) -> StoreResult<()> { + unreachable!() + } + async fn remove_prekey(&self, _: u32) -> StoreResult<()> { + unreachable!() + } + async fn get_max_prekey_id(&self) -> StoreResult { + unreachable!() + } + async fn store_signed_prekey(&self, _: u32, _: &[u8]) -> StoreResult<()> { + unreachable!() + } + async fn load_signed_prekey(&self, _: u32) -> StoreResult>> { + unreachable!() + } + async fn load_all_signed_prekeys(&self) -> StoreResult)>> { + unreachable!() + } + async fn remove_signed_prekey(&self, _: u32) -> StoreResult<()> { + unreachable!() + } + /// A flush writes through here, so the payload a later read samples is + /// the one the flush persisted, as a real backend behaves. + async fn put_sender_key(&self, _: &str, record: &[u8]) -> StoreResult<()> { + self.set_payload(Some(record.to_vec())); + Ok(()) + } + async fn delete_sender_key(&self, _: &str) -> StoreResult<()> { + self.set_payload(None); + Ok(()) + } + } + + fn sender_key_record_with_chain(chain_id: u32) -> SenderKeyRecord { + use crate::libsignal::protocol::KeyPair; + let mut rng = rand::make_rng::(); + let kp = KeyPair::generate(&mut rng); + let mut record = SenderKeyRecord::new_empty(); + record + .add_sender_key_state( + 3, + chain_id, + 0, + &[7u8; 32], + kp.public_key, + Some(kp.private_key), + ) + .expect("valid sender key state"); + record + } + + fn chain_id_of(record: &SenderKeyRecord) -> u32 { + record + .sender_key_state() + .expect("record must carry a state") + .chain_id() + } + + #[tokio::test] + async fn cold_sender_key_miss_loads_from_the_backend() { + let cache = SignalStoreCache::new(); + let backend = crate::store::in_memory::InMemoryBackend::new(); + let name = SenderKeyName::from_parts("19995550001@g.us", "19995550002@s.whatsapp.net:0"); + let stored = sender_key_record_with_chain(7); + backend + .put_sender_key( + name.cache_key(), + &stored.serialize().expect("serialize record"), + ) + .await + .expect("seed backend"); + + let loaded = cache + .get_sender_key(&name, &backend) + .await + .expect("cold load") + .expect("record present"); + assert_eq!(chain_id_of(&loaded), 7); + + // Second read is a cache hit and must agree with the first. + let warm = cache + .get_sender_key(&name, &backend) + .await + .expect("warm load") + .expect("record present"); + assert_eq!(chain_id_of(&warm), 7); + } + + #[tokio::test] + async fn concurrent_cold_sender_key_readers_share_one_cached_value() { + let cache = Arc::new(SignalStoreCache::new()); + let stored = sender_key_record_with_chain(11); + let backend = Arc::new(GatedSenderKeyLookup::new( + 2, + Some(stored.serialize().expect("serialize record")), + )); + let name = Arc::new(SenderKeyName::from_parts( + "19995550003@g.us", + "19995550004@s.whatsapp.net:0", + )); + + let readers: Vec<_> = (0..2) + .map(|_| { + let (cache, backend, name) = (cache.clone(), backend.clone(), name.clone()); + tokio::spawn(async move { cache.get_sender_key(&name, &*backend).await }) + }) + .collect(); + + // Both readers must reach the backend: with the lock held across the + // round-trip the second would still be queued and this would not + // rendezvous. + backend.arrived.wait().await; + backend.release.wait().await; + + for reader in readers { + let record = reader + .await + .expect("reader task") + .expect("cold load") + .expect("record present"); + assert_eq!(chain_id_of(&record), 11); + } + assert_eq!(backend.hits(), 2, "both readers should consult the backend"); + + let cached = cache + .get_sender_key(&name, &*backend) + .await + .expect("warm load") + .expect("record present"); + assert_eq!(chain_id_of(&cached), 11, "cache must hold a single value"); + } + + #[tokio::test] + async fn a_late_sender_key_reader_does_not_overwrite_a_concurrent_put() { + let cache = Arc::new(SignalStoreCache::new()); + let backend = Arc::new(GatedSenderKeyLookup::new( + 1, + Some( + sender_key_record_with_chain(1) + .serialize() + .expect("serialize record"), + ), + )); + let name = Arc::new(SenderKeyName::from_parts( + "19995550005@g.us", + "19995550006@s.whatsapp.net:0", + )); + + let reader = tokio::spawn({ + let (cache, backend, name) = (cache.clone(), backend.clone(), name.clone()); + async move { cache.get_sender_key(&name, &*backend).await } + }); + + backend.arrived.wait().await; + // The reader is parked inside the backend: a writer must be able to + // take the cache mutex right now, and must win the re-check. + cache + .put_sender_key(&name, sender_key_record_with_chain(2)) + .await; + backend.release.wait().await; + + let observed = reader + .await + .expect("reader task") + .expect("cold load") + .expect("record present"); + assert_eq!(chain_id_of(&observed), 2, "reader must yield to the writer"); + + let cached = cache + .get_sender_key(&name, &*backend) + .await + .expect("warm load") + .expect("record present"); + assert_eq!(chain_id_of(&cached), 2, "stale backend read must not land"); + } + + #[tokio::test] + async fn a_late_sender_key_reader_does_not_resurrect_a_concurrent_delete() { + let cache = Arc::new(SignalStoreCache::new()); + let backend = Arc::new(GatedSenderKeyLookup::new( + 1, + Some( + sender_key_record_with_chain(3) + .serialize() + .expect("serialize record"), + ), + )); + let name = Arc::new(SenderKeyName::from_parts( + "19995550007@g.us", + "19995550008@s.whatsapp.net:0", + )); + + let reader = tokio::spawn({ + let (cache, backend, name) = (cache.clone(), backend.clone(), name.clone()); + async move { cache.get_sender_key(&name, &*backend).await } + }); + + backend.arrived.wait().await; + cache.delete_sender_key(name.cache_key()).await; + backend.release.wait().await; + + assert!( + reader + .await + .expect("reader task") + .expect("cold load") + .is_none(), + "reader must observe the delete, not the row it read before it" + ); + assert!( + cache + .get_sender_key(&name, &*backend) + .await + .expect("warm load") + .is_none(), + "the tombstone must survive the late reader" + ); + } + + /// The dangerous shape behind a re-check that only looks for an entry: a + /// newer record is written, flushed and then dropped by a clean removal + /// while a cold read is in flight. The slot is absent again at re-check, + /// but the reader's bytes predate the write, and a clean removal keeps the + /// incarnation, so installing them would be trusted as an exact reload and + /// could resume an already-published iteration. + #[tokio::test] + async fn a_write_dropped_by_eviction_is_not_replaced_by_the_stale_read() { + let cache = Arc::new(SignalStoreCache::new()); + let backend = Arc::new(GatedSenderKeyLookup::new( + 1, + Some( + sender_key_record_with_chain(1) + .serialize() + .expect("serialize record"), + ), + )); + let name = Arc::new(SenderKeyName::from_parts( + "19995550009@g.us", + "19995550010@s.whatsapp.net:0", + )); + + let reader = tokio::spawn({ + let (cache, backend, name) = (cache.clone(), backend.clone(), name.clone()); + async move { cache.get_sender_key(&name, &*backend).await } + }); + + backend.arrived.wait().await; + // A newer chain lands, is flushed (so the backend now holds it), then + // leaves the cache entirely as a clean entry. + cache + .put_sender_key(&name, sender_key_record_with_chain(2)) + .await; + backend.set_payload(Some( + sender_key_record_with_chain(2) + .serialize() + .expect("serialize record"), + )); + cache.drop_clean_sender_key_for_test(name.cache_key()).await; + backend.release.wait().await; + + let observed = reader + .await + .expect("reader task") + .expect("cold load") + .expect("record present"); + assert_eq!( + chain_id_of(&observed), + 2, + "the pre-write bytes must not survive a removal that happened after them" + ); + let cached = cache + .get_sender_key(&name, &*backend) + .await + .expect("warm load") + .expect("record present"); + assert_eq!(chain_id_of(&cached), 2, "a stale record must not be cached"); + } + + /// A cold read that is cancelled mid-backend must leave no bookkeeping + /// behind: the removal window is fixed-size and reader-agnostic, so a + /// dropped future cannot strand state that would pin its chain to the slow + /// path or grow the cache. + #[tokio::test] + async fn a_cancelled_cold_read_leaves_no_bookkeeping() { + let cache = Arc::new(SignalStoreCache::new()); + let backend = Arc::new(GatedSenderKeyLookup::with_rounds( + 1, + 1, + Some( + sender_key_record_with_chain(4) + .serialize() + .expect("serialize record"), + ), + )); + let name = Arc::new(SenderKeyName::from_parts( + "19995550017@g.us", + "19995550018@s.whatsapp.net:0", + )); + + let reader = tokio::spawn({ + let (cache, backend, name) = (cache.clone(), backend.clone(), name.clone()); + async move { cache.get_sender_key(&name, &*backend).await } + }); + backend.arrived.wait().await; + // Drop the future while it is parked inside the backend. + reader.abort(); + let _ = reader.await; + + // A fresh read of the same chain must take the unlocked path and + // install on its first attempt, exactly as if nothing had happened. + let hits_before = backend.hits(); + let observed = cache + .get_sender_key(&name, &*backend) + .await + .expect("cold load") + .expect("record present"); + assert_eq!(chain_id_of(&observed), 4); + assert_eq!( + backend.hits() - hits_before, + 1, + "a cancelled read must not push later reads onto the retry path" + ); + } + + /// The keyed removal path is what the other race tests drive. `clear` and + /// `retain` cannot name the keys they drop, so they take a separate branch + /// that concedes every in-flight reader — and that is the branch where a + /// missing bump would let pre-write bytes be adopted as an exact reload. + /// Drives it through the real flush-then-`clear_after_flush` sequence. + #[tokio::test] + async fn a_write_dropped_by_clear_after_flush_is_not_replaced_by_the_stale_read() { + let cache = Arc::new(SignalStoreCache::new()); + let backend = Arc::new(GatedSenderKeyLookup::new( + 1, + Some( + sender_key_record_with_chain(1) + .serialize() + .expect("serialize record"), + ), + )); + let name = Arc::new(SenderKeyName::from_parts( + "19995550019@g.us", + "19995550020@s.whatsapp.net:0", + )); + + let reader = tokio::spawn({ + let (cache, backend, name) = (cache.clone(), backend.clone(), name.clone()); + async move { cache.get_sender_key(&name, &*backend).await } + }); + + backend.arrived.wait().await; + // A newer chain lands, is flushed (which writes it through to the + // backend), and teardown then drops it from the cache — the opaque + // removal path, not the keyed one. + cache + .put_sender_key(&name, sender_key_record_with_chain(2)) + .await; + cache.flush(&*backend).await.expect("flush"); + cache.clear_after_flush().await; + backend.release.wait().await; + + let observed = reader + .await + .expect("reader task") + .expect("cold load") + .expect("record present"); + assert_eq!( + chain_id_of(&observed), + 2, + "an unnamed removal must still reject bytes that predate the write" + ); + } + + /// The removal signal is per key: churn on other chains, which is the + /// normal state of a cache at its eviction watermark, must not cost an + /// unrelated cold reader its unlocked install. + #[tokio::test] + async fn churn_on_other_chains_does_not_force_a_reread() { + let cache = Arc::new(SignalStoreCache::new()); + let backend = Arc::new(GatedSenderKeyLookup::new( + 1, + Some( + sender_key_record_with_chain(5) + .serialize() + .expect("serialize record"), + ), + )); + let name = Arc::new(SenderKeyName::from_parts( + "19995550013@g.us", + "19995550014@s.whatsapp.net:0", + )); + let other = SenderKeyName::from_parts("19995550015@g.us", "19995550016@s.whatsapp.net:0"); + + let reader = tokio::spawn({ + let (cache, backend, name) = (cache.clone(), backend.clone(), name.clone()); + async move { cache.get_sender_key(&name, &*backend).await } + }); + + backend.arrived.wait().await; + // A whole write-and-drop cycle on a different chain. + cache + .put_sender_key(&other, sender_key_record_with_chain(9)) + .await; + cache + .drop_clean_sender_key_for_test(other.cache_key()) + .await; + backend.release.wait().await; + + let observed = reader + .await + .expect("reader task") + .expect("cold load") + .expect("record present"); + assert_eq!(chain_id_of(&observed), 5); + assert_eq!( + backend.hits(), + 1, + "an unrelated chain's removal must not invalidate this read" + ); + } + + /// Losing the epoch check on every unlocked attempt drops through to the + /// read taken under the lock, which cannot be raced. That path installs + /// without an epoch check precisely because nothing can intervene, so it + /// needs its own coverage rather than inheriting the loop's. + #[tokio::test] + async fn a_read_that_loses_every_race_falls_back_to_the_locked_path() { + let cache = Arc::new(SignalStoreCache::new()); + let backend = Arc::new(GatedSenderKeyLookup::with_rounds( + 1, + // Gate exactly the unlocked attempts; the locked fallback then runs + // ungated, as a real backend would. + SENDER_KEY_UNLOCKED_READ_ATTEMPTS, + Some( + sender_key_record_with_chain(1) + .serialize() + .expect("serialize record"), + ), + )); + let name = Arc::new(SenderKeyName::from_parts( + "19995550011@g.us", + "19995550012@s.whatsapp.net:0", + )); + + let reader = tokio::spawn({ + let (cache, backend, name) = (cache.clone(), backend.clone(), name.clone()); + async move { cache.get_sender_key(&name, &*backend).await } + }); + + // Invalidate this key on every attempt, so no unlocked install survives. + for chain in 2..=(SENDER_KEY_UNLOCKED_READ_ATTEMPTS as u32 + 1) { + backend.arrived.wait().await; + cache + .put_sender_key(&name, sender_key_record_with_chain(chain)) + .await; + backend.set_payload(Some( + sender_key_record_with_chain(chain) + .serialize() + .expect("serialize record"), + )); + cache.drop_clean_sender_key_for_test(name.cache_key()).await; + backend.release.wait().await; + } + + let observed = reader + .await + .expect("reader task") + .expect("cold load") + .expect("record present"); + let latest = SENDER_KEY_UNLOCKED_READ_ATTEMPTS as u32 + 1; + assert_eq!( + chain_id_of(&observed), + latest, + "the locked fallback must return the current record" + ); + assert!( + backend.hits() > SENDER_KEY_UNLOCKED_READ_ATTEMPTS, + "the locked fallback must have read the backend itself" + ); + } + + #[tokio::test] + async fn the_user_index_never_misses_state_across_the_mutation_paths() { + let cache = SignalStoreCache::new(); + let backend = crate::store::in_memory::InMemoryBackend::new(); + + // Every public path that can put an address into either scanned store + // must leave the user answerable without a backend probe. + let session_user = "19995551001"; + let session_addr = + ProtocolAddress::new(&format!("{session_user}@s.whatsapp.net"), 0.into()); + cache + .put_session(&session_addr, SessionRecord::new_fresh()) + .await; + assert!( + cache + .has_state_for_user(session_user, &backend) + .await + .unwrap() + ); + + let deleted_user = "19995551002"; + let deleted_addr = + ProtocolAddress::new(&format!("{deleted_user}@s.whatsapp.net"), 0.into()); + cache.delete_session(&deleted_addr).await; + assert!( + cache + .has_state_for_user(deleted_user, &backend) + .await + .unwrap() + ); + + let identity_user = "19995551003"; + let identity_addr = + ProtocolAddress::new(&format!("{identity_user}@s.whatsapp.net"), 0.into()); + cache.put_identity(&identity_addr, &[3u8; 32]).await; + assert!( + cache + .has_state_for_user(identity_user, &backend) + .await + .unwrap() + ); + + let probed_user = "19995551004"; + let probed_addr = ProtocolAddress::new(&format!("{probed_user}@s.whatsapp.net"), 0.into()); + cache.has_session(&probed_addr, &backend).await.unwrap(); + assert!( + cache + .has_state_for_user(probed_user, &backend) + .await + .unwrap() + ); + + let checked_out_user = "19995551005"; + let checked_out_addr = + ProtocolAddress::new(&format!("{checked_out_user}@s.whatsapp.net"), 0.into()); + let (_, checkout) = cache + .checkout_session(&checked_out_addr, &backend) + .await + .unwrap(); + assert!( + cache + .has_state_for_user(checked_out_user, &backend) + .await + .unwrap() + ); + cache.cancel_session_checkout(&checked_out_addr, checkout); + assert!( + cache + .has_state_for_user(checked_out_user, &backend) + .await + .unwrap() + ); + + // A user with no state anywhere still answers false. + assert!( + !cache + .has_state_for_user("19995559999", &backend) + .await + .unwrap() + ); + + // An addressed-device JID renders as `user:device@server.N`, so both + // `user` and `user:device` prefix-match it under the scan predicate. + // Only the first is an index key; the second must be conceded rather + // than denied, which is the one false negative available here. + let device_addr = ProtocolAddress::new("19995551006:5@c.us", 0.into()); + cache + .put_session(&device_addr, SessionRecord::new_fresh()) + .await; + assert!( + cache + .has_state_for_user("19995551006", &backend) + .await + .unwrap() + ); + + let with_device = "19995551006:5"; + assert!( + protocol_address_matches_user(device_addr.as_str(), with_device), + "the scan predicate matches this user, so the index must not deny it" + ); + assert!( + cache + .has_state_for_user(with_device, &backend) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn the_user_index_survives_eviction_and_defers_to_the_backend_when_cold() { + let cache = SignalStoreCache::with_max_entries(4); + let backend = crate::store::in_memory::InMemoryBackend::new(); + + // A dirty entry is never evicted, so this user must stay answerable + // from the index no matter how much churn follows. + let pinned_user = "19995552000"; + let pinned_addr = ProtocolAddress::new(&format!("{pinned_user}@s.whatsapp.net"), 0.into()); + cache + .put_session(&pinned_addr, SessionRecord::new_fresh()) + .await; + + // A user whose only entry is clean, so eviction can drop it. Give it + // durable state so the probe has something truthful to report. + let evicted_user = "19995552001"; + let evicted_addr = + ProtocolAddress::new(&format!("{evicted_user}@s.whatsapp.net"), 0.into()); + backend + .put_identity(evicted_addr.as_str(), [9u8; 32]) + .await + .unwrap(); + cache.has_session(&evicted_addr, &backend).await.unwrap(); + + // Churn well past the high watermark so eviction and index compaction + // both run repeatedly. + for i in 100..400 { + let addr = ProtocolAddress::new(&format!("1999555{i:04}@s.whatsapp.net"), 0.into()); + cache.has_session(&addr, &backend).await.unwrap(); + } + + assert!( + cache + .has_state_for_user(pinned_user, &backend) + .await + .unwrap(), + "a still-cached user must stay answerable from the index" + ); + assert!( + cache + .has_state_for_user(evicted_user, &backend) + .await + .unwrap(), + "an evicted user's durable state must still be found via the probe" + ); + assert!( + !cache + .has_state_for_user("19995559998", &backend) + .await + .unwrap(), + "a user with no state anywhere must answer false" + ); + + // A lossy reset drops the index with the cache; the probe is then the + // only source of truth, and must still find durable state. + cache.clear().await; + assert!( + cache + .has_state_for_user(evicted_user, &backend) + .await + .unwrap(), + "durable state must be found through the probe with a cold index" + ); + } + async fn wait_for_lock_waiter(lock: &Arc>, baseline: usize) { for _ in 0..10_000 { if Arc::strong_count(lock) > baseline {