diff --git a/src/cache_config.rs b/src/cache_config.rs index e0cb3fc38..e01aa4149 100644 --- a/src/cache_config.rs +++ b/src/cache_config.rs @@ -204,10 +204,9 @@ pub struct CacheConfig { pub session_locks_capacity: u64, /// Per-chat lane capacity (combined lock + queue). Default: 5000. pub chat_lanes_capacity: u64, - /// Per-group cold sender-key distribution lock capacity. Default: 512 - /// (far above any realistic number of groups distributing at once; an - /// evicted live lock only lets one extra send repeat that group's - /// fan-out, the ratchet stays correct under the chain lock). + /// Per-group cold sender-key distribution lock capacity. Default: 512. + /// Soft cap: a live lane is never evicted, so the map may briefly exceed + /// this under concurrent fan-out instead of breaking tracker ordering. pub group_distribution_locks_capacity: u64, /// Per-chat resend rate-limiter capacity: one token-bucket entry per group /// recently driving retry resends. Keep above the count of concurrently diff --git a/src/client.rs b/src/client.rs index f2c0044c3..25216e009 100644 --- a/src/client.rs +++ b/src/client.rs @@ -225,6 +225,11 @@ pub struct MemoryReport { // -- Capacity-only caches (coordination, counts only) -- pub session_locks: u64, pub chat_lanes: u64, + pub group_distribution_locks: u64, + /// Cumulative capacity evictions; poll successive reports to derive a rate. + pub group_distribution_lock_evictions: u64, + /// Cumulative attempts that kept a live lane and temporarily exceeded capacity. + pub group_distribution_lock_eviction_blocks: u64, pub resend_rate_limiter_chats: u64, // -- Unbounded collections -- pub response_waiters: usize, @@ -295,6 +300,13 @@ impl std::fmt::Display for MemoryReport { writeln!(f, "--- Capacity-only caches ---")?; writeln!(f, " session_locks: {}", self.session_locks)?; writeln!(f, " chat_lanes: {}", self.chat_lanes)?; + writeln!( + f, + " group_dist_locks: {} (evicted: {}, blocked: {})", + self.group_distribution_locks, + self.group_distribution_lock_evictions, + self.group_distribution_lock_eviction_blocks + )?; writeln!( f, " resend_rl_chats: {}", diff --git a/src/client/accessors.rs b/src/client/accessors.rs index 9485b4b2a..747245ad3 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -156,6 +156,7 @@ impl Client { .group_devices_memo .memory_stats(|k, v| k.heap_bytes() + v.heap_bytes()) .await; + let group_distribution_locks = self.group_distribution_locks.capacity_stats().await; // Each count read into a local so no two guards are ever held at once. let response_waiters = self.response_waiters_guard().len(); @@ -178,6 +179,9 @@ impl Client { pdo_requested: self.pdo_requested.entry_count(), session_locks: self.session_locks.entry_count(), chat_lanes: self.chat_lanes.entry_count(), + group_distribution_locks: group_distribution_locks.entries, + group_distribution_lock_evictions: group_distribution_locks.evictions, + group_distribution_lock_eviction_blocks: group_distribution_locks.eviction_blocks, resend_rate_limiter_chats: self.resend_rate_limiter.entry_count(), response_waiters, node_waiters: self.node_waiter_count.load(Ordering::Relaxed), diff --git a/src/client/device_registry.rs b/src/client/device_registry.rs index df04421c2..8cf648491 100644 --- a/src/client/device_registry.rs +++ b/src/client/device_registry.rs @@ -1032,7 +1032,7 @@ impl Client { mod tests { use super::*; use crate::lid_pn_cache::LearningSource; - use crate::test_utils::create_test_client_with_failing_http; + use crate::test_utils::{create_test_client_with_failing_http, wait_for_lock_waiter}; use std::sync::Arc; async fn create_test_client() -> Arc { @@ -2700,7 +2700,8 @@ mod tests { use wacore::types::jid::JidExt; let client = create_test_client().await; - let group = "120363000000000001@g.us"; + let group: Jid = "120363000000000001@g.us".parse().unwrap(); + let group_id = group.to_string(); let own_lid = Jid::from_str("193832511623409:13@lid").unwrap(); client .persistence_manager @@ -2709,7 +2710,7 @@ mod tests { ))) .await; - let sk_name = SenderKeyName::from_parts(group, own_lid.to_protocol_address().as_str()); + let sk_name = SenderKeyName::from_parts(&group_id, own_lid.to_protocol_address().as_str()); client .signal_cache .put_sender_key(&sk_name, SenderKeyRecord::new_empty()) @@ -2718,7 +2719,7 @@ mod tests { client .persistence_manager .set_sender_key_status( - group, + &group_id, &[ ("271060335329480:0@lid", true), ("77610646245392:0@lid", true), @@ -2728,7 +2729,7 @@ mod tests { .unwrap(); client - .rotate_sender_key_on_participant_remove(group, &["271060335329480"]) + .rotate_sender_key_on_participant_remove(&group, &["271060335329480"]) .await; let device_snapshot = client.persistence_manager.get_device_snapshot(); @@ -2744,7 +2745,7 @@ mod tests { let rows = client .persistence_manager - .get_sender_key_devices(group) + .get_sender_key_devices(&group_id) .await .unwrap(); assert!(rows.is_empty(), "sender_key_devices must be cleared"); @@ -2760,7 +2761,8 @@ mod tests { use wacore::types::jid::JidExt; let client = create_test_client().await; - let group = "120363000000000001@g.us"; + let group: Jid = "120363000000000001@g.us".parse().unwrap(); + let group_id = group.to_string(); let own_lid = Jid::from_str("193832511623409:13@lid").unwrap(); client .persistence_manager @@ -2769,7 +2771,7 @@ mod tests { ))) .await; - let sk_name = SenderKeyName::from_parts(group, own_lid.to_protocol_address().as_str()); + let sk_name = SenderKeyName::from_parts(&group_id, own_lid.to_protocol_address().as_str()); client .signal_cache .put_sender_key(&sk_name, SenderKeyRecord::new_empty()) @@ -2777,12 +2779,12 @@ mod tests { client .persistence_manager - .set_sender_key_status(group, &[("271060335329480:0@lid", false)]) + .set_sender_key_status(&group_id, &[("271060335329480:0@lid", false)]) .await .unwrap(); client - .rotate_sender_key_on_participant_remove(group, &["271060335329480"]) + .rotate_sender_key_on_participant_remove(&group, &["271060335329480"]) .await; let device_snapshot = client.persistence_manager.get_device_snapshot(); @@ -2796,4 +2798,186 @@ mod tests { "sender key must survive when removed had no key" ); } + + #[tokio::test] + async fn rotation_waits_for_in_flight_sender_key_advance() { + use wacore::libsignal::protocol::{ + KeyPair, SENDERKEY_MESSAGE_CURRENT_VERSION, SenderKeyRecord, group_encrypt, + }; + use wacore::libsignal::store::sender_key_name::SenderKeyName; + use wacore::types::jid::JidExt; + + let client = create_test_client().await; + let group: Jid = "120363000000000003@g.us".parse().unwrap(); + let group_id = group.to_string(); + let own_lid: Jid = "193832511623410:13@lid".parse().unwrap(); + client + .persistence_manager + .process_command(crate::store::commands::DeviceCommand::SetLid(Some( + own_lid.clone(), + ))) + .await; + + let name = SenderKeyName::from_parts(&group_id, own_lid.to_protocol_address().as_str()); + let mut rng = rand::make_rng::(); + let key_pair = KeyPair::generate(&mut rng); + let mut record = SenderKeyRecord::new_empty(); + record + .add_sender_key_state( + SENDERKEY_MESSAGE_CURRENT_VERSION, + 7, + 0, + &[9; 32], + key_pair.public_key, + Some(key_pair.private_key), + ) + .unwrap(); + client.signal_cache.put_sender_key(&name, record).await; + + let chain_lock = client.signal_cache.sender_key_lock(&name).await; + let held = chain_lock.lock().await; + let lock_refs = Arc::strong_count(&chain_lock); + let started = Arc::new(tokio::sync::Barrier::new(2)); + let rotation = tokio::spawn({ + let client = client.clone(); + let group = group.clone(); + let started = started.clone(); + async move { + started.wait().await; + client.force_rotate_own_sender_key(&group).await; + } + }); + + started.wait().await; + wait_for_lock_waiter(&chain_lock, lock_refs).await; + let snapshot = client.persistence_manager.get_device_snapshot(); + assert!( + client + .signal_cache + .get_sender_key(&name, &*snapshot.backend) + .await + .unwrap() + .is_some(), + "rotation must wait for the in-flight advance" + ); + + let mut sender_key_store = client.sender_key_adapter().await; + group_encrypt( + &mut sender_key_store, + &name, + b"in-flight ciphertext", + &mut rng, + ) + .await + .expect("advance under the held chain lock"); + drop(held); + + tokio::time::timeout(std::time::Duration::from_secs(5), rotation) + .await + .expect("rotation must resume") + .expect("rotation task"); + assert!( + client + .signal_cache + .get_sender_key(&name, &*snapshot.backend) + .await + .unwrap() + .is_none(), + "rotation must retire the state written by the in-flight advance" + ); + } + + #[tokio::test] + async fn participant_rotation_audit_waits_for_group_distribution_guard() { + use wacore::libsignal::protocol::SenderKeyRecord; + use wacore::libsignal::store::sender_key_name::SenderKeyName; + use wacore::types::jid::JidExt; + + let client = create_test_client().await; + let group: Jid = "120363000000000004@g.us".parse().unwrap(); + let group_id = group.to_string(); + let own_lid: Jid = "193832511623411:13@lid".parse().unwrap(); + client + .persistence_manager + .process_command(crate::store::commands::DeviceCommand::SetLid(Some( + own_lid.clone(), + ))) + .await; + + let name = SenderKeyName::from_parts(&group_id, own_lid.to_protocol_address().as_str()); + client + .signal_cache + .put_sender_key(&name, SenderKeyRecord::new_empty()) + .await; + client + .persistence_manager + .set_sender_key_status(&group_id, &[("271060335329481:0@lid", true)]) + .await + .unwrap(); + + let held = client.group_distribution_lock(&group).await; + let lock = client + .group_distribution_locks + .get(&group) + .await + .expect("cached distribution lock"); + let lock_refs = Arc::strong_count(&lock); + let started = Arc::new(tokio::sync::Barrier::new(2)); + let rotation = tokio::spawn({ + let client = client.clone(); + let group = group.clone(); + let started = started.clone(); + async move { + started.wait().await; + client + .rotate_sender_key_on_participant_remove(&group, &["271060335329481"]) + .await; + } + }); + + started.wait().await; + wait_for_lock_waiter(&lock, lock_refs).await; + let snapshot = client.persistence_manager.get_device_snapshot(); + assert!( + client + .signal_cache + .get_sender_key(&name, &*snapshot.backend) + .await + .unwrap() + .is_some(), + "rotation must not delete before the active distribution ends" + ); + assert_eq!( + client + .persistence_manager + .get_sender_key_devices(&group_id) + .await + .unwrap() + .len(), + 1, + "rotation must not clear tracking before it owns the distribution lane" + ); + + drop(held); + tokio::time::timeout(std::time::Duration::from_secs(5), rotation) + .await + .expect("rotation must resume") + .expect("rotation task"); + assert!( + client + .signal_cache + .get_sender_key(&name, &*snapshot.backend) + .await + .unwrap() + .is_none() + ); + assert!( + client + .persistence_manager + .get_sender_key_devices(&group_id) + .await + .unwrap() + .is_empty() + ); + } } diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index f958d03f2..b515062a5 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -267,12 +267,10 @@ impl Client { group_devices_memo: Cache::builder() .max_capacity(GROUP_DEVICES_MEMO_CAPACITY) .build(), - // Evicting a lock whose guard is still held only lets one extra - // send re-run that group's fan-out (the pre-single-flight - // behavior); the sender-key chain lock still guarantees ratchet - // correctness. + // A live lane also protects recipient-tracker reset/update ordering. group_distribution_locks: Cache::builder() .max_capacity(cache_config.group_distribution_locks_capacity.max(1)) + .evict_guard(|m| Arc::strong_count(m) <= 1) .build(), skdm_warm_memo: Cache::builder() .max_capacity(GROUP_DEVICES_MEMO_CAPACITY) diff --git a/src/client/sender_keys.rs b/src/client/sender_keys.rs index e48cb6e59..6488e74cf 100644 --- a/src/client/sender_keys.rs +++ b/src/client/sender_keys.rs @@ -83,6 +83,49 @@ impl Client { Ok(()) } + /// Redistribution must not inherit delivery marks from an earlier pass. + pub(crate) async fn reset_sender_key_device_tracking(&self, group_jid: &str) -> Result<()> { + if let Err(clear_error) = self + .persistence_manager + .clear_sender_key_devices(group_jid) + .await + { + // Cold marks preserve the reset when row deletion is unavailable. + let rows = self + .persistence_manager + .get_sender_key_devices(group_jid) + .await + .map_err(|fallback_error| { + anyhow::anyhow!( + "sender-key tracker clear failed ({clear_error}); \ + fallback read failed: {fallback_error}" + ) + })?; + if !rows.is_empty() { + let entries: Vec<(&str, bool)> = rows + .iter() + .map(|(device_jid, _)| (device_jid.as_str(), false)) + .collect(); + self.persistence_manager + .set_sender_key_status(group_jid, &entries) + .await + .map_err(|fallback_error| { + anyhow::anyhow!( + "sender-key tracker clear failed ({clear_error}); \ + cold-mark fallback failed: {fallback_error}" + ) + })?; + } + log::warn!( + "reset_sender_key_device_tracking: clear failed for {group_jid}; \ + marked {} existing rows cold: {clear_error}", + rows.len() + ); + } + self.sender_key_device_cache.invalidate(group_jid).await; + Ok(()) + } + /// Forward-secrecy rotation when participants leave a group. Mirrors WA /// Web's `removeParticipantInfo` (`GroupParticipantHelpers.js`): if any /// removed user had `has_key=true`, delete the bot's own sender key for @@ -92,18 +135,20 @@ impl Client { #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.session.rotate_sender_key_on_remove", level = "debug", skip_all, fields(removed = removed_user_ids.len())))] pub(crate) async fn rotate_sender_key_on_participant_remove( &self, - group_jid: &str, + group_jid: &Jid, removed_user_ids: &[&str], ) { if removed_user_ids.is_empty() { return; } + let group_id = group_jid.to_string(); + let distribution_guard = self.group_distribution_lock(group_jid).await; // Read failure → rotate anyway. Better to pay the redistribute cost // than leave the sender key in place after a removal we couldn't audit. let (rows, read_failed) = match self .persistence_manager - .get_sender_key_devices(group_jid) + .get_sender_key_devices(&group_id) .await { Ok(r) => (r, false), @@ -127,7 +172,10 @@ impl Client { return; } - self.force_rotate_own_sender_key(group_jid).await; + self.rotate_own_sender_key_state(&group_id).await; + drop(distribution_guard); + self.flush_signal_cache_batch_safe_logged("rotate_on_participant_remove", None) + .await; } /// Unconditional forward-secrecy rotation: delete the bot's own sender key @@ -143,28 +191,32 @@ impl Client { skip_all ) )] - pub(crate) async fn force_rotate_own_sender_key(&self, group_jid: &str) { + pub(crate) async fn force_rotate_own_sender_key(&self, group_jid: &Jid) { + let group_id = group_jid.to_string(); + let distribution_guard = self.group_distribution_lock(group_jid).await; + self.rotate_own_sender_key_state(&group_id).await; + drop(distribution_guard); + + self.flush_signal_cache_batch_safe_logged("force_rotate_own_sender_key", None) + .await; + } + + /// A send must not publish a new distribution between the audit and reset. + async fn rotate_own_sender_key_state(&self, group_id: &str) { use wacore::libsignal::store::sender_key_name::SenderKeyName; use wacore::types::jid::JidExt; let snapshot = self.persistence_manager.get_device_snapshot(); for own_jid in snapshot.lid.iter().chain(snapshot.pn.iter()) { let sk_name = - SenderKeyName::from_parts(group_jid, own_jid.to_protocol_address().as_str()); + SenderKeyName::from_parts(group_id, own_jid.to_protocol_address().as_str()); self.signal_cache .delete_sender_key(sk_name.cache_key()) .await; } - self.flush_signal_cache_batch_safe_logged("force_rotate_own_sender_key", None) - .await; - if let Err(e) = self - .persistence_manager - .clear_sender_key_devices(group_jid) - .await - { - log::warn!("force_rotate_own_sender_key: clear DB failed: {e}"); + if let Err(e) = self.reset_sender_key_device_tracking(group_id).await { + log::warn!("rotate_own_sender_key_state: reset failed for {group_id}: {e}"); } - self.sender_key_device_cache.invalidate(group_jid).await; } /// Take a sent message for retry handling. Checks L1 cache first (if enabled), @@ -535,6 +587,80 @@ mod tests { ); } + #[tokio::test] + async fn clear_failure_uses_durable_cold_mark_fallback() { + let client = create_test_client().await; + let group = "120363000000000005@g.us"; + let device = "111000111000112:0@lid"; + client + .persistence_manager + .set_sender_key_status(group, &[(device, true)]) + .await + .unwrap(); + + let rows = client + .persistence_manager + .get_sender_key_devices(group) + .await + .unwrap(); + let cached = client + .sender_key_device_cache + .get_or_init(group, async { + Arc::new(SenderKeyDeviceMap::from_db_rows(&rows)) + }) + .await; + + client + .persistence_manager + .fail_sender_key_device_clears_for_tests(true); + client + .persistence_manager + .fail_sender_key_device_status_writes_for_tests(true); + assert!( + client + .reset_sender_key_device_tracking(group) + .await + .is_err() + ); + let after_failure = client + .sender_key_device_cache + .get_or_init(group, async { panic!("failed clear must not invalidate") }) + .await; + assert!(Arc::ptr_eq(&cached, &after_failure)); + assert_eq!( + client + .persistence_manager + .get_sender_key_devices(group) + .await + .unwrap(), + vec![(device.to_string(), true)] + ); + + client + .persistence_manager + .fail_sender_key_device_status_writes_for_tests(false); + client + .reset_sender_key_device_tracking(group) + .await + .unwrap(); + let rows = client + .persistence_manager + .get_sender_key_devices(group) + .await + .unwrap(); + let after_fallback = client + .sender_key_device_cache + .get_or_init(group, async { + Arc::new(SenderKeyDeviceMap::from_db_rows(&rows)) + }) + .await; + assert!(!Arc::ptr_eq(&cached, &after_fallback)); + assert_eq!( + after_fallback.device_has_key("111000111000112", 0), + Some(false) + ); + } + // When every named device is our own, nothing is kept: no DB write, no flip, // and (crucially) no generation bump that would churn the warm memo. #[tokio::test] diff --git a/src/client/tests.rs b/src/client/tests.rs index 83f2d80ad..b59319476 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -2786,6 +2786,44 @@ async fn test_custom_cache_config_is_respected() { assert!(!client.is_logged_in()); } +#[tokio::test] +async fn held_group_distribution_lane_survives_capacity_pressure() { + let config = crate::cache_config::CacheConfig { + group_distribution_locks_capacity: 1, + ..Default::default() + }; + let client = crate::test_utils::create_test_client_with_config( + "group_distribution_eviction", + Arc::new(MockHttpClient), + config, + ) + .await; + + let first: Jid = "120363000000000011@g.us".parse().unwrap(); + let second: Jid = "120363000000000012@g.us".parse().unwrap(); + let third: Jid = "120363000000000013@g.us".parse().unwrap(); + let held = client.group_distribution_lock(&first).await; + + drop(client.group_distribution_lock(&second).await); + drop(client.group_distribution_lock(&third).await); + + let first_again = client + .group_distribution_locks + .get(&first) + .await + .expect("held lane must remain cached"); + assert!( + first_again.try_lock().is_none(), + "capacity pressure must not mint a second live lane" + ); + let report = client.memory_report().await; + assert_eq!(report.group_distribution_locks, 2); + assert_eq!(report.group_distribution_lock_evictions, 1); + assert_eq!(report.group_distribution_lock_eviction_blocks, 2); + drop(held); + assert!(first_again.try_lock().is_some()); +} + /// Proves that `is_connected()` no longer gives false negatives under mutex /// contention. Before the fix, `try_lock()` would fail when another task held /// the noise_socket mutex, causing `is_connected()` to return `false` even @@ -3383,6 +3421,9 @@ async fn memory_report_on_fresh_client() { let report = client.memory_report().await; assert_eq!(report.recent_messages.entries, 0); assert_eq!(report.recent_messages.bytes, 0); + assert_eq!(report.group_distribution_locks, 0); + assert_eq!(report.group_distribution_lock_evictions, 0); + assert_eq!(report.group_distribution_lock_eviction_blocks, 0); assert_eq!(report.signal_sessions.entries, 0); assert_eq!(report.response_waiters, 0); diff --git a/src/features/groups.rs b/src/features/groups.rs index 425137513..9ef8ef4c4 100644 --- a/src/features/groups.rs +++ b/src/features/groups.rs @@ -586,7 +586,7 @@ impl<'a> Groups<'a> { self.client.invalidate_persisted_group_metadata(jid).await; } self.client - .rotate_sender_key_on_participant_remove(&jid.to_string(), &accepted) + .rotate_sender_key_on_participant_remove(jid, &accepted) .await; } Ok(result) diff --git a/src/features/signal.rs b/src/features/signal.rs index cf6cd4354..10285d539 100644 --- a/src/features/signal.rs +++ b/src/features/signal.rs @@ -4,8 +4,8 @@ use thiserror::Error; use wacore::libsignal::protocol::{ - CiphertextMessage, PreKeySignalMessage, SignalMessage, SignalProtocolError, UsePQRatchet, - message_decrypt, message_encrypt, + CiphertextMessage, PreKeySignalMessage, SenderKeyStore, SignalMessage, SignalProtocolError, + UsePQRatchet, message_decrypt, message_encrypt, }; use wacore::message_processing::EncType; use wacore::messages::MessageUtils; @@ -224,8 +224,7 @@ impl<'a> Signal<'a> { /// Returns raw padded plaintext. Use [`MessageUtils::unpad_message_ref`] /// with the stanza's `v` attribute if WhatsApp message unpadding is needed. /// - /// Not safe to call concurrently with `encrypt_group_message` for the - /// same group — sender key state is not internally locked. + /// Concurrent mutations of the same sender-key chain are serialized. pub async fn decrypt_group_message( &self, group_jid: &Jid, @@ -236,6 +235,11 @@ impl<'a> Signal<'a> { make_sender_key_name(group_jid, &sender_jid.to_non_ad().to_protocol_address()); let mut adapter = self.client.signal_adapter().await; + let chain_lock = adapter + .sender_key_store + .sender_key_lock(&sender_key_name) + .await; + let _chain_guard = chain_lock.lock().await; let plaintext = wacore::libsignal::protocol::group_decrypt( ciphertext, @@ -244,6 +248,7 @@ impl<'a> Signal<'a> { ) .await?; + drop(_chain_guard); self.client.flush_signal_cache_batch_safe().await?; Ok(plaintext.to_vec()) diff --git a/src/handlers/notification/device.rs b/src/handlers/notification/device.rs index 11b0a486d..474c6269a 100644 --- a/src/handlers/notification/device.rs +++ b/src/handlers/notification/device.rs @@ -267,6 +267,8 @@ pub(crate) async fn handle_identity_change(client: &Arc, node: &NodeRef< client.signal_cache.delete_identity(cand).await; } + let status_jid = Jid::status_broadcast(); + let distribution_guard = client.group_distribution_lock(&status_jid).await; let status_group = "status@broadcast"; for own_jid in device_snapshot.pn.iter().chain(device_snapshot.lid.iter()) { let sk_name = @@ -276,6 +278,7 @@ pub(crate) async fn handle_identity_change(client: &Arc, node: &NodeRef< .delete_sender_key(sk_name.cache_key()) .await; } + drop(distribution_guard); client .flush_signal_cache_batch_safe_logged("identity change", None) diff --git a/src/handlers/notification/groups.rs b/src/handlers/notification/groups.rs index 541b497fa..b6d983ca4 100644 --- a/src/handlers/notification/groups.rs +++ b/src/handlers/notification/groups.rs @@ -195,10 +195,7 @@ pub(crate) async fn handle_group_notification(client: &Arc, node: Arc { @@ -229,7 +226,7 @@ pub(crate) async fn handle_group_notification(client: &Arc, node: Arc {} diff --git a/src/message/special.rs b/src/message/special.rs index b5ca760df..66355c81f 100644 --- a/src/message/special.rs +++ b/src/message/special.rs @@ -217,6 +217,8 @@ impl Client { // Only the sender-key store is needed here, so build it standalone instead of // the full five-store adapter. let mut sender_key_store = self.sender_key_adapter().await; + let chain_lock = sender_key_store.sender_key_lock(&sender_key_name).await; + let _chain_guard = chain_lock.lock().await; if let Err(e) = process_sender_key_distribution_message(&sender_key_name, &skdm, &mut sender_key_store) diff --git a/src/message/tests.rs b/src/message/tests.rs index 96a231877..a1a848745 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -3,7 +3,7 @@ use super::*; use crate::store::SqliteStore; use crate::store::persistence_manager::PersistenceManager; -use crate::test_utils::MockHttpClient; +use crate::test_utils::{MockHttpClient, wait_for_lock_waiter}; use crate::types::message::EditAttribute; use std::sync::Arc; use wacore_binary::builder::NodeBuilder; @@ -6685,6 +6685,144 @@ async fn group_skmsg_decrypts_under_sender_key_lock() { ); } +#[tokio::test] +async fn skdm_processing_waits_for_sender_key_lock() { + let client = crate::test_utils::create_test_client_with_name("skdm_chain_lock").await; + let mut alice = AlicePeer::new("146824178450542@lid").await; + let group: Jid = "120363408782575462@g.us".parse().expect("group"); + let skdm = alice.create_group_skdm(&group).await; + let bytes = skdm + .axolotl_sender_key_distribution_message + .expect("SKDM bytes"); + let sender = alice.jid.clone(); + let sender_key_name = make_sender_key_name(&group, &sender.to_non_ad().to_protocol_address()); + + let lock = client.signal_cache.sender_key_lock(&sender_key_name).await; + let held = lock.lock().await; + let lock_refs = Arc::strong_count(&lock); + let started = Arc::new(tokio::sync::Barrier::new(2)); + let task = tokio::spawn({ + let client = client.clone(); + let group = group.clone(); + let sender = sender.clone(); + let started = started.clone(); + async move { + started.wait().await; + client + .handle_sender_key_distribution_message(&group, &sender, &bytes) + .await; + } + }); + + started.wait().await; + wait_for_lock_waiter(&lock, lock_refs).await; + let snapshot = client.persistence_manager.get_device_snapshot(); + assert!( + client + .signal_cache + .get_sender_key(&sender_key_name, &*snapshot.backend) + .await + .unwrap() + .is_none(), + "SKDM must not mutate a checked-out chain" + ); + + drop(held); + tokio::time::timeout(std::time::Duration::from_secs(5), task) + .await + .expect("SKDM processing must resume") + .expect("SKDM task"); + assert!( + client + .signal_cache + .get_sender_key(&sender_key_name, &*snapshot.backend) + .await + .unwrap() + .is_some(), + "SKDM must install after the chain is released" + ); +} + +#[tokio::test] +async fn public_group_decrypt_waits_for_sender_key_lock() { + let client = crate::test_utils::create_test_client_with_name("public_group_decrypt_lock").await; + let mut alice = AlicePeer::new("146824178450543@lid").await; + let group: Jid = "120363408782575463@g.us".parse().expect("group"); + let skdm = alice.create_group_skdm(&group).await; + let skdm_bytes = skdm + .axolotl_sender_key_distribution_message + .expect("SKDM bytes"); + client + .handle_sender_key_distribution_message(&group, &alice.jid, &skdm_bytes) + .await; + + let plaintext = b"serialized group decrypt".to_vec(); + let ciphertext = alice.encrypt_group_message(&group, &plaintext).await; + let sender = alice.jid.clone(); + let sender_key_name = make_sender_key_name(&group, &sender.to_non_ad().to_protocol_address()); + let lock = client.signal_cache.sender_key_lock(&sender_key_name).await; + let held = lock.lock().await; + let lock_refs = Arc::strong_count(&lock); + let started = Arc::new(tokio::sync::Barrier::new(2)); + let task = tokio::spawn({ + let client = client.clone(); + let group = group.clone(); + let sender = sender.clone(); + let started = started.clone(); + async move { + started.wait().await; + client + .signal() + .decrypt_group_message(&group, &sender, &ciphertext) + .await + } + }); + + started.wait().await; + wait_for_lock_waiter(&lock, lock_refs).await; + let snapshot = client.persistence_manager.get_device_snapshot(); + let before = client + .signal_cache + .get_sender_key(&sender_key_name, &*snapshot.backend) + .await + .unwrap() + .expect("installed sender key"); + assert_eq!( + before + .sender_key_state() + .unwrap() + .sender_chain_key() + .unwrap() + .iteration(), + 0, + "decrypt must not advance while another mutation owns the chain" + ); + + drop(held); + let decrypted = tokio::time::timeout(std::time::Duration::from_secs(5), task) + .await + .expect("group decrypt must resume") + .expect("group decrypt task") + .expect("group decrypt"); + assert_eq!(decrypted, plaintext); + + let after = client + .signal_cache + .get_sender_key(&sender_key_name, &*snapshot.backend) + .await + .unwrap() + .expect("advanced sender key"); + assert_eq!( + after + .sender_key_state() + .unwrap() + .sender_chain_key() + .unwrap() + .iteration(), + 1 + ); +} + /// Bad path: a group skmsg whose sender key was never distributed hits /// NoSenderKeyState inside process_group_enc_batch (still under the lock) and /// takes the retry path — one undecryptable event, no user content. diff --git a/src/portable_cache.rs b/src/portable_cache.rs index ac6962837..eb44437d1 100644 --- a/src/portable_cache.rs +++ b/src/portable_cache.rs @@ -26,6 +26,13 @@ struct CacheEntry { seq: u64, } +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(crate) struct CapacityStats { + pub entries: u64, + pub evictions: u64, + pub eviction_blocks: u64, +} + /// Portable, runtime-agnostic in-process cache. /// /// - Max capacity with FIFO eviction @@ -53,6 +60,8 @@ struct CacheInner { order: BTreeMap, /// Next FIFO sequence to assign. next_seq: u64, + capacity_evictions: u64, + capacity_eviction_blocks: u64, } impl CacheInner @@ -64,6 +73,8 @@ where map: HashMap::new(), order: BTreeMap::new(), next_seq: 0, + capacity_evictions: 0, + capacity_eviction_blocks: 0, } } @@ -83,7 +94,9 @@ where // Unguarded caches keep the single-pass pop_first() fast path. None => match self.order.pop_first() { Some((_, oldest_key)) => { - self.map.remove(&oldest_key); + if self.map.remove(&oldest_key).is_some() { + self.capacity_evictions = self.capacity_evictions.saturating_add(1); + } } None => break, }, @@ -101,11 +114,17 @@ where } match victim_seq { Some(seq) => { - if let Some(oldest_key) = self.order.remove(&seq) { - self.map.remove(&oldest_key); + if let Some(oldest_key) = self.order.remove(&seq) + && self.map.remove(&oldest_key).is_some() + { + self.capacity_evictions = self.capacity_evictions.saturating_add(1); } } - None => break, + None => { + self.capacity_eviction_blocks = + self.capacity_eviction_blocks.saturating_add(1); + break; + } } } } @@ -372,6 +391,15 @@ where .unwrap_or(0) } + pub(crate) async fn capacity_stats(&self) -> CapacityStats { + let guard = self.inner.read().await; + CapacityStats { + entries: guard.map.len() as u64, + evictions: guard.capacity_evictions, + eviction_blocks: guard.capacity_eviction_blocks, + } + } + /// Reliable awaited snapshot of `(Arc, V)` pairs. Prefer this over /// [`iter`](Self::iter) in async contexts: `iter` is best-effort (a /// `try_read` spin that yields an empty snapshot under write contention), @@ -594,6 +622,14 @@ mod tests { assert!(cache.get("a").await.is_none()); assert_eq!(cache.get("b").await, Some(2)); assert_eq!(cache.get("d").await, Some(4)); + assert_eq!( + cache.capacity_stats().await, + CapacityStats { + entries: 3, + evictions: 1, + eviction_blocks: 0, + } + ); } #[tokio::test] @@ -659,6 +695,14 @@ mod tests { 3, "all entries held -> cache exceeds capacity instead of evicting a live lock" ); + assert_eq!( + cache.capacity_stats().await, + CapacityStats { + entries: 3, + evictions: 0, + eviction_blocks: 1, + } + ); // Drop the external refs; the next insert now evicts back down to capacity. drop(held); @@ -670,6 +714,14 @@ mod tests { 2, "once entries are released, eviction resumes down to capacity" ); + assert_eq!( + cache.capacity_stats().await, + CapacityStats { + entries: 2, + evictions: 2, + eviction_blocks: 1, + } + ); } #[tokio::test] diff --git a/src/retry.rs b/src/retry.rs index 3fd93001e..c7a2f994e 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -379,6 +379,7 @@ impl Client { // WA Web rotateKey: unknown device (not in participant list, not LID) → // force full sender key rotation by clearing all sender key device tracking. // This is separate from updateLocalSignalSession and specific to group retries. + let mut rotated_sender_key = false; if is_group_or_status && !info.requester.is_lid() && !info.chat.is_status_broadcast() { let group_jid = info.chat.to_string(); let is_known_participant = cached_group_info @@ -392,6 +393,7 @@ impl Client { info.requester.observe(), group_jid ); + let _distribution_guard = self.group_distribution_lock(&info.chat).await; // WA Web: deleteGroupSenderKeyInfo(groupWid, ownWid) — delete our own // sender key for forward secrecy. When addressing mode is known, @@ -425,16 +427,16 @@ impl Client { // DB first, then cache invalidate — prevents a concurrent // resolve_skdm_targets from reviving stale cache entries. - if let Err(e) = self - .persistence_manager - .clear_sender_key_devices(&group_jid) - .await - { + if let Err(e) = self.reset_sender_key_device_tracking(&group_jid).await { log::warn!("Failed to clear sender key devices for rotation: {}", e); } - self.sender_key_device_cache.invalidate(&group_jid).await; + rotated_sender_key = true; } } + if rotated_sender_key { + self.flush_signal_cache_batch_safe_logged("unknown-participant rotation", None) + .await; + } // Mirror WAWebUpdateLocalSignalSession for all chat types: markForgetSenderKey // (group/status) + processKeyBundle + regId-mismatch delete + base-key logic. @@ -2537,6 +2539,114 @@ mod tests { ); } + #[tokio::test] + async fn unknown_participant_rotation_is_durable_before_throttled_return() { + use wacore::libsignal::protocol::{SENDERKEY_MESSAGE_CURRENT_VERSION, SenderKeyRecord}; + use wacore::libsignal::store::sender_key_name::SenderKeyName; + use wacore_binary::builder::NodeBuilder; + + let backend = crate::test_utils::create_test_backend().await; + let pm = Arc::new(PersistenceManager::new(backend.clone()).await.unwrap()); + let mut config = crate::cache_config::CacheConfig::default(); + config.recent_messages.capacity = 1_000; + let (client, _rx) = Client::new_with_cache_config( + Arc::new(crate::runtime_impl::TokioRuntime), + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + config, + ) + .await; + + let own_lid: Jid = "100000000001040:13@lid".parse().unwrap(); + client + .persistence_manager + .process_command(crate::store::commands::DeviceCommand::SetLid(Some( + own_lid.clone(), + ))) + .await; + let group: Jid = "120363021033254950@g.us".parse().unwrap(); + let group_id = group.to_string(); + let sender_key_name = + SenderKeyName::from_parts(&group_id, own_lid.to_protocol_address().as_str()); + let mut rng = rand::make_rng::(); + let key_pair = KeyPair::generate(&mut rng); + let mut record = SenderKeyRecord::new_empty(); + record + .add_sender_key_state( + SENDERKEY_MESSAGE_CURRENT_VERSION, + 9, + 0, + &[7; 32], + key_pair.public_key, + Some(key_pair.private_key), + ) + .unwrap(); + client + .signal_cache + .put_sender_key(&sender_key_name, record) + .await; + client.flush_signal_cache().await.unwrap(); + assert!( + backend + .get_sender_key(sender_key_name.cache_key()) + .await + .unwrap() + .is_some() + ); + + let msg_id = "ROTATEFLUSH001"; + client + .add_recent_message( + &group, + msg_id, + &wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }, + None, + ) + .await; + client.set_resend_rate_limit(1, 0); + assert!(client.resend_rate_limiter.try_acquire(&group).await); + + let requester: Jid = "15551234002@s.whatsapp.net".parse().unwrap(); + let node = NodeBuilder::new("receipt") + .attr("participant", &requester) + .children([NodeBuilder::new("retry") + .attr("id", msg_id) + .attr("count", "1") + .build()]) + .build(); + let node_ref = crate::test_utils::node_to_owned_ref(&node); + let receipt = Receipt::builder() + .source(crate::types::message::MessageSource { + chat: group.clone(), + sender: requester, + is_group: true, + ..Default::default() + }) + .message_ids(vec![msg_id.to_string()]) + .timestamp(wacore::time::now_utc()) + .r#type(crate::types::presence::ReceiptType::Retry) + .offline(false) + .build(); + + client + .handle_retry_receipt(&receipt, &node_ref) + .await + .unwrap(); + assert!( + backend + .get_sender_key(sender_key_name.cache_key()) + .await + .unwrap() + .is_none(), + "early retry return must not leave the retired key durable" + ); + } + /// Atomicity guard for the per-peer session lock the retry caller wraps /// around the recreate check+stamp. The cache's get+insert is not atomic, and /// same-peer retries for different message_ids dispatch concurrently, so diff --git a/src/send/mod.rs b/src/send/mod.rs index 87b97f4f5..001ad8bf1 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -828,6 +828,7 @@ impl Client { let device_store_arc = self.persistence_manager.get_device_arc().await; let to_str = to.to_string(); + let distribution_guard = self.group_distribution_lock(&to).await; let force_skdm = { use wacore::libsignal::store::sender_key_name::SenderKeyName; @@ -844,6 +845,10 @@ impl Client { .await? .is_some(); + if !key_exists { + self.reset_sender_key_device_tracking(&to_str).await?; + } + !key_exists }; @@ -916,18 +921,7 @@ impl Client { { log::warn!("No sender key for status broadcast, forcing distribution."); - if let Err(e) = self - .persistence_manager - .clear_sender_key_devices(&to_str) - .await - { - log::warn!( - "Failed to clear status SKDM recipients for {}: {:?}", - to_str, - e - ); - } - self.sender_key_device_cache.invalidate(&to_str).await; + self.reset_sender_key_device_tracking(&to_str).await?; let mut store_adapter_retry = self.signal_adapter_from(device_store_arc.clone()); @@ -990,6 +984,7 @@ impl Client { self.update_sender_key_devices(&to_str, &prepared.skdm_devices) .await; + drop(distribution_guard); for user in &prepared.stale_device_users { self.invalidate_device_cache(user).await; @@ -1342,12 +1337,10 @@ impl Client { // distribution path. If the clear fails, fall back to deleting the bot's // own sender key for the chat — the next send will see `!key_exists` and // force_skdm without depending on the tracker. + let mut flush_fallback = false; if jid.is_group() || jid.is_status_broadcast() { - let cleared = self - .persistence_manager - .clear_sender_key_devices(&jid_str) - .await; - if let Err(e) = cleared { + let distribution_guard = self.group_distribution_lock(jid).await; + if let Err(e) = self.reset_sender_key_device_tracking(&jid_str).await { log::warn!( "phash mismatch: clear_sender_key_devices failed: {e} — \ deleting own sender key as fallback to force redistribution" @@ -1360,12 +1353,17 @@ impl Client { SenderKeyName::from_parts(&jid_str, own.to_protocol_address().as_str()); self.signal_cache.delete_sender_key(sk.cache_key()).await; } - let _ = self - .flush_signal_cache_batch_safe_logged("phash-mismatch-fallback", None) - .await; + flush_fallback = true; } + drop(distribution_guard); + } else { + self.sender_key_device_cache.invalidate(&jid_str).await; + } + if flush_fallback { + let _ = self + .flush_signal_cache_batch_safe_logged("phash-mismatch-fallback", None) + .await; } - self.sender_key_device_cache.invalidate(&jid_str).await; if invalidate_group_cache { self.get_group_cache().await.invalidate(jid).await; } @@ -1714,6 +1712,9 @@ impl Client { distribution_guard = Some(self.group_distribution_lock(&to).await); let (key_exists, needs_rotation) = read_sender_key_state().await?; force_skdm = force_key_distribution || !key_exists || needs_rotation; + if !key_exists || needs_rotation { + self.reset_sender_key_device_tracking(&to_str).await?; + } if needs_rotation { log::info!( "Periodic sender-key rotation for {} (chain iteration >= {SENDER_KEY_ROTATION_THRESHOLD})", @@ -1722,14 +1723,6 @@ impl Client { self.signal_cache .delete_sender_key(sender_key_name.cache_key()) .await; - if let Err(e) = self - .persistence_manager - .clear_sender_key_devices(&to_str) - .await - { - log::warn!("periodic rotation: clear_sender_key_devices failed: {e}"); - } - self.sender_key_device_cache.invalidate(&to_str).await; } if !force_skdm { distribution_guard = None; @@ -1886,14 +1879,7 @@ impl Client { let (retry_force, retry_targets, retry_all) = match warm_targets { Some((all, needs)) => (false, Some(needs), Some(all)), None => { - if let Err(e) = self - .persistence_manager - .clear_sender_key_devices(&to_str) - .await - { - log::warn!("Failed to clear SKDM recipients: {:?}", e); - } - self.sender_key_device_cache.invalidate(&to_str).await; + self.reset_sender_key_device_tracking(&to_str).await?; (true, None, None) } }; @@ -2276,6 +2262,7 @@ pub(crate) fn dm_stanza_to(recipient_bare: &Jid, to: &Jid) -> Jid { #[cfg(test)] mod tests { use super::*; + use crate::test_utils::wait_for_lock_waiter; use std::str::FromStr; #[test] @@ -2381,6 +2368,55 @@ mod tests { ); } + #[tokio::test] + async fn status_send_waits_for_distribution_guard() { + let client = crate::test_utils::create_test_client().await; + let own_pn: Jid = "15551234001@s.whatsapp.net".parse().unwrap(); + let own_lid: Jid = "100000000000001@lid".parse().unwrap(); + client + .persistence_manager + .process_command(crate::store::commands::DeviceCommand::SetId(Some(own_pn))) + .await; + client + .persistence_manager + .process_command(crate::store::commands::DeviceCommand::SetLid(Some(own_lid))) + .await; + + let status = Jid::status_broadcast(); + let held = client.group_distribution_lock(&status).await; + let lock = client + .group_distribution_locks + .get(&status) + .await + .expect("cached distribution lock"); + let lock_refs = std::sync::Arc::strong_count(&lock); + let mut task = tokio::spawn({ + let client = client.clone(); + async move { + let recipient: Jid = "100000000000002@lid".parse().unwrap(); + client + .send_status_message( + wa::Message { + conversation: Some("serialized status".into()), + ..Default::default() + }, + std::slice::from_ref(&recipient), + crate::features::status::StatusSendOptions::default(), + ) + .await + } + }); + + wait_for_lock_waiter(&lock, lock_refs).await; + assert!(!task.is_finished(), "status send must wait for the lane"); + drop(held); + + let _ = tokio::time::timeout(std::time::Duration::from_secs(5), &mut task) + .await + .expect("status send must resume") + .expect("status task"); + } + // A logged-out send goes through send_message_impl, whose internal // `ClientError::NotLoggedIn` is threaded as `anyhow`. The wrapper must // surface the typed `SendError::NotLoggedIn`, not the `Internal` catch-all, diff --git a/src/store/persistence_manager.rs b/src/store/persistence_manager.rs index 2e9328bd3..1bf74ec10 100644 --- a/src/store/persistence_manager.rs +++ b/src/store/persistence_manager.rs @@ -22,6 +22,10 @@ pub struct PersistenceManager { save_notify: Arc, /// Set to true when the background saver halts due to repeated flush failures. saver_halted: Arc, + #[cfg(test)] + fail_sender_key_device_clears: AtomicBool, + #[cfg(test)] + fail_sender_key_device_status_writes: AtomicBool, } impl PersistenceManager { @@ -63,6 +67,10 @@ impl PersistenceManager { dirty: Arc::new(AtomicBool::new(false)), save_notify: Arc::new(Event::new()), saver_halted: Arc::new(AtomicBool::new(false)), + #[cfg(test)] + fail_sender_key_device_clears: AtomicBool::new(false), + #[cfg(test)] + fail_sender_key_device_status_writes: AtomicBool::new(false), }) } @@ -274,13 +282,40 @@ impl PersistenceManager { group_jid: &str, entries: &[(&str, bool)], ) -> Result<(), StoreError> { + #[cfg(test)] + if self + .fail_sender_key_device_status_writes + .load(Ordering::Acquire) + { + return Err(StoreError::Io(std::io::Error::other( + "injected sender-key tracker status-write failure", + ))); + } self.backend.set_sender_key_status(group_jid, entries).await } pub async fn clear_sender_key_devices(&self, group_jid: &str) -> Result<(), StoreError> { + #[cfg(test)] + if self.fail_sender_key_device_clears.load(Ordering::Acquire) { + return Err(StoreError::Io(std::io::Error::other( + "injected sender-key tracker clear failure", + ))); + } self.backend.clear_sender_key_devices(group_jid).await } + #[cfg(test)] + pub(crate) fn fail_sender_key_device_clears_for_tests(&self, fail: bool) { + self.fail_sender_key_device_clears + .store(fail, Ordering::Release); + } + + #[cfg(test)] + pub(crate) fn fail_sender_key_device_status_writes_for_tests(&self, fail: bool) { + self.fail_sender_key_device_status_writes + .store(fail, Ordering::Release); + } + pub async fn delete_sender_key_device_rows( &self, device_jids: &[&str], diff --git a/src/test_utils.rs b/src/test_utils.rs index b1ecc2cd5..4dbc99e9e 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -13,6 +13,17 @@ pub fn node_to_owned_ref(node: &Node) -> Arc { Arc::new(OwnedNodeRef::new(bytes).expect("OwnedNodeRef::new should succeed")) } } + +pub async fn wait_for_lock_waiter(lock: &Arc>, baseline: usize) { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while Arc::strong_count(lock) <= baseline { + tokio::task::yield_now().await; + } + }) + .await + .expect("task must reach the contested lock"); +} + use crate::http::{HttpClient, HttpRequest, HttpResponse}; use crate::runtime_impl::TokioRuntime; use crate::store::SqliteStore; diff --git a/wacore/libsignal/src/protocol/group_cipher.rs b/wacore/libsignal/src/protocol/group_cipher.rs index 6579f3f69..42e61065c 100644 --- a/wacore/libsignal/src/protocol/group_cipher.rs +++ b/wacore/libsignal/src/protocol/group_cipher.rs @@ -200,6 +200,8 @@ fn get_sender_key(state: &mut SenderKeyState, iteration: u32) -> Result Result( sender_key_name: &SenderKeyName, sender_key_store: &mut dyn SenderKeyStore, diff --git a/wacore/libsignal/src/protocol/storage/traits.rs b/wacore/libsignal/src/protocol/storage/traits.rs index 6d6f4dc1e..31e22326b 100644 --- a/wacore/libsignal/src/protocol/storage/traits.rs +++ b/wacore/libsignal/src/protocol/storage/traits.rs @@ -163,9 +163,9 @@ pub trait SenderKeyStore: ThreadSafe { sender_key_name: &SenderKeyName, ) -> Result>; - /// Serializes the load/advance/store of one sender-key chain so concurrent - /// encrypts to the same `(group, sender)` can't reuse a chain iteration. - /// Default is uncontended; stores over shared state override it. + /// Serializes every load/mutate/store of one sender-key chain so concurrent + /// encrypt, decrypt, distribution and rotation operations cannot overwrite + /// each other. Default is uncontended; stores over shared state override it. async fn sender_key_lock( &self, _sender_key_name: &SenderKeyName, diff --git a/wacore/src/store/signal_cache.rs b/wacore/src/store/signal_cache.rs index 4e21b2f77..2071a4762 100644 --- a/wacore/src/store/signal_cache.rs +++ b/wacore/src/store/signal_cache.rs @@ -686,7 +686,10 @@ impl SignalStoreCache { lock } + /// Prevent an in-flight mutation from storing the retired chain again. pub async fn delete_sender_key(&self, cache_key: &str) { + let lock = self.shared_named_lock(cache_key).await; + let _guard = lock.lock().await; let mut state = self.sender_keys.lock().await; state.delete(cache_key); } @@ -1018,6 +1021,16 @@ mod sender_key_lock_tests { use super::*; use crate::libsignal::store::sender_key_name::SenderKeyName; + async fn wait_for_lock_waiter(lock: &Arc>, baseline: usize) { + for _ in 0..10_000 { + if Arc::strong_count(lock) > baseline { + return; + } + tokio::task::yield_now().await; + } + panic!("task did not reach the contested lock"); + } + #[tokio::test] async fn same_name_shares_one_lock() { let cache = SignalStoreCache::new(); @@ -1047,6 +1060,52 @@ mod sender_key_lock_tests { assert!(lock.try_lock().is_some(), "released lock must reacquire"); } + #[tokio::test] + async fn delete_waits_for_the_chain_lock() { + let cache = Arc::new(SignalStoreCache::new()); + let backend = crate::store::in_memory::InMemoryBackend::new(); + let name = SenderKeyName::from_parts("g@g.us", "u@s.whatsapp.net:0"); + cache + .put_sender_key(&name, SenderKeyRecord::new_empty()) + .await; + + let lock = cache.sender_key_lock(&name).await; + let held = lock.lock().await; + let lock_refs = Arc::strong_count(&lock); + let started = Arc::new(async_lock::Barrier::new(2)); + let task = tokio::spawn({ + let cache = cache.clone(); + let started = started.clone(); + let cache_key = name.cache_key().to_string(); + async move { + started.wait().await; + cache.delete_sender_key(&cache_key).await; + } + }); + + started.wait().await; + wait_for_lock_waiter(&lock, lock_refs).await; + assert!( + cache + .get_sender_key(&name, &backend) + .await + .unwrap() + .is_some(), + "delete must wait for the in-flight chain mutation" + ); + + drop(held); + task.await.expect("delete task"); + assert!( + cache + .get_sender_key(&name, &backend) + .await + .unwrap() + .is_none(), + "delete must run after the mutation releases the chain" + ); + } + #[tokio::test] async fn warm_sender_key_hit_shares_arc_not_deep_clone() { let cache = SignalStoreCache::new();