From fad388adead7330684c43904254420fb9f6a144d Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 9 Jun 2026 18:56:47 +0000 Subject: [PATCH 1/2] perf(send): establish sessions before taking the sender-key chain lock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit prepare_group_stanza held the per-(group, sender) chain lock across the whole SKDM path: device resolution and the prekey fetch + X3DH inside encrypt_for_devices — network round-trips — ran inside the critical section, so concurrent sends to the same group serialized behind an RTT (or its retries) whenever any target device lacked a session. Split encrypt_for_devices into its two halves: ensure_sessions_for_devices (network: LID-first lookup, prekey fetch, parallel X3DH; touches only session/identity state) and encrypt_for_devices_with_sessions (CPU: the pairwise fan-out). The combined function remains as a composition for the DM path. The group path now resolves devices and ensures sessions before the lock; the lock covers only SKDM creation + pairwise encrypt + skmsg — the chain-consistency invariant it exists for. Failure semantics preserved: session-setup errors log and continue without distribution (WA Web GroupSkmsgJob), and the sender-key record is still created under the lock so the skmsg always encrypts. The regression test probes the actual chain lock from inside the mock resolver's fetch; it fails against the previous code. Test mem stores now share state across clones (Arc), matching production store semantics so spawned-task session writes are visible. --- wacore/src/send/encrypt.rs | 97 ++++++++++++++++++-- wacore/src/send/group.rs | 154 ++++++++++++++++++------------- wacore/src/send/tests.rs | 181 +++++++++++++++++++++++++++++++++++-- 3 files changed, 355 insertions(+), 77 deletions(-) diff --git a/wacore/src/send/encrypt.rs b/wacore/src/send/encrypt.rs index 7f7766400..86c96ed13 100644 --- a/wacore/src/send/encrypt.rs +++ b/wacore/src/send/encrypt.rs @@ -286,9 +286,14 @@ fn push_encrypt_result( /// by [`ENCRYPT_FANOUT_CONCURRENCY`]. Each task clones the store handles /// (Arc bumps under the hood); the shared cache provides interior mutability. /// +/// Composition of [`ensure_sessions_for_devices`] (network: prekey fetch + +/// X3DH for missing sessions) and [`encrypt_for_devices_with_sessions`] +/// (CPU: the pairwise encrypt fan-out). Callers that must not hold a lock +/// across network I/O (the group sender-key chain lock) call the two phases +/// directly with the lock taken only around the second. +/// /// Callers must hold per-device session locks before calling this function — /// concurrent ratchet mutations will corrupt Signal session state. -#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.encrypt_fanout", level = "debug", skip_all, fields(count = devices.len()), err(Debug)))] pub async fn encrypt_for_devices<'a, S, I, P, SP>( runtime: &dyn Runtime, stores: &mut SignalStores<'a, S, I, P, SP>, @@ -298,6 +303,46 @@ pub async fn encrypt_for_devices<'a, S, I, P, SP>( hide_decrypt_fail: bool, mediatype: Option<&str>, ) -> Result +where + S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static, + I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static, + P: crate::libsignal::protocol::PreKeyStore + Send + Sync, + SP: crate::libsignal::protocol::SignedPreKeyStore + Send + Sync, +{ + let plan = ensure_sessions_for_devices(runtime, stores, resolver, devices).await?; + encrypt_for_devices_with_sessions( + runtime, + stores, + devices, + plaintext_to_encrypt, + hide_decrypt_fail, + mediatype, + plan, + ) + .await +} + +/// Session material prepared for one encrypt fan-out: per-index LID +/// encryption overrides (mirroring the `devices` slice it was built from) +/// plus whether any device 406'd during prekey fetch. Produced only by +/// [`ensure_sessions_for_devices`]; consumed by +/// [`encrypt_for_devices_with_sessions`] over the same `devices` slice. +pub struct SessionPlan { + encryption_overrides: Vec>, + pub had_unregistered_device: bool, +} + +/// Resolve LID overrides and establish missing Signal sessions (prekey +/// fetch + X3DH) for `devices`. This is the network half of the encrypt +/// fan-out and touches only session/identity state — never a sender-key +/// chain — so group sends run it before taking the chain lock. +#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.ensure_sessions", level = "debug", skip_all, fields(count = devices.len()), err(Debug)))] +pub async fn ensure_sessions_for_devices<'a, S, I, P, SP>( + runtime: &dyn Runtime, + stores: &mut SignalStores<'a, S, I, P, SP>, + resolver: &dyn SendContextResolver, + devices: &[Jid], +) -> Result where S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static, I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static, @@ -489,6 +534,44 @@ where } } + Ok(SessionPlan { + encryption_overrides, + had_unregistered_device: had_406, + }) +} + +/// CPU half of the encrypt fan-out: pairwise-encrypt `plaintext_to_encrypt` +/// for each device using sessions prepared by [`ensure_sessions_for_devices`] +/// over the same `devices` slice. No resolver, no network — safe to run +/// under locks that must not span I/O. A device whose session is still +/// missing (e.g. its bundle was absent) fails its encrypt and is skipped, +/// matching the combined path's behavior. +#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.encrypt_fanout", level = "debug", skip_all, fields(count = devices.len()), err(Debug)))] +pub async fn encrypt_for_devices_with_sessions<'a, S, I, P, SP>( + runtime: &dyn Runtime, + stores: &mut SignalStores<'a, S, I, P, SP>, + devices: &[Jid], + plaintext_to_encrypt: &[u8], + hide_decrypt_fail: bool, + mediatype: Option<&str>, + plan: SessionPlan, +) -> Result +where + S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static, + I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static, + P: crate::libsignal::protocol::PreKeyStore + Send + Sync, + SP: crate::libsignal::protocol::SignedPreKeyStore + Send + Sync, +{ + debug_assert_eq!( + plan.encryption_overrides.len(), + devices.len(), + "SessionPlan built for a different device list" + ); + let SessionPlan { + encryption_overrides, + had_unregistered_device, + } = plan; + let mut participant_nodes = Vec::with_capacity(devices.len()); let mut includes_prekey_message = false; let mut encrypted_devices = Vec::with_capacity(devices.len()); @@ -502,8 +585,9 @@ where // a FuturesUnordered, and two store clones), with no parallelism to gain. // Encrypt inline. let device_jid = devices[0].clone(); - let addr = encryption_overrides[0] - .as_ref() + let addr = encryption_overrides + .first() + .and_then(|o| o.as_ref()) .unwrap_or(&devices[0]) .to_protocol_address(); let res = encrypt_one_device( @@ -536,8 +620,9 @@ where // The encryption JID is only needed to build the Signal address, so // derive it here from a borrow rather than cloning the whole Jid into // the task (device_jid is still cloned because it's returned). - let addr = encryption_overrides[idx] - .as_ref() + let addr = encryption_overrides + .get(idx) + .and_then(|o| o.as_ref()) .unwrap_or(&devices[idx]) .to_protocol_address(); let plaintext = plaintext_arc.clone(); @@ -587,6 +672,6 @@ where participant_nodes, includes_prekey_message, encrypted_devices, - had_unregistered_device: had_406, + had_unregistered_device, }) } diff --git a/wacore/src/send/group.rs b/wacore/src/send/group.rs index 49bbc8d9a..8a6386871 100644 --- a/wacore/src/send/group.rs +++ b/wacore/src/send/group.rs @@ -172,17 +172,9 @@ pub async fn prepare_group_stanza< let mut phash_for_stanza: Option = None; let mut skdm_encrypted_devices: Vec = Vec::new(); - // Build the chain name once and hold its lock across SKDM creation + the - // skmsg encrypt, so concurrent same-(group, sender) sends can't split the - // key between the SKDM and the skmsg (nor reuse a chain iteration). - let sender_key_name = make_sender_key_name(&to_jid, &own_sending_jid.to_protocol_address()); - let chain_lock = stores - .sender_key_store - .sender_key_lock(&sender_key_name) - .await; - let _chain_guard = chain_lock.lock().await; - - // Determine if we need to distribute SKDM and to which devices + // Determine if we need to distribute SKDM and to which devices. + // Resolved before the chain lock below: device resolution is network I/O + // on state independent of the sender-key chain. let distribution_list: Option> = if let Some(target_devices) = skdm_target_devices { // Use the specific list of devices that need SKDM if target_devices.is_empty() { @@ -338,69 +330,109 @@ pub async fn prepare_group_stanza< let mut had_unregistered_devices = false; + // Establish missing pairwise sessions (prekey fetch + X3DH) for the SKDM + // targets before taking the chain lock, so the critical section below + // never spans a network RTT — concurrent sends to the same group would + // otherwise serialize behind it. WA Web's GroupSkmsgJob wraps + // ensureE2ESessions in try/catch — logs but does NOT rethrow: a session + // setup failure must not prevent the group message from being sent. + let session_plan = match distribution_list.as_deref() { + Some(list) => match ensure_sessions_for_devices(runtime, stores, resolver, list).await { + Ok(plan) => Some(plan), + Err(e) => { + log::warn!( + "SKDM session setup failed for group {}, continuing without distribution: {e}", + to_jid.observe() + ); + if is_device_unregistered_error(&e) { + had_unregistered_devices = true; + } + None + } + }, + None => None, + }; + + // Build the chain name once and hold its lock across SKDM creation + the + // skmsg encrypt, so concurrent same-(group, sender) sends can't split the + // key between the SKDM and the skmsg (nor reuse a chain iteration). The + // lock covers only chain-touching steps; session setup and device + // resolution stay outside it. + let sender_key_name = make_sender_key_name(&to_jid, &own_sending_jid.to_protocol_address()); + let chain_lock = stores + .sender_key_store + .sender_key_lock(&sender_key_name) + .await; + let _chain_guard = chain_lock.lock().await; + if let Some(ref distribution_list) = distribution_list { + // Created even when session setup failed (plan None): the sender-key + // record must exist so the skmsg below still encrypts, preserving the + // continue-without-distribution semantics. let axolotl_skdm_bytes = create_sender_key_distribution_message_for_group( stores.sender_key_store, &sender_key_name, ) .await?; - let skdm_wrapper_msg = wa::Message { - sender_key_distribution_message: Some(wa::message::SenderKeyDistributionMessage { - group_id: Some(to_jid.to_string()), - axolotl_sender_key_distribution_message: Some(axolotl_skdm_bytes), - }), - ..Default::default() - }; - let skdm_plaintext_to_encrypt = MessageUtils::encode_and_pad(&skdm_wrapper_msg); - - // WA Web's GroupSkmsgJob wraps ensureE2ESessions in try/catch — logs error - // but does NOT rethrow. SKDM distribution failure must not prevent the group - // message from being sent. Only successfully encrypted devices are tracked. - // Must match the rule applied to the main skmsg payload below: if SKDM carries - // `decrypt-fail="hide"` but the payload does not (e.g. AdminRevoke), recipients - // without a sender key never decrypt the skmsg and the revoke is silently dropped. - let skdm_hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit.as_ref(), message); - match encrypt_for_devices( - runtime, - stores, - resolver, - distribution_list, - &skdm_plaintext_to_encrypt, - skdm_hide_decrypt_fail, - None, - ) - .await - { - Ok(result) => { - includes_prekey_message = includes_prekey_message || result.includes_prekey_message; - if result.had_unregistered_device { - had_unregistered_devices = true; - } - skdm_encrypted_devices = result.encrypted_devices; + if let Some(plan) = session_plan { + let skdm_wrapper_msg = wa::Message { + sender_key_distribution_message: Some(wa::message::SenderKeyDistributionMessage { + group_id: Some(to_jid.to_string()), + axolotl_sender_key_distribution_message: Some(axolotl_skdm_bytes), + }), + ..Default::default() + }; + let skdm_plaintext_to_encrypt = MessageUtils::encode_and_pad(&skdm_wrapper_msg); + + // SKDM distribution failure must not prevent the group message from + // being sent. Only successfully encrypted devices are tracked. + // Must match the rule applied to the main skmsg payload below: if SKDM carries + // `decrypt-fail="hide"` but the payload does not (e.g. AdminRevoke), recipients + // without a sender key never decrypt the skmsg and the revoke is silently dropped. + let skdm_hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit.as_ref(), message); + match encrypt_for_devices_with_sessions( + runtime, + stores, + distribution_list, + &skdm_plaintext_to_encrypt, + skdm_hide_decrypt_fail, + None, + plan, + ) + .await + { + Ok(result) => { + includes_prekey_message = + includes_prekey_message || result.includes_prekey_message; + if result.had_unregistered_device { + had_unregistered_devices = true; + } + skdm_encrypted_devices = result.encrypted_devices; - if !result.participant_nodes.is_empty() { - message_children.push( - NodeBuilder::new("participants") - .children(result.participant_nodes) - .build(), - ); - if includes_prekey_message && let Some(acc) = account { + if !result.participant_nodes.is_empty() { message_children.push( - NodeBuilder::new("device-identity") - .bytes(acc.encode_to_vec()) + NodeBuilder::new("participants") + .children(result.participant_nodes) .build(), ); + if includes_prekey_message && let Some(acc) = account { + message_children.push( + NodeBuilder::new("device-identity") + .bytes(acc.encode_to_vec()) + .build(), + ); + } } } - } - Err(e) => { - log::warn!( - "SKDM distribution failed for group {}, continuing without it: {e}", - to_jid.observe() - ); - if is_device_unregistered_error(&e) { - had_unregistered_devices = true; + Err(e) => { + log::warn!( + "SKDM distribution failed for group {}, continuing without it: {e}", + to_jid.observe() + ); + if is_device_unregistered_error(&e) { + had_unregistered_devices = true; + } } } } diff --git a/wacore/src/send/tests.rs b/wacore/src/send/tests.rs index 4214395ff..b99855868 100644 --- a/wacore/src/send/tests.rs +++ b/wacore/src/send/tests.rs @@ -348,6 +348,16 @@ fn build_member_label_message_preserves_unicode() { assert_eq!(ml.label.as_deref(), Some("🚀 BOT")); } +/// Probe installed by chain-lock tests: records whether the sender-key chain +/// lock was held while `fetch_prekeys_for_identity_check` ran (it must not be +/// — the fetch is network I/O hoisted out of the chain critical section). +#[derive(Clone, Default)] +struct ChainLockProbe { + lock: std::sync::Arc>, + fetched_under_lock: std::sync::Arc, + fetch_calls: std::sync::Arc, +} + /// Mock implementation of SendContextResolver for testing struct MockSendContextResolver { /// Pre-key bundles to return: JID -> Option @@ -358,6 +368,7 @@ struct MockSendContextResolver { phone_to_lid: HashMap, /// JIDs reported via `on_local_identity_change` (send-path detection). identity_changes: std::sync::Mutex>, + chain_lock_probe: Option, } impl MockSendContextResolver { @@ -367,9 +378,15 @@ impl MockSendContextResolver { devices: Vec::new(), phone_to_lid: HashMap::new(), identity_changes: std::sync::Mutex::new(Vec::new()), + chain_lock_probe: None, } } + fn with_chain_lock_probe(mut self, probe: ChainLockProbe) -> Self { + self.chain_lock_probe = Some(probe); + self + } + fn captured_identity_changes(&self) -> Vec { self.identity_changes.lock().unwrap().clone() } @@ -417,6 +434,16 @@ impl SendContextResolver for MockSendContextResolver { &self, jids: &[Jid], ) -> Result> { + if let Some(probe) = &self.chain_lock_probe { + probe + .fetch_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if probe.lock.try_lock().is_none() { + probe + .fetched_under_lock + .store(true, std::sync::atomic::Ordering::SeqCst); + } + } let mut result = HashMap::new(); for jid in jids { if let Some(bundle_opt) = self.prekey_bundles.get(jid) @@ -2706,8 +2733,12 @@ mod mark_full_distribution_list { type SigResult = crate::libsignal::protocol::error::Result; + // Clones share state (Arc), mirroring production stores: the encrypt + // fan-out spawns tasks over store clones and their writes must be + // visible to the original ("the shared cache provides interior + // mutability"). #[derive(Clone, Default)] - struct MemSessionStore(HashMap>); + struct MemSessionStore(std::sync::Arc>>>); #[async_trait::async_trait] impl SessionStore for MemSessionStore { async fn load_session( @@ -2716,18 +2747,20 @@ mod mark_full_distribution_list { ) -> SigResult> { Ok(self .0 + .lock() + .unwrap() .get(a) .and_then(|b| crate::libsignal::protocol::SessionRecord::deserialize(b).ok())) } async fn has_session(&self, a: &ProtocolAddress) -> SigResult { - Ok(self.0.contains_key(a)) + Ok(self.0.lock().unwrap().contains_key(a)) } async fn store_session( &mut self, a: &ProtocolAddress, r: crate::libsignal::protocol::SessionRecord, ) -> SigResult<()> { - self.0.insert(a.clone(), r.serialize()?); + self.0.lock().unwrap().insert(a.clone(), r.serialize()?); Ok(()) } } @@ -2736,7 +2769,7 @@ mod mark_full_distribution_list { struct MemIdentityStore { pair: IdentityKeyPair, reg_id: u32, - known: HashMap, + known: std::sync::Arc>>, } #[async_trait::async_trait] impl IdentityKeyStore for MemIdentityStore { @@ -2751,7 +2784,7 @@ mod mark_full_distribution_list { a: &ProtocolAddress, id: &IdentityKey, ) -> SigResult { - self.known.insert(a.clone(), *id); + self.known.lock().unwrap().insert(a.clone(), *id); Ok(IdentityChange::from_changed(false)) } async fn is_trusted_identity( @@ -2763,12 +2796,17 @@ mod mark_full_distribution_list { Ok(true) } async fn get_identity(&self, a: &ProtocolAddress) -> SigResult> { - Ok(self.known.get(a).copied()) + Ok(self.known.lock().unwrap().get(a).copied()) } } #[derive(Default)] - struct MemSenderKeyStore(HashMap); + struct MemSenderKeyStore { + records: HashMap, + // Shared per-name locks (like production stores override it), so tests + // can observe whether the chain lock is held during resolver calls. + locks: std::sync::Mutex>>>, + } #[async_trait::async_trait] impl SenderKeyStore for MemSenderKeyStore { async fn store_sender_key( @@ -2776,11 +2814,22 @@ mod mark_full_distribution_list { n: &SenderKeyName, r: SenderKeyRecord, ) -> SigResult<()> { - self.0.insert(n.clone(), r); + self.records.insert(n.clone(), r); Ok(()) } async fn load_sender_key(&self, n: &SenderKeyName) -> SigResult> { - Ok(self.0.get(n).cloned()) + Ok(self.records.get(n).cloned()) + } + async fn sender_key_lock( + &self, + n: &SenderKeyName, + ) -> std::sync::Arc> { + self.locks + .lock() + .unwrap() + .entry(n.clone()) + .or_default() + .clone() } } @@ -2866,7 +2915,7 @@ mod mark_full_distribution_list { let mut is = MemIdentityStore { pair: sender, reg_id: 42, - known: HashMap::new(), + known: Default::default(), }; process_prekey_bundle( &a.to_protocol_address(), @@ -2964,6 +3013,118 @@ mod mark_full_distribution_list { "a key-distributing group send must carry a phash" ); } + + /// Regression: the prekey fetch (network RTT) must run BEFORE the + /// sender-key chain lock is taken, so concurrent sends to the same group + /// don't serialize behind a slow fetch. The probe try_locks the actual + /// chain lock from inside the resolver's fetch and records a violation. + #[tokio::test] + async fn prekey_fetch_runs_outside_chain_lock() { + use std::sync::atomic::Ordering::SeqCst; + + let group: Jid = "120363000000000002@g.us".parse().unwrap(); + let own_jid: Jid = "559900000000@s.whatsapp.net".parse().unwrap(); + let own_lid: Jid = "100000000000000@lid".parse().unwrap(); + // B has no session but its bundle IS available — forces the prekey + // fetch + X3DH path on this send. + let b: Jid = "559933334444:0@s.whatsapp.net".parse().unwrap(); + + let mut rng = rand::make_rng::(); + let mut ss = MemSessionStore::default(); + let mut is = MemIdentityStore { + pair: IdentityKeyPair::generate(&mut rng), + reg_id: 7, + known: Default::default(), + }; + let mut sks = MemSenderKeyStore::default(); + let chain_name = + crate::types::jid::make_sender_key_name(&group, &own_jid.to_protocol_address()); + let probe = ChainLockProbe { + lock: sks.sender_key_lock(&chain_name).await, + ..Default::default() + }; + let mut pks = UnusedPreKeyStore; + let spks = UnusedSignedPreKeyStore; + let mut stores = SignalStores { + sender_key_store: &mut sks, + session_store: &mut ss, + identity_store: &mut is, + prekey_store: &mut pks, + signed_prekey_store: &spks, + }; + + // Verifiable bundle (create_mock_bundle's zeroed signature fails X3DH). + let receiver = IdentityKeyPair::generate(&mut rng); + let spk = KeyPair::generate(&mut rng); + let opk = KeyPair::generate(&mut rng); + let sig = receiver + .private_key() + .calculate_signature(&spk.public_key.serialize(), &mut rng) + .unwrap(); + let bundle = PreKeyBundle::new( + 1, + 1u32.into(), + Some((1u32.into(), opk.public_key)), + 1u32.into(), + spk.public_key, + sig.to_vec(), + *receiver.identity_key(), + ) + .unwrap(); + + let resolver = MockSendContextResolver::new() + .with_bundle(b.clone(), bundle) + .with_chain_lock_probe(probe.clone()); + let rt = TokioTestRuntime; + + let group_info = + GroupInfo::new(vec![own_jid.to_non_ad(), b.to_non_ad()], AddressingMode::Pn); + let msg = wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }; + + let prepared = prepare_group_stanza( + &rt, + &mut stores, + &resolver, + &group_info, + &own_jid, + &own_lid, + None, + group, + &msg, + "TESTREQID2".into(), + false, + Some(vec![b.clone()]), + None, + None, + &[], + ) + .await + .expect("prepare_group_stanza should succeed"); + + assert!( + probe.fetch_calls.load(SeqCst) >= 1, + "test must exercise the prekey fetch path" + ); + assert!( + !probe.fetched_under_lock.load(SeqCst), + "prekey fetch must not run under the sender-key chain lock" + ); + + // End-to-end: the session established before the lock produced a + // pairwise SKDM for B under the lock. + let participants = prepared + .node + .get_optional_child("participants") + .expect("participants node with the SKDM fan-out"); + assert_eq!( + participants.children().map(|c| c.len()).unwrap_or(0), + 1, + "B must receive a pairwise SKDM via the pre-established session" + ); + } } /// Item 3 — phash device-set construction. The set hashed is the full From 767459153d3a4eed5e49c05f95c29db07a27fc04 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 9 Jun 2026 19:07:38 +0000 Subject: [PATCH 2/2] fix(send): serialize same-group session setup with a dedicated lock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-up: hoisting ensure_sessions_for_devices out of the chain lock let two concurrent cold sends to the same group race prekey fetch + X3DH writes to the same per-device sessions (previously serialized as a side effect of the chain lock). Add SenderKeyStore::session_setup_lock — a per-group lock held only across session setup, never the chain critical section. Cold same-group sends serialize their setup again; warm sends never take it, so the chain lock stays network-free. Default impl is uncontended (mirrors sender_key_lock); the signal cache shares the chain-lock map under a disjoint '::setup'-suffixed key. The regression test now also asserts the fetch runs UNDER the setup lock, alongside the existing not-under-chain-lock assertion. --- src/store/signal_adapter.rs | 7 +++ .../libsignal/src/protocol/storage/traits.rs | 13 +++++ wacore/src/send/group.rs | 54 +++++++++++-------- wacore/src/send/tests.rs | 27 ++++++++++ wacore/src/store/signal_cache.rs | 19 ++++++- 5 files changed, 96 insertions(+), 24 deletions(-) diff --git a/src/store/signal_adapter.rs b/src/store/signal_adapter.rs index 4a8e86753..26bc17a38 100644 --- a/src/store/signal_adapter.rs +++ b/src/store/signal_adapter.rs @@ -309,6 +309,13 @@ impl wacore::libsignal::protocol::SenderKeyStore for SenderKeyAdapter { ) -> std::sync::Arc> { self.0.cache.sender_key_lock(sender_key_name).await } + + async fn session_setup_lock( + &self, + sender_key_name: &SenderKeyName, + ) -> std::sync::Arc> { + self.0.cache.session_setup_lock(sender_key_name).await + } } #[cfg(test)] diff --git a/wacore/libsignal/src/protocol/storage/traits.rs b/wacore/libsignal/src/protocol/storage/traits.rs index e15f8f4a2..6d6f4dc1e 100644 --- a/wacore/libsignal/src/protocol/storage/traits.rs +++ b/wacore/libsignal/src/protocol/storage/traits.rs @@ -172,6 +172,19 @@ pub trait SenderKeyStore: ThreadSafe { ) -> std::sync::Arc> { std::sync::Arc::new(async_lock::Mutex::new(())) } + + /// Serializes per-group session setup (prekey fetch + X3DH) so concurrent + /// cold sends to the same group can't race writes to the same per-device + /// sessions. Distinct from [`sender_key_lock`](Self::sender_key_lock): + /// held only across setup — never the chain-advancing critical section — + /// so it may span network I/O without blocking warm sends. Default is + /// uncontended; stores over shared state override it. + async fn session_setup_lock( + &self, + _sender_key_name: &SenderKeyName, + ) -> std::sync::Arc> { + std::sync::Arc::new(async_lock::Mutex::new(())) + } } /// Mixes in all the store interfaces defined in this module. diff --git a/wacore/src/send/group.rs b/wacore/src/send/group.rs index 8a6386871..1d604aa2d 100644 --- a/wacore/src/send/group.rs +++ b/wacore/src/send/group.rs @@ -330,35 +330,45 @@ pub async fn prepare_group_stanza< let mut had_unregistered_devices = false; + let sender_key_name = make_sender_key_name(&to_jid, &own_sending_jid.to_protocol_address()); + // Establish missing pairwise sessions (prekey fetch + X3DH) for the SKDM - // targets before taking the chain lock, so the critical section below - // never spans a network RTT — concurrent sends to the same group would - // otherwise serialize behind it. WA Web's GroupSkmsgJob wraps - // ensureE2ESessions in try/catch — logs but does NOT rethrow: a session - // setup failure must not prevent the group message from being sent. + // targets before taking the chain lock, so the chain critical section + // below never spans a network RTT — concurrent sends to the same group + // would otherwise serialize behind it. The setup lock serializes this + // phase per group instead: two cold sends can't race fetch + X3DH writes + // to the same per-device sessions, while warm sends (no SKDM) never take + // it. WA Web's GroupSkmsgJob wraps ensureE2ESessions in try/catch — logs + // but does NOT rethrow: a session setup failure must not prevent the + // group message from being sent. let session_plan = match distribution_list.as_deref() { - Some(list) => match ensure_sessions_for_devices(runtime, stores, resolver, list).await { - Ok(plan) => Some(plan), - Err(e) => { - log::warn!( - "SKDM session setup failed for group {}, continuing without distribution: {e}", - to_jid.observe() - ); - if is_device_unregistered_error(&e) { - had_unregistered_devices = true; + Some(list) => { + let setup_lock = stores + .sender_key_store + .session_setup_lock(&sender_key_name) + .await; + let _setup_guard = setup_lock.lock().await; + match ensure_sessions_for_devices(runtime, stores, resolver, list).await { + Ok(plan) => Some(plan), + Err(e) => { + log::warn!( + "SKDM session setup failed for group {}, continuing without distribution: {e}", + to_jid.observe() + ); + if is_device_unregistered_error(&e) { + had_unregistered_devices = true; + } + None } - None } - }, + } None => None, }; - // Build the chain name once and hold its lock across SKDM creation + the - // skmsg encrypt, so concurrent same-(group, sender) sends can't split the - // key between the SKDM and the skmsg (nor reuse a chain iteration). The - // lock covers only chain-touching steps; session setup and device - // resolution stay outside it. - let sender_key_name = make_sender_key_name(&to_jid, &own_sending_jid.to_protocol_address()); + // Hold the chain lock across SKDM creation + the skmsg encrypt, so + // concurrent same-(group, sender) sends can't split the key between the + // SKDM and the skmsg (nor reuse a chain iteration). The lock covers only + // chain-touching steps; session setup and device resolution stay outside. let chain_lock = stores .sender_key_store .sender_key_lock(&sender_key_name) diff --git a/wacore/src/send/tests.rs b/wacore/src/send/tests.rs index b99855868..d3da1ea18 100644 --- a/wacore/src/send/tests.rs +++ b/wacore/src/send/tests.rs @@ -354,7 +354,9 @@ fn build_member_label_message_preserves_unicode() { #[derive(Clone, Default)] struct ChainLockProbe { lock: std::sync::Arc>, + setup_lock: std::sync::Arc>, fetched_under_lock: std::sync::Arc, + fetched_without_setup_lock: std::sync::Arc, fetch_calls: std::sync::Arc, } @@ -443,6 +445,12 @@ impl SendContextResolver for MockSendContextResolver { .fetched_under_lock .store(true, std::sync::atomic::Ordering::SeqCst); } + // The setup lock must be HELD here (try_lock succeeds = violation). + if probe.setup_lock.try_lock().is_some() { + probe + .fetched_without_setup_lock + .store(true, std::sync::atomic::Ordering::SeqCst); + } } let mut result = HashMap::new(); for jid in jids { @@ -2806,6 +2814,8 @@ mod mark_full_distribution_list { // Shared per-name locks (like production stores override it), so tests // can observe whether the chain lock is held during resolver calls. locks: std::sync::Mutex>>>, + setup_locks: + std::sync::Mutex>>>, } #[async_trait::async_trait] impl SenderKeyStore for MemSenderKeyStore { @@ -2831,6 +2841,17 @@ mod mark_full_distribution_list { .or_default() .clone() } + async fn session_setup_lock( + &self, + n: &SenderKeyName, + ) -> std::sync::Arc> { + self.setup_locks + .lock() + .unwrap() + .entry(n.clone()) + .or_default() + .clone() + } } // Outgoing group encryption never consumes our own prekeys, and device B @@ -3041,6 +3062,7 @@ mod mark_full_distribution_list { crate::types::jid::make_sender_key_name(&group, &own_jid.to_protocol_address()); let probe = ChainLockProbe { lock: sks.sender_key_lock(&chain_name).await, + setup_lock: sks.session_setup_lock(&chain_name).await, ..Default::default() }; let mut pks = UnusedPreKeyStore; @@ -3112,6 +3134,11 @@ mod mark_full_distribution_list { !probe.fetched_under_lock.load(SeqCst), "prekey fetch must not run under the sender-key chain lock" ); + assert!( + !probe.fetched_without_setup_lock.load(SeqCst), + "prekey fetch must run under the per-group session-setup lock \ + (serializes same-group cold sends' session writes)" + ); // End-to-end: the session established before the lock produced a // pairwise SKDM for B under the lock. diff --git a/wacore/src/store/signal_cache.rs b/wacore/src/store/signal_cache.rs index 08f117dc7..e1e3b6388 100644 --- a/wacore/src/store/signal_cache.rs +++ b/wacore/src/store/signal_cache.rs @@ -535,8 +535,23 @@ impl SignalStoreCache { /// Shared lock for the `name` chain. Same name returns the same lock so a /// concurrent encrypt can't read a chain iteration another is advancing. pub async fn sender_key_lock(&self, name: &SenderKeyName) -> Arc> { + self.shared_named_lock(name.cache_key()).await + } + + /// Shared per-group session-setup lock (see + /// `SenderKeyStore::session_setup_lock`). Lives in the chain-lock map + /// under a suffixed key; chain cache_keys end in a numeric device id, so + /// the key spaces are disjoint. + pub async fn session_setup_lock(&self, name: &SenderKeyName) -> Arc> { + let mut key = String::with_capacity(name.cache_key().len() + 8); + key.push_str(name.cache_key()); + key.push_str("::setup"); + self.shared_named_lock(&key).await + } + + async fn shared_named_lock(&self, key: &str) -> Arc> { let mut map = self.sender_key_locks.lock().await; - if let Some(lock) = map.get(name.cache_key()) { + if let Some(lock) = map.get(key) { return lock.clone(); } // Drop idle locks (held only by the map) once the map grows large. @@ -544,7 +559,7 @@ impl SignalStoreCache { map.retain(|_, lock| Arc::strong_count(lock) > 1); } let lock = Arc::new(Mutex::new(())); - map.insert(Arc::from(name.cache_key()), lock.clone()); + map.insert(Arc::from(key), lock.clone()); lock }