diff --git a/src/message.rs b/src/message.rs index d2a8f29ee..bf756ba1a 100644 --- a/src/message.rs +++ b/src/message.rs @@ -80,13 +80,17 @@ pub(crate) struct SessionBatchOutcome { had_failure: bool, } -#[derive(Clone, Copy, Debug, Default)] -struct MigrationDecryptOutcome { - decrypted: bool, - duplicate: bool, - dispatched: bool, - skdm_only: bool, - plaintext_failed: bool, +/// Outcome of a PN→LID migration retry decrypt. On `Decrypted` the plaintext +/// has already been pushed onto the caller's deferred-handling buffer (it runs +/// after the session lock is released), so no dispatch flags travel back here. +#[derive(Clone, Copy, Debug)] +enum MigrationDecryptResult { + /// Decrypted; plaintext buffered for post-lock handling. + Decrypted, + /// Server redelivered an already-processed message. + Duplicate, + /// Migration didn't apply or still failed; caller sends a retry receipt. + NotDecrypted, } #[derive(Clone, Copy, Debug, Default)] @@ -95,6 +99,14 @@ pub(crate) struct PlaintextHandleOutcome { skdm_only: bool, } +/// A decrypted session plaintext buffered during the locked decrypt loop and +/// handled after the per-sender session lock is released. +struct DeferredPlaintext { + enc_type: &'static str, + plaintext: Vec, + padding_version: u8, +} + fn should_process_skmsg_after_session( session_payloads_empty: bool, session_outcome: SessionBatchOutcome, diff --git a/src/message/receive.rs b/src/message/receive.rs index db605eb68..d7c09f938 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -593,19 +593,23 @@ impl Client { // the SignalProtocolStoreAdapter's per-session locks (prevents ratchet counter races). let signal_address = sender_encryption_jid.to_protocol_address(); - // `session_guard` is held across the entire batch but dropped around - // calls into `try_pn_to_lid_migration_decrypt` because that function's - // migration loop re-enters this same mutex (non-reentrant). + // `session_guard` is held across the decrypt loop (and released before the + // plaintext drain below); it's also dropped around calls into + // `try_pn_to_lid_migration_decrypt` because that function's migration loop + // re-enters this same mutex (non-reentrant). let session_mutex = self.session_lock_for(signal_address.as_str()).await; let mut session_guard: Option> = Some(session_mutex.lock_arc().await); - // Started after the lock so the histogram is crypto-only, not lock/queue wait. + // Started after the lock so the histogram excludes lock/queue wait. let _t = wacore::telemetry::timer(wacore::telemetry::DECRYPT_DURATION); let mut adapter = self.signal_adapter().await; let mut rng = rand::make_rng::(); let mut outcome = SessionBatchOutcome::default(); + // Buffer plaintexts to handle after the ratchet lock drops (see the drain + // below for why that's safe). + let mut deferred: Vec = Vec::new(); // Local identity-change detection fires once per batch: the first pkmsg // saves the new key (ReplacedExisting); the rest are NewOrUnchanged. let mut local_identity_reacted = false; @@ -730,35 +734,14 @@ impl Client { local_identity_reacted = true; self.react_to_local_identity_change(sender_encryption_jid); } - let padded_plaintext = decrypted.plaintext; - match self - .clone() - .handle_decrypted_plaintext( - enc_type, - &padded_plaintext, - padding_version, - info, - ) - .await - { - Ok(plaintext_outcome) => { - outcome.decrypted = true; - outcome.dispatched |= plaintext_outcome.dispatched; - outcome.skdm_only |= plaintext_outcome.skdm_only; - } - Err(e) => { - log::warn!( - "[msg:{}] Failed processing plaintext from {}: {e:?}", - info.id, - info.source.sender.observe() - ); - outcome.decrypted = true; - outcome.plaintext_failed = true; - outcome.had_failure = true; - outcome.undecryptable |= - self.handle_plaintext_failure(info, decrypt_fail_mode).await; - } - } + // Decrypt succeeded (the ratchet advanced); defer the + // plaintext handling until the lock is dropped. + outcome.decrypted = true; + deferred.push(DeferredPlaintext { + enc_type, + plaintext: decrypted.plaintext, + padding_version, + }); } Err(e) => { // Handle DuplicatedMessage: This is expected when messages are redelivered during reconnection @@ -865,34 +848,14 @@ impl Client { local_identity_reacted = true; self.react_to_local_identity_change(sender_encryption_jid); } - let padded_plaintext = decrypted.plaintext; - match self - .clone() - .handle_decrypted_plaintext( - enc_type, - &padded_plaintext, - padding_version, - info, - ) - .await - { - Ok(plaintext_outcome) => { - outcome.decrypted = true; - outcome.dispatched |= plaintext_outcome.dispatched; - outcome.skdm_only |= plaintext_outcome.skdm_only; - } - Err(e) => { - log::warn!( - "Failed processing plaintext after identity retry: {e:?}" - ); - outcome.decrypted = true; - outcome.plaintext_failed = true; - outcome.had_failure = true; - outcome.undecryptable |= self - .handle_plaintext_failure(info, decrypt_fail_mode) - .await; - } - } + // Decrypt succeeded after the identity retry; + // defer plaintext handling until the lock drops. + outcome.decrypted = true; + deferred.push(DeferredPlaintext { + enc_type, + plaintext: decrypted.plaintext, + padding_version, + }); } Err(retry_err) => { // Handle DuplicatedMessage in retry path: This commonly happens during reconnection @@ -913,7 +876,7 @@ impl Client { } else if matches!(retry_err, SignalProtocolError::InvalidPreKeyId) { // Session may exist under PN address after identity change - let migration_outcome = self + match self .try_pn_to_lid_migration_decrypt( sender_encryption_jid, &signal_address, @@ -925,39 +888,32 @@ impl Client { info, &session_mutex, &mut session_guard, + &mut deferred, ) - .await; - if migration_outcome.decrypted - || migration_outcome.duplicate - || migration_outcome.plaintext_failed + .await { - outcome.decrypted |= migration_outcome.decrypted; - outcome.duplicate |= migration_outcome.duplicate; - outcome.dispatched |= migration_outcome.dispatched; - outcome.skdm_only |= migration_outcome.skdm_only; - outcome.plaintext_failed |= - migration_outcome.plaintext_failed; - outcome.had_failure |= migration_outcome.plaintext_failed; - if migration_outcome.plaintext_failed { + MigrationDecryptResult::Decrypted => { + outcome.decrypted = true; + } + MigrationDecryptResult::Duplicate => { + outcome.duplicate = true; + } + MigrationDecryptResult::NotDecrypted => { + log::debug!( + "[msg:{}] InvalidPreKeyId after identity change for {}. \ + Sending retry receipt with fresh keys.", + info.id, + address + ); + outcome.had_failure = true; outcome.undecryptable |= self - .handle_plaintext_failure(info, decrypt_fail_mode) + .handle_decrypt_failure( + info, + RetryReason::InvalidKeyId, + decrypt_fail_mode, + ) .await; } - } else { - log::debug!( - "[msg:{}] InvalidPreKeyId after identity change for {}. \ - Sending retry receipt with fresh keys.", - info.id, - address - ); - outcome.had_failure = true; - outcome.undecryptable |= self - .handle_decrypt_failure( - info, - RetryReason::InvalidKeyId, - decrypt_fail_mode, - ) - .await; } } else { log::error!( @@ -1003,7 +959,7 @@ impl Client { } // Try PN→LID session migration before sending retry receipt if let SignalProtocolError::SessionNotFound(_) = e { - let migration_outcome = self + match self .try_pn_to_lid_migration_decrypt( sender_encryption_jid, &signal_address, @@ -1015,23 +971,19 @@ impl Client { info, &session_mutex, &mut session_guard, + &mut deferred, ) - .await; - if migration_outcome.decrypted - || migration_outcome.duplicate - || migration_outcome.plaintext_failed + .await { - outcome.decrypted |= migration_outcome.decrypted; - outcome.duplicate |= migration_outcome.duplicate; - outcome.dispatched |= migration_outcome.dispatched; - outcome.skdm_only |= migration_outcome.skdm_only; - outcome.plaintext_failed |= migration_outcome.plaintext_failed; - outcome.had_failure |= migration_outcome.plaintext_failed; - if migration_outcome.plaintext_failed { - outcome.undecryptable |= - self.handle_plaintext_failure(info, decrypt_fail_mode).await; + MigrationDecryptResult::Decrypted => { + outcome.decrypted = true; + continue; } - continue; + MigrationDecryptResult::Duplicate => { + outcome.duplicate = true; + continue; + } + MigrationDecryptResult::NotDecrypted => {} } debug!( @@ -1051,7 +1003,7 @@ impl Client { ) { // whatsmeow migrates PN sessions before decrypt; a fresh // LID record can otherwise shadow the sender's PN ratchet. - let migration_outcome = self + match self .try_pn_to_lid_migration_decrypt( sender_encryption_jid, &signal_address, @@ -1063,23 +1015,19 @@ impl Client { info, &session_mutex, &mut session_guard, + &mut deferred, ) - .await; - if migration_outcome.decrypted - || migration_outcome.duplicate - || migration_outcome.plaintext_failed + .await { - outcome.decrypted |= migration_outcome.decrypted; - outcome.duplicate |= migration_outcome.duplicate; - outcome.dispatched |= migration_outcome.dispatched; - outcome.skdm_only |= migration_outcome.skdm_only; - outcome.plaintext_failed |= migration_outcome.plaintext_failed; - outcome.had_failure |= migration_outcome.plaintext_failed; - if migration_outcome.plaintext_failed { - outcome.undecryptable |= - self.handle_plaintext_failure(info, decrypt_fail_mode).await; + MigrationDecryptResult::Decrypted => { + outcome.decrypted = true; + continue; } - continue; + MigrationDecryptResult::Duplicate => { + outcome.duplicate = true; + continue; + } + MigrationDecryptResult::NotDecrypted => {} } // WAWebMsgProcessingDecryptionHandler classifies both as @@ -1108,7 +1056,7 @@ impl Client { // session exists under a PN address (legacy migration). // Migrating lets Signal use the existing ratchet state // instead of looking up the consumed one-time prekey. - let migration_outcome = self + match self .try_pn_to_lid_migration_decrypt( sender_encryption_jid, &signal_address, @@ -1120,23 +1068,19 @@ impl Client { info, &session_mutex, &mut session_guard, + &mut deferred, ) - .await; - if migration_outcome.decrypted - || migration_outcome.duplicate - || migration_outcome.plaintext_failed + .await { - outcome.decrypted |= migration_outcome.decrypted; - outcome.duplicate |= migration_outcome.duplicate; - outcome.dispatched |= migration_outcome.dispatched; - outcome.skdm_only |= migration_outcome.skdm_only; - outcome.plaintext_failed |= migration_outcome.plaintext_failed; - outcome.had_failure |= migration_outcome.plaintext_failed; - if migration_outcome.plaintext_failed { - outcome.undecryptable |= - self.handle_plaintext_failure(info, decrypt_fail_mode).await; + MigrationDecryptResult::Decrypted => { + outcome.decrypted = true; + continue; } - continue; + MigrationDecryptResult::Duplicate => { + outcome.duplicate = true; + continue; + } + MigrationDecryptResult::NotDecrypted => {} } log::debug!( @@ -1182,6 +1126,43 @@ impl Client { } } } + + // Release the per-sender session lock before handling plaintext: + // handle_decrypted_plaintext never touches this session's ratchet (only the + // decrypt above does), so a concurrent same-sender stanza can decrypt while + // this one dispatches. Safe because the buffer drains before we return (SKDM + // still precedes PASS 2's group decrypt) and per-chat order is owned by the + // serial chat-lane worker, not this guard. Matches whatsmeow. + drop(session_guard); + + for DeferredPlaintext { + enc_type, + plaintext, + padding_version, + } in deferred + { + match self + .handle_decrypted_plaintext(enc_type, &plaintext, padding_version, info) + .await + { + Ok(plaintext_outcome) => { + outcome.dispatched |= plaintext_outcome.dispatched; + outcome.skdm_only |= plaintext_outcome.skdm_only; + } + Err(e) => { + log::warn!( + "[msg:{}] Failed processing plaintext from {}: {e:?}", + info.id, + info.source.sender.observe() + ); + outcome.plaintext_failed = true; + outcome.had_failure = true; + outcome.undecryptable |= + self.handle_plaintext_failure(info, decrypt_fail_mode).await; + } + } + } + outcome } @@ -1236,7 +1217,6 @@ impl Client { } if let Err(e) = self - .clone() .handle_decrypted_plaintext( "skmsg", &padded_plaintext, @@ -1379,7 +1359,7 @@ impl Client { #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.recv.handle_plaintext", level = "debug", skip_all, fields(chat = %info.source.chat.observe(), sender = %info.source.sender.observe(), msg_id = %info.id, enc_type = %enc_type), err(Debug)))] pub(crate) async fn handle_decrypted_plaintext( - self: Arc, + self: &Arc, enc_type: &str, padded_plaintext: &[u8], padding_version: u8, @@ -1529,9 +1509,10 @@ impl Client { } } - /// Attempt PN→LID session migration and retry decryption. - /// Returns whether decryption succeeded after migration and whether it - /// reached user dispatch. + /// Attempt PN→LID session migration and retry decryption. On success the + /// plaintext is pushed onto `deferred` for post-lock handling by the caller + /// (see `process_session_enc_batch`), so this returns only the decrypt + /// disposition. /// /// Manages the per-address session lock around the migration loop: /// drops the caller's guard (migration re-enters that mutex and @@ -1547,20 +1528,21 @@ impl Client { parsed_message: &wacore::libsignal::protocol::CiphertextMessage, adapter: &mut crate::store::signal_adapter::SignalProtocolStoreAdapter, rng: &mut rand::rngs::StdRng, - enc_type: &str, + enc_type: &'static str, padding_version: u8, info: &Arc, session_mutex: &Arc>, session_guard: &mut Option>, - ) -> MigrationDecryptOutcome { + deferred: &mut Vec, + ) -> MigrationDecryptResult { use wacore::libsignal::protocol::{UsePQRatchet, message_decrypt}; if !sender_jid.is_lid() { - return MigrationDecryptOutcome::default(); + return MigrationDecryptResult::NotDecrypted; } let Some(pn) = self.lid_pn_cache.get_phone_number(&sender_jid.user).await else { - return MigrationDecryptOutcome::default(); + return MigrationDecryptResult::NotDecrypted; }; // Release the address lock so the migration loop can acquire it for @@ -1569,8 +1551,10 @@ impl Client { let migrated = self .migrate_signal_sessions_on_lid_discovery(&pn, &sender_jid.user) .await; - // Re-acquire for the retry decrypt and hand the guard back to the - // caller for subsequent payloads in the batch. + // Re-acquire for the retry decrypt and hand the guard back to the caller + // for subsequent payloads in the batch. Every return below is past this + // reacquire, so the caller always gets the guard back as `Some`; losing + // that would silently drop same-sender serialization. *session_guard = Some(session_mutex.lock_arc().await); // Nothing moved namespaces, so the retry would hit the exact same @@ -1582,7 +1566,7 @@ impl Client { info.id, info.source.sender.observe() ); - return MigrationDecryptOutcome::default(); + return MigrationDecryptResult::NotDecrypted; } match message_decrypt( @@ -1612,47 +1596,28 @@ impl Client { .buffer_consumed_prekey(prekey_id, signal_address) .await; } - let padded_plaintext = decrypted.plaintext; - match self - .clone() - .handle_decrypted_plaintext(enc_type, &padded_plaintext, padding_version, info) - .await - { - Ok(plaintext_outcome) => MigrationDecryptOutcome { - decrypted: true, - dispatched: plaintext_outcome.dispatched, - skdm_only: plaintext_outcome.skdm_only, - ..Default::default() - }, - Err(e) => { - log::warn!( - "[msg:{}] Failed processing plaintext after migration: {e:?}", - info.id - ); - MigrationDecryptOutcome { - decrypted: true, - plaintext_failed: true, - ..Default::default() - } - } - } + // Buffer for post-lock handling, keeping the same ordering as the + // batch's main path. + deferred.push(DeferredPlaintext { + enc_type, + plaintext: decrypted.plaintext, + padding_version, + }); + MigrationDecryptResult::Decrypted } Err(SignalProtocolError::DuplicatedMessage(chain, counter)) => { log::debug!( "[msg:{}] Already processed (chain {chain}, counter {counter}) after migration", info.id ); - MigrationDecryptOutcome { - duplicate: true, - ..Default::default() - } + MigrationDecryptResult::Duplicate } Err(retry_err) => { log::warn!( "[msg:{}] Decryption still failed after PN→LID migration: {retry_err:?}", info.id ); - MigrationDecryptOutcome::default() + MigrationDecryptResult::NotDecrypted } } } diff --git a/src/message/tests.rs b/src/message/tests.rs index 2f928e06c..29f417283 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -7431,7 +7431,6 @@ async fn app_state_sync_key_share_honored_only_from_self() { create_test_message_info("5510000@s.whatsapp.net", "AKS1", "5510000@s.whatsapp.net"); info.source.is_from_me = false; client - .clone() .handle_decrypted_plaintext("msg", &padded, 2, &Arc::new(info)) .await .unwrap(); @@ -7448,7 +7447,6 @@ async fn app_state_sync_key_share_honored_only_from_self() { ); info.source.is_from_me = true; client - .clone() .handle_decrypted_plaintext("msg", &padded, 2, &Arc::new(info)) .await .unwrap(); @@ -7497,7 +7495,6 @@ async fn lid_migration_mapping_sync_honored_only_from_self() { create_test_message_info("5510000@s.whatsapp.net", "LMS1", "5510000@s.whatsapp.net"); info.source.is_from_me = false; client - .clone() .handle_decrypted_plaintext("msg", &padded, 2, &Arc::new(info)) .await .unwrap(); @@ -7514,7 +7511,6 @@ async fn lid_migration_mapping_sync_honored_only_from_self() { ); info.source.is_from_me = true; client - .clone() .handle_decrypted_plaintext("msg", &padded, 2, &Arc::new(info)) .await .unwrap(); @@ -10259,3 +10255,154 @@ async fn addon_decrypts_right_after_capture_without_flush() { "an add-on right after the capture must find the secret" ); } + +// =========================================================================== +// Empirical steady-state session-decrypt benchmark (ignored; run explicitly). +// +// cargo test -p whatsapp-rust --release bench_session_decrypt_throughput \ +// -- --ignored --nocapture +// +// Wall-clock gives throughput; for deterministic CPU-instructions and +// per-callsite allocations run the SAME test binary under valgrind +// (callgrind / dhat) with a small BENCH_N. +// =========================================================================== + +async fn bench_feed( + client: &Arc, + peer: &Jid, + enc_type: &'static str, + bytes: Vec, +) -> bool { + let enc_node = NodeBuilder::new("enc") + .attr("type", enc_type) + .bytes(bytes) + .build(); + let enc_ref = enc_node.as_node_ref(); + let payloads: Vec = vec![EncPayload::from_node_ref(&enc_ref).unwrap()]; + let info = Arc::new(MessageInfo { + source: crate::types::message::MessageSource { + sender: peer.clone(), + chat: peer.clone(), + ..Default::default() + }, + ..Default::default() + }); + let outcome = client + .clone() + .process_session_enc_batch( + &payloads, + &info, + peer, + crate::types::events::DecryptFailMode::Show, + ) + .await; + outcome.decrypted && !outcome.undecryptable +} + +fn bench_vm_kb(key: &str) -> u64 { + std::fs::read_to_string("/proc/self/status") + .ok() + .and_then(|s| { + s.lines() + .find(|l| l.starts_with(key)) + .and_then(|l| l.split_whitespace().nth(1)) + .and_then(|v| v.parse().ok()) + }) + .unwrap_or(0) +} + +#[tokio::test(flavor = "current_thread")] +#[ignore = "benchmark: run explicitly with --ignored --nocapture"] +async fn bench_session_decrypt_throughput() { + let n: usize = std::env::var("BENCH_N") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(20_000); + let warmup: usize = std::env::var("BENCH_WARMUP") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(500); + + let client = crate::test_utils::create_test_client_with_name("bench_decrypt").await; + ensure_bob_paired(&client).await; + let (bundle, _bob) = bobs_prekey_bundle(&client).await; + let bob_addr = { + let s = client.persistence_manager.get_device_snapshot(); + s.lid + .as_ref() + .or(s.pn.as_ref()) + .expect("own jid") + .to_protocol_address() + }; + + let mut alice = AlicePeer::new("10000000000001:0@s.whatsapp.net").await; + alice.install_bob_session(&bob_addr, &bundle).await; + + // One real pkmsg decrypt to establish Bob's reciprocal session. + let pkmsg = alice.encrypt_text(&bob_addr, "establish").await; + let pk_bytes = match &pkmsg { + CiphertextMessage::PreKeySignalMessage(m) => m.serialized().to_vec(), + _ => panic!("first message must be pkmsg"), + }; + assert!( + bench_feed(&client, &alice.jid, "pkmsg", pk_bytes).await, + "establish decrypt must succeed" + ); + + // Steady state: force plain SignalMessages (the common case) from here on. + { + let record = alice + .sessions + .0 + .get_mut(&bob_addr) + .expect("alice session for bob"); + if let Some(state) = record.session_state_mut() { + state.clear_unacknowledged_pre_key_message(); + } + } + + // Pre-encrypt warmup + N steady-state `msg` ciphertexts (Alice ratchets). + let total = warmup + n; + let mut msgs: Vec> = Vec::with_capacity(total); + for i in 0..total { + let ct = alice + .encrypt_text(&bob_addr, &format!("benchmark payload {i}")) + .await; + match ct { + CiphertextMessage::SignalMessage(m) => msgs.push(m.serialized().to_vec()), + other => panic!( + "expected steady-state msg, got {:?}", + std::mem::discriminant(&other) + ), + } + } + + for b in msgs.iter().take(warmup) { + assert!( + bench_feed(&client, &alice.jid, "msg", b.clone()).await, + "warmup decrypt" + ); + } + + let rss0 = bench_vm_kb("VmRSS"); + let t0 = wacore::time::Instant::now(); + let mut ok = 0usize; + for b in msgs.iter().skip(warmup) { + if bench_feed(&client, &alice.jid, "msg", b.clone()).await { + ok += 1; + } + } + let elapsed = t0.elapsed(); + let rss1 = bench_vm_kb("VmRSS"); + let hwm = bench_vm_kb("VmHWM"); + + assert_eq!(ok, n, "all steady-state messages must decrypt"); + let per_ns = elapsed.as_nanos() as f64 / n as f64; + let thrpt = n as f64 / elapsed.as_secs_f64(); + println!( + "\n=== BENCH session decrypt (msg) ===\n\ + n={n} elapsed={elapsed:.2?} per_msg={per_ns:.0}ns throughput={thrpt:.0} msg/s\n\ + rss_start={rss0}KB rss_end={rss1}KB rss_delta={}KB vm_hwm={hwm}KB\n", + rss1.saturating_sub(rss0) + ); +}