diff --git a/src/client.rs b/src/client.rs index a2d966f04..f4670ff40 100644 --- a/src/client.rs +++ b/src/client.rs @@ -610,6 +610,17 @@ pub struct Client { pub(crate) group_devices_memo: Cache>, + /// Last `(devices, sender-key-device map)` Arc pair with an empty `needs_skdm`, + /// so a warm repeat send skips `filter_skdm_targets`. `Weak` keeps the pointer + /// comparison ABA-safe, matching `GroupDevicesMemo`. + pub(crate) skdm_warm_memo: Cache< + Jid, + ( + std::sync::Weak, + std::sync::Weak, + ), + >, + /// Router for dispatching stanzas to their appropriate handlers pub(crate) stanza_router: crate::handlers::router::StanzaRouter, diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 0e161fd73..ddff40955 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -243,6 +243,9 @@ impl Client { group_devices_memo: Cache::builder() .max_capacity(GROUP_DEVICES_MEMO_CAPACITY) .build(), + skdm_warm_memo: Cache::builder() + .max_capacity(GROUP_DEVICES_MEMO_CAPACITY) + .build(), stanza_router: Self::create_stanza_router(), synchronous_ack: false, http_client, diff --git a/src/send/mod.rs b/src/send/mod.rs index fc903a042..a7ea4f8c7 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -1048,12 +1048,34 @@ impl Client { .await { Ok(all_devices) => { + // Skip the O(devices) filter_skdm_targets scan when the same + // (devices, sender-key-map) Arc pair was already fully warm. Both + // Arcs swap on any warm-state or membership change, so a stale skip + // is impossible. Needs the device memo for a stable devices Arc. + if self.group_devices_memo_enabled + && let Some((dw, cw)) = self.skdm_warm_memo.get(group).await + && std::ptr::eq(dw.as_ptr(), std::sync::Arc::as_ptr(&all_devices)) + && std::ptr::eq(cw.as_ptr(), std::sync::Arc::as_ptr(&cached_map)) + { + return Some((all_devices, Vec::new())); + } let needs_skdm = self.filter_skdm_targets( group_jid, all_devices.devices(), &cached_map, own_sending_jid, ); + if needs_skdm.is_empty() && self.group_devices_memo_enabled { + self.skdm_warm_memo + .insert( + group.clone(), + ( + std::sync::Arc::downgrade(&all_devices), + std::sync::Arc::downgrade(&cached_map), + ), + ) + .await; + } Some((all_devices, needs_skdm)) } Err(e) => { diff --git a/wacore/src/send/encrypt.rs b/wacore/src/send/encrypt.rs index 30d87de1d..1776cd743 100644 --- a/wacore/src/send/encrypt.rs +++ b/wacore/src/send/encrypt.rs @@ -718,57 +718,71 @@ pub async fn encrypt_for_devices_with_sessions_raw( .await; push_raw_result(res, &mut encrypted, &mut includes_prekey_message); } else { - // Parallel encrypt fan-out across tokio tasks bounded by - // ENCRYPT_FANOUT_CONCURRENCY; collected in completion order so the - // fastest encrypts ship first. + // One task per chunk, not per device: the per-device fan-out allocated a + // task + oneshot + two store clones for every recipient. Same parallelism, + // spawns bounded by ENCRYPT_FANOUT_CONCURRENCY. Wire order is irrelevant + // (phash sorts before hashing on both ends). let plaintext_arc: std::sync::Arc<[u8]> = std::sync::Arc::from(plaintext_to_encrypt); let total = devices.len(); - let mut next_spawn = 0usize; + let num_chunks = ENCRYPT_FANOUT_CONCURRENCY.min(total); - let make_encrypt_task = |idx: usize| { - let device_jid = devices[idx].clone(); - // 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 - .get(idx) - .and_then(|o| o.as_ref()) - .unwrap_or(&devices[idx]) - .to_protocol_address(); + let mut in_flight: FuturesUnordered<_> = FuturesUnordered::new(); + // Index partitioning gives exactly num_chunks slices (keeps the configured + // parallelism) and no-ops on an empty device set instead of dividing by zero. + for chunk_idx in 0..num_chunks { + let chunk_start = chunk_idx * total / num_chunks; + let chunk_end = (chunk_idx + 1) * total / num_chunks; + // The 'static task can't borrow devices/encryption_overrides. + let jobs: Vec<(ProtocolAddress, Jid)> = (chunk_start..chunk_end) + .map(|idx| { + let addr = encryption_overrides + .get(idx) + .and_then(|o| o.as_ref()) + .unwrap_or(&devices[idx]) + .to_protocol_address(); + (addr, devices[idx].clone()) + }) + .collect(); let plaintext = plaintext_arc.clone(); + // clone_box shares the Arc-backed backend, so the sequential ratchet + // advances persist despite one clone serving the whole chunk. let mut session_store = stores.session_store.clone_box(); let mut identity_store = stores.identity_store.clone_box(); - spawn_oneshot(runtime, async move { - encrypt_one_device( - &plaintext, - &addr, - &mut *session_store, - &mut *identity_store, - device_jid, - ) - .await - }) - }; - - let mut in_flight: FuturesUnordered<_> = FuturesUnordered::new(); - while next_spawn < total && in_flight.len() < ENCRYPT_FANOUT_CONCURRENCY { - in_flight.push(make_encrypt_task(next_spawn)); - next_spawn += 1; + in_flight.push(spawn_oneshot(runtime, async move { + let mut out = Vec::with_capacity(jobs.len()); + for (addr, device_jid) in jobs { + out.push( + encrypt_one_device( + &plaintext, + &addr, + &mut *session_store, + &mut *identity_store, + device_jid, + ) + .await, + ); + } + out + })); } while let Some(spawn_result) = in_flight.next().await { match spawn_result { - Ok(res) => push_raw_result(res, &mut encrypted, &mut includes_prekey_message), + Ok(results) => { + for res in results { + push_raw_result(res, &mut encrypted, &mut includes_prekey_message); + } + } Err(SpawnCanceled) => { - log::warn!("Encrypt task did not deliver a result; skipping device."); + // A whole chunk drops (not one device); its members stay + // un-warm and are re-targeted next send. + log::warn!( + "Encrypt chunk did not deliver a result; up to ~{} device(s) skipped this send.", + total.div_ceil(num_chunks) + ); } } - - if next_spawn < total { - in_flight.push(make_encrypt_task(next_spawn)); - next_spawn += 1; - } } } diff --git a/wacore/src/send/group.rs b/wacore/src/send/group.rs index 2246b95c8..be9f39a3b 100644 --- a/wacore/src/send/group.rs +++ b/wacore/src/send/group.rs @@ -371,20 +371,31 @@ pub async fn prepare_group_stanza( None => None, }; - // 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. + // Padding is chain-independent; compute it before the lock so the + // per-(group,sender) serialization point covers only the ratchet steps. + let plaintext = match &shared_content { + Some(content) => { + MessageUtils::pad_with_context_from_encoded(content, reporting_context.as_ref()) + } + None => MessageUtils::encode_and_pad_with_context(message, reporting_context.as_ref()), + }; + + // The lock spans SKDM creation, the pairwise SKDM fan-out, and the skmsg + // encrypt. Creating the SKDM snapshots the sender key and the skmsg uses it, + // so those must be atomic. The fan-out also mutates the shared per-device + // Signal sessions, which the group path (unlike the DM path) does not lock, + // so it has to stay serialized here too or concurrent same-group sends race + // those sessions. Dropped after the encrypt so only the stanza build runs + // off the serialization point. let chain_lock = stores .sender_key_store .sender_key_lock(&sender_key_name) .await; - let _chain_guard = chain_lock.lock().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. + // record must exist so the skmsg below still encrypts. let axolotl_skdm_bytes = create_sender_key_distribution_message_for_group( stores.sender_key_store, &sender_key_name, @@ -460,14 +471,6 @@ pub async fn prepare_group_stanza( } } - // Reuse the shared encode (also fed to the reporting token) for the skmsg plaintext - // instead of encoding the message a second time; the mci-hoist path re-encodes. - let plaintext = match &shared_content { - Some(content) => { - MessageUtils::pad_with_context_from_encoded(content, reporting_context.as_ref()) - } - None => MessageUtils::encode_and_pad_with_context(message, reporting_context.as_ref()), - }; let skmsg = encrypt_group_message( stores.sender_key_store, &sender_key_name, @@ -476,6 +479,9 @@ pub async fn prepare_group_stanza( ) .await?; + // Release before the chain-independent stanza build. + drop(chain_guard); + let skmsg_ciphertext = skmsg.into_serialized(); let mediatype = media_type_from_message(message); diff --git a/wacore/src/send/tests.rs b/wacore/src/send/tests.rs index 8673124bd..51267e54f 100644 --- a/wacore/src/send/tests.rs +++ b/wacore/src/send/tests.rs @@ -3479,6 +3479,176 @@ mod local_identity_change_on_send { } } + /// Prekey bundle with a valid signed-prekey signature (create_mock_bundle's + /// zeroed signature fails X3DH, so it can't establish a real session). + fn verifiable_bundle(rng: &mut rand::rngs::StdRng) -> PreKeyBundle { + let identity = IdentityKeyPair::generate(rng); + let spk = KeyPair::generate(rng); + let opk = KeyPair::generate(rng); + let sig = identity + .private_key() + .calculate_signature(&spk.public_key.serialize(), rng) + .unwrap(); + PreKeyBundle::new( + 1, + 1u32.into(), + Some((1u32.into(), opk.public_key)), + 1u32.into(), + spk.public_key, + sig.to_vec(), + *identity.identity_key(), + ) + .unwrap() + } + + fn raw_fanout_stores<'a>( + sender_key_store: &'a mut MemSenderKeyStore, + session_store: &'a mut MemSessionStore, + identity_store: &'a mut MemIdentityStore, + prekey_store: &'a mut UnusedPreKeyStore, + signed_prekey_store: &'a UnusedSignedPreKeyStore, + ) -> SignalStores<'a> { + SignalStores { + sender_key_store, + session_store, + identity_store, + prekey_store, + signed_prekey_store, + } + } + + /// Establish a real Signal session for each device directly on the stores + /// (the module's per-value MemSessionStore would lose sessions written + /// through the fan-out's clone_box, so setup must not go through spawns). + async fn stores_with_sessions(devices: &[Jid]) -> (MemSessionStore, MemIdentityStore) { + let mut rng = rand::make_rng::(); + let mut session_store = MemSessionStore::default(); + let mut identity_store = MemIdentityStore { + pair: IdentityKeyPair::generate(&mut rng), + known: HashMap::new(), + }; + for d in devices { + process_prekey_bundle( + &d.to_protocol_address(), + &mut session_store, + &mut identity_store, + &verifiable_bundle(&mut rng), + &mut rng, + UsePQRatchet::No, + ) + .await + .expect("session established"); + } + (session_store, identity_store) + } + + /// Happy path: the chunked fan-out returns a ciphertext for every device, + /// spanning more than one ENCRYPT_FANOUT_CONCURRENCY chunk. + #[tokio::test] + async fn encrypt_for_devices_with_sessions_raw_encrypts_every_device() { + let devices: Vec = (0..20u16) + .map(|i| Jid::pn_device(format!("1555000{i:04}"), 0)) + .collect(); + + let (mut session_store, mut identity_store) = stores_with_sessions(&devices).await; + let mut prekey_store = UnusedPreKeyStore; + let signed_prekey_store = UnusedSignedPreKeyStore; + let mut sender_key_store = MemSenderKeyStore::default(); + let mut stores = raw_fanout_stores( + &mut sender_key_store, + &mut session_store, + &mut identity_store, + &mut prekey_store, + &signed_prekey_store, + ); + let rt = TokioTestRuntime; + + let raw = encrypt_for_devices_with_sessions_raw( + &rt, + &mut stores, + &devices, + b"payload", + SessionPlan::assume_ready(devices.len()), + ) + .await + .expect("fan-out succeeds"); + + assert_eq!(raw.devices.len(), devices.len()); + assert!(raw.includes_prekey_message, "fresh sessions emit pkmsg"); + } + + /// Bad path: a device without a session is skipped while the rest still + /// encrypt. + #[tokio::test] + async fn encrypt_for_devices_with_sessions_raw_skips_sessionless_device() { + let device_ok = Jid::pn_device("15550000000", 0); + let device_bad = Jid::pn_device("15550000001", 0); + + let (mut session_store, mut identity_store) = + stores_with_sessions(std::slice::from_ref(&device_ok)).await; + let mut prekey_store = UnusedPreKeyStore; + let signed_prekey_store = UnusedSignedPreKeyStore; + let mut sender_key_store = MemSenderKeyStore::default(); + let mut stores = raw_fanout_stores( + &mut sender_key_store, + &mut session_store, + &mut identity_store, + &mut prekey_store, + &signed_prekey_store, + ); + let rt = TokioTestRuntime; + + let devices = vec![device_ok.clone(), device_bad]; + let raw = encrypt_for_devices_with_sessions_raw( + &rt, + &mut stores, + &devices, + b"payload", + SessionPlan::assume_ready(devices.len()), + ) + .await + .expect("fan-out succeeds despite the sessionless device"); + + assert_eq!(raw.devices.len(), 1); + assert_eq!(raw.devices[0].device_jid, device_ok); + } + + /// Regression: the chunked fan-out must return empty, not divide by zero, for + /// an empty device set (reachable on the cold force-SKDM path). + #[tokio::test] + async fn encrypt_for_devices_with_sessions_raw_handles_empty_device_set() { + let mut rng = rand::make_rng::(); + let mut session_store = MemSessionStore::default(); + let mut identity_store = MemIdentityStore { + pair: IdentityKeyPair::generate(&mut rng), + known: HashMap::new(), + }; + let mut prekey_store = UnusedPreKeyStore; + let signed_prekey_store = UnusedSignedPreKeyStore; + let mut sender_key_store = MemSenderKeyStore::default(); + let mut stores = SignalStores { + sender_key_store: &mut sender_key_store, + session_store: &mut session_store, + identity_store: &mut identity_store, + prekey_store: &mut prekey_store, + signed_prekey_store: &signed_prekey_store, + }; + let rt = TokioTestRuntime; + + let raw = encrypt_for_devices_with_sessions_raw( + &rt, + &mut stores, + &[], + b"x", + SessionPlan::assume_ready(0), + ) + .await + .expect("empty fan-out must succeed, not panic"); + + assert!(raw.devices.is_empty()); + assert!(!raw.includes_prekey_message); + } + /// The send path must report a replaced identity via the resolver when /// establishing a session whose bundle carries a new identity key for an /// address we already knew (peer reinstall). Mirrors WA Web saveIdentity