diff --git a/src/history_sync.rs b/src/history_sync.rs index 550757003..265c9de97 100644 --- a/src/history_sync.rs +++ b/src/history_sync.rs @@ -164,11 +164,7 @@ impl HistorySecretSeedCollector { Err(_) => return, }; for sender in senders.into_iter().flatten() { - let sender_id = if sender.is_same_chat_as(chat) { - Arc::clone(&chat_id) - } else { - Arc::from(sender.to_non_ad_string()) - }; + let sender_id = MsgSecretEntry::sender_id_for(chat, &chat_id, &sender); self.entries.push(MsgSecretEntry { chat: Arc::clone(&chat_id), sender: sender_id, diff --git a/src/message/msg_secret.rs b/src/message/msg_secret.rs index 1d4916f77..78a784ec1 100644 --- a/src/message/msg_secret.rs +++ b/src/message/msg_secret.rs @@ -31,34 +31,43 @@ impl Client { let class = wacore::msg_secret::classify(msg, chat_is_bot); let message_ts = u64::try_from(info.timestamp.timestamp()).ok(); - // Build both aliases (primary, plus the bot-DM LID key) and write them - // in one batch so a partial write can't leave only one stored. - let mut entries = Vec::with_capacity(2); - if let Some(entry) = self.build_msg_secret_entry( + let primary = self.build_msg_secret_entry( &info.source.chat, &info.source.sender, &info.id, secret_bytes, class, message_ts, - ) { - entries.push(entry); - } + ); + // The bot-DM LID key is the only alias a capture ever adds, so a plain + // chat writes exactly one row. Deciding that before touching a Vec keeps + // the common capture off the batch allocation entirely. + let mut bot_alias = None; if chat_is_bot && let Some(sender) = self.dm_sender_identity_for(&info.source.chat).await && sender.to_non_ad() != info.source.sender.to_non_ad() - && let Some(entry) = self.build_msg_secret_entry( + { + bot_alias = self.build_msg_secret_entry( &info.source.chat, &sender, &info.id, secret_bytes, class, message_ts, - ) - { - entries.push(entry); + ); + } + + match (primary, bot_alias) { + // Both aliases go out in one batch so a partial write can't leave + // only one of them stored. + (Some(primary), Some(alias)) => { + self.persist_msg_secret_entries(vec![primary, alias]).await + } + (Some(entry), None) | (None, Some(entry)) => { + self.msg_secret_buffer.queue_one(entry).await + } + (None, None) => {} } - self.persist_msg_secret_entries(entries).await; } /// Build one retention entry, applying the policy gates and computing the @@ -92,14 +101,14 @@ impl Client { message_ts, wacore::time::now_secs(), ); - Some(wacore::store::traits::MsgSecretEntry { - chat: chat.to_non_ad_string().into(), - sender: sender.to_non_ad_string().into(), - msg_id: msg_id.into(), - secret: *secret, + Some(wacore::store::traits::MsgSecretEntry::new( + chat, + sender, + msg_id, + *secret, expires_at, - message_ts: message_ts.and_then(|t| i64::try_from(t).ok()).unwrap_or(0), - }) + message_ts.and_then(|t| i64::try_from(t).ok()).unwrap_or(0), + )) } /// Queue a batch of secret aliases on the write-behind buffer: immediately diff --git a/src/message/receive.rs b/src/message/receive.rs index b1943f92a..f014027b7 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -22,6 +22,21 @@ impl ParsedSessionMessage { } } +/// Append to one of the three per-kind enc buckets, allocating that bucket only +/// once it actually receives a payload. +/// +/// A stanza's enc nodes are overwhelmingly all one kind, so reserving all three +/// buckets up front spent two allocations per message on buffers that stayed +/// empty for their whole lifetime. Reserving the stanza's full enc count on the +/// first push keeps a mixed stanza at one allocation per non-empty bucket, +/// exactly as before. +fn push_enc_payload(bucket: &mut Vec, stanza_enc_count: usize, payload: EncPayload) { + if bucket.capacity() == 0 { + bucket.reserve_exact(stanza_enc_count); + } + bucket.push(payload); +} + async fn decrypt_session_message( message: &mut ParsedSessionMessage, signal_address: &wacore::libsignal::protocol::ProtocolAddress, @@ -222,9 +237,9 @@ impl Client { return None; } - let mut session_payloads = Vec::with_capacity(all_enc_nodes.len()); - let mut group_payloads = Vec::with_capacity(all_enc_nodes.len()); - let mut bot_payloads = Vec::with_capacity(all_enc_nodes.len()); + let mut session_payloads = Vec::new(); + let mut group_payloads = Vec::new(); + let mut bot_payloads = Vec::new(); let mut max_sender_retry_count = 0; let mut has_hide_fail = false; let mut had_unknown_enc = false; @@ -301,13 +316,14 @@ impl Client { } }; - if payload.enc_type.is_bot_secret() { - bot_payloads.push(payload); + let bucket = if payload.enc_type.is_bot_secret() { + &mut bot_payloads } else if payload.enc_type.is_session() { - session_payloads.push(payload); + &mut session_payloads } else { - group_payloads.push(payload); - } + &mut group_payloads + }; + push_enc_payload(bucket, all_enc_nodes.len(), payload); } // WA Web diagnostic: validate skmsg is not first in multi-enc messages. @@ -1750,6 +1766,76 @@ impl Client { } } +#[cfg(test)] +mod enc_bucket_tests { + use super::push_enc_payload; + use crate::message::EncPayload; + use wacore::message_processing::EncType; + + fn payload(enc_type: EncType) -> EncPayload { + EncPayload { + ciphertext: bytes::Bytes::from_static(b"ct"), + enc_type, + padding_version: 2, + } + } + + /// The DM shape: one enc node, one bucket used, two buckets untouched. The + /// used bucket must be sized to the stanza exactly (not to `Vec`'s default + /// growth step) and the unused ones must own no buffer at all. + #[test] + fn a_one_enc_stanza_sizes_one_bucket_exactly_and_leaves_the_rest_empty() { + let mut used: Vec = Vec::new(); + let unused: Vec = Vec::new(); + + push_enc_payload(&mut used, 1, payload(EncType::Message)); + + assert_eq!(used.len(), 1); + assert_eq!( + used.capacity(), + 1, + "the bucket must be reserved to the stanza's enc count, not grown" + ); + assert_eq!(unused.capacity(), 0, "an empty bucket must own no buffer"); + assert!(unused.is_empty()); + } + + /// The first push must reserve the stanza's whole enc count, so a multi-enc + /// stanza still pays exactly one allocation for the bucket it fills. Eight + /// is past `Vec`'s own first growth step, so plain pushes would reallocate. + #[test] + fn the_first_push_reserves_the_whole_stanza() { + const ENC_COUNT: usize = 8; + let mut bucket: Vec = Vec::new(); + push_enc_payload(&mut bucket, ENC_COUNT, payload(EncType::Message)); + let reserved = bucket.as_ptr(); + assert_eq!(bucket.capacity(), ENC_COUNT); + + for _ in 1..ENC_COUNT { + push_enc_payload(&mut bucket, ENC_COUNT, payload(EncType::PreKeyMessage)); + } + assert_eq!(bucket.len(), ENC_COUNT); + assert_eq!( + bucket.as_ptr(), + reserved, + "filling up to the stanza's enc count must not reallocate" + ); + assert_eq!(bucket[0].enc_type, EncType::Message); + assert_eq!(bucket[ENC_COUNT - 1].enc_type, EncType::PreKeyMessage); + } + + /// A degenerate count must not make the helper skip its reservation and + /// leave the bucket re-reserving on every later push. + #[test] + fn a_zero_count_still_stores_the_payload() { + let mut bucket: Vec = Vec::new(); + push_enc_payload(&mut bucket, 0, payload(EncType::SenderKey)); + assert_eq!(bucket.len(), 1); + push_enc_payload(&mut bucket, 0, payload(EncType::SenderKey)); + assert_eq!(bucket.len(), 2); + } +} + #[cfg(test)] mod tests { use crate::test_utils::create_test_client_with_failing_http; diff --git a/src/msg_secret_buffer.rs b/src/msg_secret_buffer.rs index 204507cf6..fccfe71f1 100644 --- a/src/msg_secret_buffer.rs +++ b/src/msg_secret_buffer.rs @@ -616,6 +616,31 @@ mod tests { buf.wait_flushed().await; } + /// Building a row costs one allocation per *distinct* identifier and + /// nothing else. The naive spelling (`to_non_ad_string().into()` per JID) + /// cost five for a direct message; this is the per-message saving on both + /// the inbound capture and the outbound persist. + #[test] + fn entry_construction_allocates_once_per_distinct_identifier() { + let chat: wacore_binary::Jid = "5511987650001@s.whatsapp.net".parse().unwrap(); + let peer_device: wacore_binary::Jid = "5511987650001:33@s.whatsapp.net".parse().unwrap(); + let me: wacore_binary::Jid = "5511987650002@s.whatsapp.net".parse().unwrap(); + let secret = [0u8; wacore::reporting_token::MESSAGE_SECRET_SIZE]; + + // A DM: chat and sender are the same user, so the row holds two + // allocations (the shared identifier and the message id). + let dm = crate::test_alloc::min_allocs(2, || { + MsgSecretEntry::new(&chat, &peer_device, "3EB0AABBCCDDEEFF0011", secret, 0, 0) + }); + assert_eq!(dm, 2, "a direct-message row: shared identifier + msg id"); + + // Outbound: the sender is us, a different user, so it needs its own. + let outbound = crate::test_alloc::min_allocs(3, || { + MsgSecretEntry::new(&chat, &me, "3EB0AABBCCDDEEFF0011", secret, 0, 0) + }); + assert_eq!(outbound, 3, "distinct users cost one identifier each"); + } + /// Cloning a buffered entry must stay allocation-free: identifiers share /// their Arc allocations and the protocol-sized secret lives inline. #[test] diff --git a/src/send/mod.rs b/src/send/mod.rs index 198cb2217..5c423f33b 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -2432,14 +2432,9 @@ impl Client { u64::try_from(now).ok(), now, ); - let entry = wacore::store::traits::MsgSecretEntry { - chat: chat.to_non_ad_string().into(), - sender: sender.to_non_ad_string().into(), - msg_id: msg_id.into(), - secret: *secret, - expires_at, - message_ts: now, - }; + let entry = wacore::store::traits::MsgSecretEntry::new( + chat, sender, msg_id, *secret, expires_at, now, + ); // Same write-behind buffer as inbound captures: visible immediately, // flushed off the send path (msmsg replies read buffer-first). self.msg_secret_buffer.queue_one(entry).await; diff --git a/wacore/binary/src/jid.rs b/wacore/binary/src/jid.rs index 72a253405..04b4d59aa 100644 --- a/wacore/binary/src/jid.rs +++ b/wacore/binary/src/jid.rs @@ -659,6 +659,20 @@ impl Jid { buf } + /// [`Self::to_non_ad_string`] as a shareable `Arc`, in exactly one + /// allocation. Going through the `String` first costs two — the buffer, then + /// the `Arc` its bytes are copied into — and the message-secret rows + /// build two of these per message. + pub fn to_non_ad_arc_str(&self) -> std::sync::Arc { + let mut writer = JidStackWriter::new(); + if write_jid_fallible(&mut writer, &self.user, self.server, 0, 0).is_ok() { + return std::sync::Arc::from(writer.as_str()); + } + // A user part too long for the stack buffer (never seen on the wire) + // still renders, just back through the heap. + std::sync::Arc::from(self.to_non_ad_string()) + } + /// Check if this JID matches the user or their LID. /// Useful for checking if a participant is "us" in group messages. #[inline] @@ -1976,6 +1990,53 @@ mod tests { } } + /// The stack-buffered `Arc` form must render exactly what the `String` + /// form does, including for inputs that overflow the stack buffer and fall + /// back to the heap, and for multibyte user parts (the buffer is bounded in + /// bytes, and a split fragment would be invalid UTF-8). + #[test] + fn to_non_ad_arc_str_matches_to_non_ad_string() { + let long_user = "9".repeat(80); + let multibyte_user = "ẞünïcodé-ñ".repeat(3); + let owned = [ + format!("{long_user}:12@s.whatsapp.net"), + format!("{multibyte_user}@g.us"), + format!("{multibyte_user}@s.whatsapp.net"), + ]; + let cases = [ + "1234567890:33@s.whatsapp.net", + "1234567890@s.whatsapp.net", + "100000012345678:25@lid", + "867051314767696:0@bot", + // Nonzero agent on a bot JID: both forms drop the agent (matching + // whatsmeow's ToNonAD), while `is_same_chat_as` still treats it as + // identity-significant. Pinning it here keeps that asymmetry + // deliberate rather than something a later edit can erase quietly. + "867051314767696.5:10@bot", + "120363021033254949@g.us", + "status@broadcast", + ] + .into_iter() + .chain(owned.iter().map(String::as_str)); + + for s in cases { + let jid: Jid = s.parse().unwrap_or_else(|e| panic!("parse {s}: {e}")); + assert_eq!( + &*jid.to_non_ad_arc_str(), + jid.to_non_ad_string().as_str(), + "mismatch for {s}" + ); + } + + // A default (empty-user) JID has no wire form to render but must still + // agree with the String path rather than panic in the stack writer. + let empty = Jid::default(); + assert_eq!( + &*empty.to_non_ad_arc_str(), + empty.to_non_ad_string().as_str() + ); + } + #[test] fn test_into_non_ad_matches_to_non_ad() { // into_non_ad (consuming) must produce a JID identical to to_non_ad (cloning). diff --git a/wacore/binary/tests/jid_non_ad_arc_alloc.rs b/wacore/binary/tests/jid_non_ad_arc_alloc.rs new file mode 100644 index 000000000..30e1df717 --- /dev/null +++ b/wacore/binary/tests/jid_non_ad_arc_alloc.rs @@ -0,0 +1,68 @@ +//! Locks the one-allocation property of [`Jid::to_non_ad_arc_str`]. The +//! message-secret rows build two of these per message, and the obvious +//! `to_non_ad_string().into()` spelling costs two allocations each: the +//! intermediate `String`, then the `Arc` its bytes are copied into. +//! +//! Single test fn on purpose: the counting allocator is process-global, so a +//! concurrently running sibling test would bleed its allocations into the +//! measurement. + +// Host-only allocation-count harness; std's 64-bit atomic is fine (never built +// for embedded targets). +#![allow(clippy::disallowed_types)] + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use wacore_binary::Jid; + +struct CountingAlloc; +static ALLOCS: AtomicU64 = AtomicU64::new(0); + +unsafe impl GlobalAlloc for CountingAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + ALLOCS.fetch_add(1, Ordering::Relaxed); + unsafe { System.alloc(layout) } + } + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) } + } +} + +#[global_allocator] +static GLOBAL: CountingAlloc = CountingAlloc; + +/// Smallest allocation delta over many windows: the counter is process-global, +/// so harness threads can bleed into any single window, but a genuine extra +/// allocation is paid on every iteration and its minimum can never drop. +fn min_allocs(iterations: u32, mut op: impl FnMut() -> T) -> u64 { + let mut min = u64::MAX; + for _ in 0..iterations { + let before = ALLOCS.load(Ordering::Relaxed); + let value = std::hint::black_box(op()); + let after = ALLOCS.load(Ordering::Relaxed); + drop(value); + min = min.min(after - before); + } + min +} + +#[test] +fn to_non_ad_arc_str_costs_one_allocation() { + // A device-qualified PN: 28 bytes of output, past CompactString's inline + // limit, so nothing here is inline-able and every allocation is real. + let jid: Jid = "5511987650001:33@s.whatsapp.net".parse().expect("parse"); + assert_eq!(&*jid.to_non_ad_arc_str(), "5511987650001@s.whatsapp.net"); + + let direct = min_allocs(200, || jid.to_non_ad_arc_str()); + assert_eq!(direct, 1, "the Arc itself is the only allocation"); + + // The spelling this replaced, measured the same way, so the test states the + // saving rather than asserting a bare number that could drift with the JID. + let via_string = min_allocs(200, || Arc::::from(jid.to_non_ad_string())); + assert_eq!( + via_string, + direct + 1, + "going through String must cost exactly the one extra allocation this removes" + ); +} diff --git a/wacore/src/messages.rs b/wacore/src/messages.rs index 431d35784..bbcee43d2 100644 --- a/wacore/src/messages.rs +++ b/wacore/src/messages.rs @@ -1125,9 +1125,12 @@ pub fn parse_message_info( .unwrap_or_default(), timestamp: crate::time::from_secs_or_now(attrs.unix_time("t")), category, + // Parse from the borrowed attribute: `From` immediately + // re-borrows it, so materializing a String first only buys a discarded + // allocation for every known variant. edit: attrs .optional_string("edit") - .map(|s| EditAttribute::from(s.to_string())) + .map(|s| EditAttribute::from(s.as_ref())) .unwrap_or_default(), is_offline, server_timestamp_us, diff --git a/wacore/src/reporting_token.rs b/wacore/src/reporting_token.rs index eac20383d..03ce08162 100644 --- a/wacore/src/reporting_token.rs +++ b/wacore/src/reporting_token.rs @@ -758,8 +758,11 @@ pub fn generate_reporting_token_from_encoded( tracing::instrument(name = "wa.send.reporting_node", level = "debug", skip_all) )] pub fn build_reporting_node(result: &ReportingTokenResult) -> Node { + // The integer goes in as an integer: `NodeValue`'s numeric conversion + // formats through `itoa` into an inline `CompactString`, while a + // `to_string()` first would heap-allocate a one-byte String per message. let token_node = NodeBuilder::new("reporting_token") - .attrs([("v", result.version.to_string())]) + .attr("v", result.version) .bytes(result.reporting_token.to_vec()) .build(); @@ -1252,6 +1255,17 @@ mod tests { ); } } + + // `v` carries the result's version, not a baked-in constant: the + // integer goes through `NodeValue`'s itoa conversion, so a version + // bump must reach the wire without touching this builder. + let bumped = ReportingTokenResult { + version: 7, + ..result + }; + let node = build_reporting_node(&bumped); + let token_node = node.get_children_by_tag("reporting_token").next().unwrap(); + assert!(token_node.attrs.get("v").is_some_and(|v| v == "7")); } #[test] diff --git a/wacore/src/send/dm.rs b/wacore/src/send/dm.rs index 3ab74fdb6..365cce242 100644 --- a/wacore/src/send/dm.rs +++ b/wacore/src/send/dm.rs @@ -240,11 +240,16 @@ pub async fn prepare_dm_stanza( )); } - let mut message_content_nodes = vec![ + // Sized for everything that can follow ``: the optional + // ``, the optional ``, and the caller's extra + // nodes. `vec![one]` reserves exactly one slot, so each later push + // reallocated and memcpy'd the whole (large) `Node` values. + let mut message_content_nodes = Vec::with_capacity(3 + extra_stanza_nodes.len()); + message_content_nodes.push( NodeBuilder::new("participants") .children(participant_nodes) .build(), - ]; + ); // DM stays lenient when pkmsg lacks an account (no pre-flight here): map the // helper's error back to omission so the wire shape is unchanged. diff --git a/wacore/src/send/encrypt.rs b/wacore/src/send/encrypt.rs index e56968bba..37bca6114 100644 --- a/wacore/src/send/encrypt.rs +++ b/wacore/src/send/encrypt.rs @@ -377,6 +377,12 @@ pub async fn encrypt_for_devices( /// [`ensure_sessions_for_devices`]; consumed by /// [`encrypt_for_devices_with_sessions`] over the same `devices` slice. pub struct SessionPlan { + /// Device count the plan was built for. Kept alongside the (possibly empty) + /// override map so the "same slice" invariant is still checkable now that an + /// override-free plan carries no vector at all. + device_count: usize, + /// Empty means "no device is overridden"; otherwise one slot per device. + /// See [`record_encryption_override`]. encryption_overrides: Vec>, pub had_unregistered_device: bool, first_error: Option, @@ -389,13 +395,38 @@ impl SessionPlan { /// and must not touch the network during the encrypt fan-out. pub fn assume_ready(device_count: usize) -> Self { Self { - encryption_overrides: vec![None; device_count], + device_count, + encryption_overrides: Vec::new(), had_unregistered_device: false, first_error: None, } } } +/// The LID address recorded for `index`, or `None` when that device encrypts +/// against its own JID. Indexing tolerates the empty (no-override) map, which +/// is what a warm send carries. +fn encryption_override_at(overrides: &[Option], index: usize) -> Option<&Jid> { + overrides.get(index).and_then(Option::as_ref) +} + +/// Record a per-index LID override, materializing the map on its first entry. +/// +/// A steady-state send overrides nothing, so the all-`None` vector it would +/// otherwise allocate (once per encrypt fan-out, twice per DM that also has +/// companion devices) never exists. +fn record_encryption_override( + overrides: &mut Vec>, + device_count: usize, + index: usize, + jid: Jid, +) { + if overrides.is_empty() { + overrides.resize(device_count, None); + } + overrides[index] = Some(jid); +} + /// 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 @@ -411,10 +442,12 @@ pub async fn ensure_sessions_for_devices( // None = use devices[i] as-is; Some(jid) = use this LID-upgraded version. // The Vec replaces a HashMap<&Jid, Jid> that paid hash + alloc per insert // and per get (~666 of each on a large group). Plain Vec> is - // direct indexing and contiguous memory. - let mut encryption_overrides: Vec> = vec![None; devices.len()]; + // direct indexing and contiguous memory. Both vectors stay unallocated + // until something is actually recorded in them, which on a warm send is + // never. + let mut encryption_overrides: Vec> = Vec::new(); // Indices into `devices` for those needing prekey fetch. - let mut indices_needing_prekeys: Vec = Vec::with_capacity(devices.len()); + let mut indices_needing_prekeys: Vec = Vec::new(); let mut had_406 = false; let mut first_error = None; @@ -437,7 +470,7 @@ pub async fn ensure_sessions_for_devices( lid_jid.observe(), device_jid.observe() ); - encryption_overrides[idx] = Some(lid_jid); + record_encryption_override(&mut encryption_overrides, devices.len(), idx, lid_jid); continue; } } @@ -459,7 +492,7 @@ pub async fn ensure_sessions_for_devices( lid_jid.observe(), device_jid.observe() ); - encryption_overrides[idx] = Some(lid_jid); + record_encryption_override(&mut encryption_overrides, devices.len(), idx, lid_jid); } indices_needing_prekeys.push(idx); } @@ -516,8 +549,8 @@ pub async fn ensure_sessions_for_devices( let make_session_task = |spawn_idx: usize| { let idx = indices_needing_prekeys[spawn_idx]; let device_jid = devices[idx].clone(); - let mut encryption_jid = encryption_overrides[idx] - .clone() + let mut encryption_jid = encryption_override_at(&encryption_overrides, idx) + .cloned() .unwrap_or_else(|| device_jid.clone()); // Normalize agent to 0 for LID JIDs to match how pre-key bundles are stored. @@ -607,6 +640,7 @@ pub async fn ensure_sessions_for_devices( } Ok(SessionPlan { + device_count: devices.len(), encryption_overrides, had_unregistered_device: had_406, first_error, @@ -722,11 +756,12 @@ async fn encrypt_for_devices_with_sessions_raw_detailed( plan: SessionPlan, ) -> Result { debug_assert_eq!( - plan.encryption_overrides.len(), + plan.device_count, devices.len(), "SessionPlan built for a different device list" ); let SessionPlan { + device_count: _, encryption_overrides, had_unregistered_device, mut first_error, @@ -744,9 +779,7 @@ async fn encrypt_for_devices_with_sessions_raw_detailed( // a FuturesUnordered, and two store clones), with no parallelism to gain. // Encrypt inline. let device_jid = devices[0].clone(); - let addr = encryption_overrides - .first() - .and_then(|o| o.as_ref()) + let addr = encryption_override_at(&encryption_overrides, 0) .unwrap_or(&devices[0]) .to_protocol_address(); let res = encrypt_one_device( @@ -782,9 +815,7 @@ async fn encrypt_for_devices_with_sessions_raw_detailed( // 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()) + let addr = encryption_override_at(&encryption_overrides, idx) .unwrap_or(&devices[idx]) .to_protocol_address(); (addr, devices[idx].clone()) @@ -849,3 +880,80 @@ async fn encrypt_for_devices_with_sessions_raw_detailed( first_error, }) } + +#[cfg(test)] +mod encryption_override_tests { + use super::{SessionPlan, encryption_override_at, record_encryption_override}; + use wacore_binary::Jid; + + fn lid(user: &str, device: u16) -> Jid { + Jid::lid_device(user.to_owned(), device) + } + + /// The steady state: nothing is overridden, so the per-device map is never + /// allocated and every lookup still answers "use the device's own JID". + #[test] + fn an_empty_map_answers_every_index_without_allocating() { + let overrides: Vec> = Vec::new(); + assert_eq!(overrides.capacity(), 0, "no override must mean no buffer"); + for index in [0, 1, 7, usize::MAX] { + assert!(encryption_override_at(&overrides, index).is_none()); + } + + let plan = SessionPlan::assume_ready(4); + assert!( + plan.encryption_overrides.is_empty(), + "a plan that overrides nothing must carry no override buffer" + ); + assert_eq!(plan.device_count, 4, "the slice length is still recorded"); + } + + /// The first recorded override materializes the full map, so later indices + /// stay addressable and earlier ones stay `None`. + #[test] + fn recording_materializes_the_whole_map_once() { + let mut overrides: Vec> = Vec::new(); + record_encryption_override(&mut overrides, 3, 2, lid("100000000000001", 5)); + assert_eq!(overrides.len(), 3); + assert!(encryption_override_at(&overrides, 0).is_none()); + assert!(encryption_override_at(&overrides, 1).is_none()); + assert_eq!( + encryption_override_at(&overrides, 2), + Some(&lid("100000000000001", 5)) + ); + + // A second record must not resize again, nor clear the first. + record_encryption_override(&mut overrides, 3, 0, lid("100000000000002", 0)); + assert_eq!(overrides.len(), 3); + assert_eq!( + encryption_override_at(&overrides, 0), + Some(&lid("100000000000002", 0)) + ); + assert_eq!( + encryption_override_at(&overrides, 2), + Some(&lid("100000000000001", 5)) + ); + + // Overwriting an index replaces it rather than appending. + record_encryption_override(&mut overrides, 3, 2, lid("100000000000003", 1)); + assert_eq!(overrides.len(), 3); + assert_eq!( + encryption_override_at(&overrides, 2), + Some(&lid("100000000000003", 1)) + ); + } + + /// A single-device fan-out is the DM hot path, and it reads index 0 off a + /// map that may not exist. + #[test] + fn a_single_device_plan_records_and_reads_index_zero() { + let mut overrides: Vec> = Vec::new(); + assert!(encryption_override_at(&overrides, 0).is_none()); + record_encryption_override(&mut overrides, 1, 0, lid("100000000000009", 33)); + assert_eq!( + encryption_override_at(&overrides, 0), + Some(&lid("100000000000009", 33)) + ); + assert!(encryption_override_at(&overrides, 1).is_none()); + } +} diff --git a/wacore/src/store/traits.rs b/wacore/src/store/traits.rs index 78b68ed09..e35fdd33e 100644 --- a/wacore/src/store/traits.rs +++ b/wacore/src/store/traits.rs @@ -14,6 +14,7 @@ use bytes::Bytes; use serde::{Deserialize, Serialize}; use std::sync::Arc; use wacore_appstate::processor::AppStateMutationMAC; +use wacore_binary::Jid; /// Inline protocol-sized message secret. The array makes invalid lengths /// unrepresentable without a heap allocation or pointer indirection per row. @@ -80,6 +81,46 @@ pub struct MsgSecretEntry { pub message_ts: i64, } +impl MsgSecretEntry { + /// The canonical non-AD sender identifier for a row whose chat identifier is + /// already in hand, sharing that allocation whenever both JIDs address the + /// same user. That covers every direct message (chat and sender are the peer) + /// and every self-authored history row, which is what the `sender` field doc + /// means by "often aliases `chat`". + pub fn sender_id_for(chat: &Jid, chat_id: &Arc, sender: &Jid) -> Arc { + if sender.is_same_chat_as(chat) { + Arc::clone(chat_id) + } else { + sender.to_non_ad_arc_str() + } + } + + /// Build a row from the JIDs the send and receive paths already carry. + /// + /// Single chokepoint for how a row's three identifiers are derived, so the + /// inbound capture and the outbound persist cannot drift on either the + /// canonicalisation or the aliasing above. + pub fn new( + chat: &Jid, + sender: &Jid, + msg_id: &str, + secret: MessageSecret, + expires_at: i64, + message_ts: i64, + ) -> Self { + let chat_id = chat.to_non_ad_arc_str(); + let sender_id = Self::sender_id_for(chat, &chat_id, sender); + Self { + chat: chat_id, + sender: sender_id, + msg_id: Arc::from(msg_id), + secret, + expires_at, + message_ts, + } + } +} + /// Device information for registry tracking. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DeviceInfo { @@ -109,6 +150,84 @@ impl DeviceInfo { } } +#[cfg(test)] +mod msg_secret_entry_tests { + use super::{Jid, MsgSecretEntry}; + use std::sync::Arc; + + fn jid(s: &str) -> Jid { + s.parse().unwrap_or_else(|e| panic!("parse {s}: {e}")) + } + + /// A direct message names the same user as chat and as sender (the sender + /// only adds a device suffix), so the row must carry one shared allocation + /// rather than two identical strings. + #[test] + fn direct_message_shares_one_identifier_allocation() { + let entry = MsgSecretEntry::new( + &jid("5511987650001@s.whatsapp.net"), + &jid("5511987650001:33@s.whatsapp.net"), + "3EB0AABBCCDDEEFF0011", + [7u8; 32], + 0, + 0, + ); + + assert_eq!(&*entry.chat, "5511987650001@s.whatsapp.net"); + assert_eq!(&*entry.sender, "5511987650001@s.whatsapp.net"); + assert!( + Arc::ptr_eq(&entry.chat, &entry.sender), + "chat and sender must share the allocation when they are the same user" + ); + assert_eq!(&*entry.msg_id, "3EB0AABBCCDDEEFF0011"); + } + + /// A group row (and an outbound row, where the sender is us) names two + /// different users, which must stay two distinct identifiers. + #[test] + fn distinct_users_keep_separate_identifiers() { + let entry = MsgSecretEntry::new( + &jid("120363021033254949@g.us"), + &jid("5511987650001:2@s.whatsapp.net"), + "M1", + [0u8; 32], + 123, + 456, + ); + + assert_eq!(&*entry.chat, "120363021033254949@g.us"); + assert_eq!(&*entry.sender, "5511987650001@s.whatsapp.net"); + assert!(!Arc::ptr_eq(&entry.chat, &entry.sender)); + assert_eq!((entry.expires_at, entry.message_ts), (123, 456)); + } + + /// Same user part, different namespace: the LID and PN forms are distinct + /// lookup keys and must never be collapsed into one. + #[test] + fn same_user_across_namespaces_is_not_aliased() { + let chat = jid("100000012345678@lid"); + let entry = MsgSecretEntry::new( + &chat, + &jid("100000012345678@s.whatsapp.net"), + "", + [0u8; 32], + 0, + 0, + ); + + assert_eq!(&*entry.chat, "100000012345678@lid"); + assert_eq!(&*entry.sender, "100000012345678@s.whatsapp.net"); + assert!(!Arc::ptr_eq(&entry.chat, &entry.sender)); + // An empty message id is a degenerate but representable key, not a panic. + assert_eq!(&*entry.msg_id, ""); + + // The standalone helper must make the same call as the constructor. + let chat_id: Arc = Arc::from("100000012345678@lid"); + let aliased = MsgSecretEntry::sender_id_for(&chat, &chat_id, &jid("100000012345678:9@lid")); + assert!(Arc::ptr_eq(&aliased, &chat_id)); + } +} + #[cfg(test)] mod device_info_tests { use super::DeviceInfo; diff --git a/wacore/src/types/message.rs b/wacore/src/types/message.rs index 137519888..cc4c6437e 100644 --- a/wacore/src/types/message.rs +++ b/wacore/src/types/message.rs @@ -382,6 +382,28 @@ mod tests { use super::*; use buffa::MessageField; + /// The stanza parser reads `edit` as a borrowed attribute and parses it + /// directly. That is only safe while the borrowed and owned constructors + /// agree on every input, including the `Unknown` fallback that is the one + /// case actually needing an allocation. + #[test] + fn edit_attribute_parses_identically_from_borrowed_and_owned() { + for wire in ["", "1", "2", "3", "7", "8", "0", "99", "revogação", " 7"] { + assert_eq!( + EditAttribute::from(wire), + EditAttribute::from(wire.to_owned()), + "mismatch for {wire:?}" + ); + } + assert_eq!(EditAttribute::from("7"), EditAttribute::SenderRevoke); + // An unrecognized value must keep its exact wire bytes so the resend + // path can echo them back verbatim. + assert_eq!( + EditAttribute::from("99"), + EditAttribute::Unknown("99".to_owned()) + ); + } + #[test] fn message_info_serde_omits_only_absent_optional_fields() { let mut info = MessageInfo::default();