diff --git a/src/client/messaging.rs b/src/client/messaging.rs index 3f2a0c39a..08ea27477 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -39,10 +39,16 @@ impl Client { /// /// Always drains `frames`, including when no socket is installed, while /// retaining its outer allocation for the persistent workers to reuse. + /// + /// Results land in `results`, which the caller owns and reuses too. A + /// returned `Vec` would allocate once per burst, and the common burst is a + /// single frame, so that allocation was the dominant cost of sending one. pub(crate) async fn send_raw_bytes_burst( &self, frames: &mut Vec>, - ) -> Result, ClientError> { + results: &mut Vec, + ) -> Result<(), ClientError> { + results.clear(); let noise_socket = match self.get_noise_socket().await { Ok(socket) => socket, Err(error) => { @@ -52,16 +58,27 @@ impl Client { }; if frames.len() == 1 { let plaintext = frames.pop().expect("length checked"); - return Ok(vec![ + results.push( noise_socket .encrypt_and_send(bytes::Bytes::from(plaintext)) .await, - ]); + ); + return Ok(()); } + // The out-parameter buys nothing here yet: `join_all` allocates its own + // storage and a `Vec` for the results, so a multi-frame burst still + // pays what it used to, plus this copy. Removing it means enqueueing + // every job before awaiting any of them and holding the receivers + // inline, which needs `encrypt_and_send` split into enqueue and await; + // that is a change to the send path rather than to this function. + // Tracked separately. Awaiting in sequence instead is not the fix: the + // frames would reach the sender one completion apart and stop batching + // into a single write. let sends = frames .drain(..) .map(|plaintext| noise_socket.encrypt_and_send(bytes::Bytes::from(plaintext))); - Ok(futures::future::join_all(sends).await) + results.extend(futures::future::join_all(sends).await); + Ok(()) } #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.node", level = "debug", skip_all, fields(tag = %node.tag), err(Debug)))] diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 6b257ae0a..d836343e0 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -697,6 +697,7 @@ impl Client { let mut batch = Vec::with_capacity(Self::MAX_ACK_BURST); let mut frames = Vec::with_capacity(Self::MAX_ACK_BURST); let mut guards = Vec::with_capacity(Self::MAX_ACK_BURST); + let mut results = Vec::with_capacity(Self::MAX_ACK_BURST); while let Ok(first) = rx.recv().await { let Some(client) = client.upgrade() else { break; @@ -752,9 +753,9 @@ impl Client { // EnteredSpan is not Send and cannot cross the await. let frame_count = frames.len(); let send_and_report = async { - match client.send_raw_bytes_burst(&mut frames).await { - Ok(results) => { - for result in results { + match client.send_raw_bytes_burst(&mut frames, &mut results).await { + Ok(()) => { + for result in results.drain(..) { if let Err(e) = result && !e.is_transport_unavailable() { diff --git a/src/client/tests.rs b/src/client/tests.rs index 93c47afd0..bf926b94e 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -3409,22 +3409,43 @@ async fn raw_bytes_burst_drains_and_reuses_input_on_happy_paths() { let mut frames = Vec::with_capacity(4); let retained_capacity = frames.capacity(); frames.push(vec![0x11; 32]); - let single = client - .send_raw_bytes_burst(&mut frames) + // Sized for the largest burst below, so a reallocation here would mean the + // callee replaced the buffer rather than filling it. + let mut results = Vec::with_capacity(4); + // Captured before the first call, not between the two: taken after it, a + // replacement made on the single-frame path would already be the buffer + // this compares against and would go unnoticed. + let results_ptr = results.as_ptr(); + client + .send_raw_bytes_burst(&mut frames, &mut results) .await .expect("installed socket"); - assert_eq!(single.len(), 1); - assert!(single.into_iter().all(|result| result.is_ok())); + assert_eq!(results.len(), 1); + assert!(results.iter().all(|result| result.is_ok())); + assert_eq!( + results.as_ptr(), + results_ptr, + "the single-frame path must fill the caller's buffer, not replace it" + ); assert!(frames.is_empty(), "the single-frame fast path must drain"); assert_eq!(frames.capacity(), retained_capacity); frames.extend((0..4).map(|index| vec![index; 32])); - let burst = client - .send_raw_bytes_burst(&mut frames) + client + .send_raw_bytes_burst(&mut frames, &mut results) .await .expect("installed socket"); - assert_eq!(burst.len(), 4); - assert!(burst.into_iter().all(|result| result.is_ok())); + assert_eq!(results.len(), 4); + assert!(results.iter().all(|result| result.is_ok())); + // Identity, not capacity: a fresh Vec of the same capacity would satisfy a + // capacity check while defeating the whole point of the out-parameter. The + // buffer is preallocated above so the second burst cannot legitimately + // reallocate it. + assert_eq!( + results.as_ptr(), + results_ptr, + "the caller's results buffer must be the same allocation, not an equal one" + ); assert!(frames.is_empty(), "the joined path must drain"); assert_eq!(frames.capacity(), retained_capacity); assert_eq!(transport.sent_count(), 5, "every frame must reach the wire"); @@ -3442,7 +3463,8 @@ async fn raw_bytes_burst_drains_input_when_disconnected() { let retained_capacity = frames.capacity(); frames.extend([vec![0x21; 32], vec![0x22; 32]]); - let result = client.send_raw_bytes_burst(&mut frames).await; + let mut results = Vec::new(); + let result = client.send_raw_bytes_burst(&mut frames, &mut results).await; assert!( matches!(result, Err(ClientError::NotConnected)), "a missing socket must remain an outer NotConnected error: {result:?}" @@ -3469,11 +3491,12 @@ async fn raw_bytes_burst_surfaces_transport_then_poisoned_per_frame() { let mut frames = Vec::with_capacity(4); let retained_capacity = frames.capacity(); frames.push(vec![0x31; 32]); - let mut failed = client - .send_raw_bytes_burst(&mut frames) + let mut results = Vec::new(); + client + .send_raw_bytes_burst(&mut frames, &mut results) .await .expect("the socket lookup itself succeeds"); - let transport_error = failed + let transport_error = results .pop() .expect("one result") .expect_err("the transport is configured to fail"); @@ -3486,11 +3509,11 @@ async fn raw_bytes_burst_surfaces_transport_then_poisoned_per_frame() { assert_eq!(frames.capacity(), retained_capacity); frames.push(vec![0x32; 32]); - let mut poisoned = client - .send_raw_bytes_burst(&mut frames) + client + .send_raw_bytes_burst(&mut frames, &mut results) .await .expect("the installed socket remains reachable"); - let poisoned_error = poisoned + let poisoned_error = results .pop() .expect("one result") .expect_err("the sender must reject work after an ambiguous write"); @@ -3524,8 +3547,9 @@ async fn raw_bytes_burst_surfaces_a_closed_sender_per_frame() { let mut frames = Vec::with_capacity(4); let retained_capacity = frames.capacity(); frames.push(vec![0x41; 32]); - let mut results = client - .send_raw_bytes_burst(&mut frames) + let mut results = Vec::new(); + client + .send_raw_bytes_burst(&mut frames, &mut results) .await .expect("the installed socket remains reachable"); let error = results diff --git a/src/message/dispatch.rs b/src/message/dispatch.rs index 9b25a7b73..a0504f5a2 100644 --- a/src/message/dispatch.rs +++ b/src/message/dispatch.rs @@ -158,6 +158,7 @@ impl Client { let mut batch = Vec::with_capacity(Self::MAX_RECEIPT_BURST); let mut frames = Vec::with_capacity(Self::MAX_RECEIPT_BURST); let mut guards = Vec::with_capacity(Self::MAX_RECEIPT_BURST); + let mut results = Vec::with_capacity(Self::MAX_RECEIPT_BURST); while let Ok(first) = rx.recv().await { let Some(client) = client.upgrade() else { break; @@ -198,9 +199,9 @@ impl Client { // single-receipt path, which this one does not use. let frame_count = frames.len(); let send_and_report = async { - match client.send_raw_bytes_burst(&mut frames).await { - Ok(results) => { - for result in results { + match client.send_raw_bytes_burst(&mut frames, &mut results).await { + Ok(()) => { + for result in results.drain(..) { if let Some(error) = delivery_receipt_burst_warning(&result) { log::warn!(target: "Client/Receipt", "Failed to send delivery receipt: {error:?}"); } diff --git a/wacore/libsignal/src/protocol/state/session.rs b/wacore/libsignal/src/protocol/state/session.rs index 28fadf37e..f19bedb03 100644 --- a/wacore/libsignal/src/protocol/state/session.rs +++ b/wacore/libsignal/src/protocol/state/session.rs @@ -97,6 +97,28 @@ pub(crate) enum ReceiverChainState { Closed { next_index: u32 }, } +/// Writes a chain key's material into the protobuf field, reusing the buffer +/// already there when it is ours to reuse. +/// +/// `Bytes` is immutable, so assigning a fresh `copy_from_slice` allocates on +/// every ratchet advance, and the ratchet advances three times per message +/// round trip. After a checkout the record is uniquely owned (the cache takes +/// it out of its `Arc` with `try_unwrap`), so the old buffer is almost always +/// reusable: take it, overwrite in place, and freeze it back. A buffer that is +/// still shared, or one of an unexpected length, falls back to allocating, +/// which is exactly the previous behaviour. +fn write_chain_key(field: &mut Option, key: &[u8]) { + if let Some(existing) = field.take() + && existing.len() == key.len() + && let Ok(mut owned) = existing.try_into_mut() + { + owned.copy_from_slice(key); + *field = Some(owned.freeze()); + return; + } + *field = Some(bytes::Bytes::copy_from_slice(key)); +} + impl SessionState { pub fn from_session_structure(session: SessionStructure) -> Self { Self { session } @@ -456,7 +478,7 @@ impl SessionState { match chain.chain_key.as_option_mut() { Some(existing) => { existing.index = Some(next_chain_key.index()); - existing.key = Some(Bytes::copy_from_slice(next_chain_key.key())); + write_chain_key(&mut existing.key, next_chain_key.key()); } None => { chain.chain_key = MessageField::some(session_structure::chain::ChainKey { @@ -600,7 +622,7 @@ impl SessionState { match target.as_option_mut() { Some(existing) => { existing.index = Some(chain_key.index()); - existing.key = Some(Bytes::copy_from_slice(chain_key.key())); + write_chain_key(&mut existing.key, chain_key.key()); } None => { *target = MessageField::some(session_structure::chain::ChainKey { @@ -2023,3 +2045,75 @@ mod tests { ); } } + +#[cfg(test)] +mod chain_key_buffer_tests { + use super::*; + use bytes::Bytes; + + /// The ratchet advances three times per message round trip, so the buffer + /// this writes into is the difference between one allocation per advance + /// and none. Reuse is only valid when the buffer is ours alone. + #[test] + fn a_uniquely_owned_buffer_is_written_in_place() { + let mut field = Some(Bytes::copy_from_slice(&[0u8; 32])); + let before = field.as_ref().expect("seeded").as_ptr(); + + write_chain_key(&mut field, &[7u8; 32]); + + let after = field.as_ref().expect("written"); + assert_eq!( + after.as_ref(), + &[7u8; 32], + "the new key must be what is read back" + ); + assert_eq!( + after.as_ptr(), + before, + "a uniquely owned buffer must be reused, not replaced" + ); + } + + /// Bad path: a buffer someone else still holds cannot be overwritten, or + /// that holder would observe a key it never asked for. Falling back to a + /// fresh allocation is the whole point of the guard. + #[test] + fn a_shared_buffer_is_never_overwritten() { + let shared = Bytes::copy_from_slice(&[1u8; 32]); + let observer = shared.clone(); + let mut field = Some(shared); + + write_chain_key(&mut field, &[9u8; 32]); + + assert_eq!( + observer.as_ref(), + &[1u8; 32], + "the other holder must still see what it had" + ); + assert_eq!(field.expect("written").as_ref(), &[9u8; 32]); + } + + /// Bad path: a stored key of the wrong length (a record from an older + /// format) must not be partially overwritten, leaving a key that is half + /// old and half new. + #[test] + fn a_buffer_of_the_wrong_length_is_replaced_whole() { + let mut field = Some(Bytes::copy_from_slice(&[3u8; 16])); + + write_chain_key(&mut field, &[4u8; 32]); + + let written = field.expect("written"); + assert_eq!(written.len(), 32, "the new key's length wins"); + assert_eq!(written.as_ref(), &[4u8; 32]); + } + + /// An empty field is the None arm: nothing to reuse, so it allocates. + #[test] + fn an_absent_buffer_is_created() { + let mut field: Option = None; + + write_chain_key(&mut field, &[5u8; 32]); + + assert_eq!(field.expect("written").as_ref(), &[5u8; 32]); + } +} diff --git a/wacore/src/messages.rs b/wacore/src/messages.rs index bbcee43d2..2c41483fa 100644 --- a/wacore/src/messages.rs +++ b/wacore/src/messages.rs @@ -10,6 +10,131 @@ use waproto::whatsapp as wa; pub struct MessageUtils; +/// Names the DSM destination without requiring it to exist as a string. +/// +/// The DSM field needs the JID's length before its bytes, so the caller used to +/// render one into a `String` just to measure it and copy it. A `Jid` can do +/// both without the intermediate: this is what lets `&Jid` and `&str` share the +/// same body instead of the format existing in two shapes. +/// +/// Implemented for `str`, the standard string wrappers, and `Jid`, and carried +/// through references of any depth. A type outside that set (a string newtype, +/// say) has two ways in: pass `&*wrapper` to go through the `str` +/// implementation, or implement this trait for it, which is why the trait is +/// public. What is deliberately *not* offered is a blanket implementation over +/// `Deref`: it collides with `Jid` and with the reference +/// implementations at once, since `&str` derefs to `str` as well. +pub trait DsmDestination { + /// Exactly the number of bytes [`Self::write_into`] appends. + /// + /// The DSM field writes this as a length prefix before the bytes, so an + /// implementation that disagrees with itself puts a prefix on the wire that + /// points past its own payload, and the peer reads the next protobuf field + /// from the wrong offset. + fn encoded_len(&self) -> usize; + + /// Appends the destination's wire form, whose length must equal + /// [`Self::encoded_len`]. + fn write_into(&self, out: &mut Vec); +} + +/// A generic parameter does not deref-coerce the way the old `&str` argument +/// did, so every shape a caller could previously pass has to be reachable by +/// an implementation instead. +/// +/// The trait is therefore implemented on the *owned* types, and the two blanket +/// implementations below carry it through references. That covers a reference +/// of any depth and mutability (`&String`, `&mut String`, `&&str`) without +/// naming each one, which a finite list cannot do. +/// +/// Recorded so it is not attempted again: a blanket +/// `impl>` over the *reference* types instead would +/// collide with the `Jid` implementation, because coherence cannot rule out +/// `Jid` gaining that `Deref`. +macro_rules! dsm_destination_via_str { + ($($ty:ty),+ $(,)?) => {$( + impl DsmDestination for $ty { + #[inline] + fn encoded_len(&self) -> usize { + str::len(self) + } + + #[inline] + fn write_into(&self, out: &mut Vec) { + out.extend_from_slice(str::as_bytes(self)); + } + } + )+}; +} + +dsm_destination_via_str!( + str, + String, + Box, + std::rc::Rc, + std::sync::Arc, + std::borrow::Cow<'_, str>, +); + +impl DsmDestination for &T { + #[inline] + fn encoded_len(&self) -> usize { + (**self).encoded_len() + } + + #[inline] + fn write_into(&self, out: &mut Vec) { + (**self).write_into(out); + } +} + +impl DsmDestination for &mut T { + #[inline] + fn encoded_len(&self) -> usize { + (**self).encoded_len() + } + + #[inline] + fn write_into(&self, out: &mut Vec) { + (**self).write_into(out); + } +} + +impl DsmDestination for wacore_binary::jid::Jid { + #[inline] + fn encoded_len(&self) -> usize { + let mut counter = DisplayLen(0); + // Infallible: the counter never errors, so the render always completes. + let _ = self.write_display_to(&mut counter); + counter.0 + } + + #[inline] + fn write_into(&self, out: &mut Vec) { + let _ = self.write_display_to(&mut Utf8Sink(out)); + } +} + +/// Counts what a render would write, so the length is known before the bytes. +struct DisplayLen(usize); + +impl core::fmt::Write for DisplayLen { + fn write_str(&mut self, s: &str) -> core::fmt::Result { + self.0 += s.len(); + Ok(()) + } +} + +/// Appends a render straight into the wire buffer. +struct Utf8Sink<'a>(&'a mut Vec); + +impl core::fmt::Write for Utf8Sink<'_> { + fn write_str(&mut self, s: &str) -> core::fmt::Result { + self.0.extend_from_slice(s.as_bytes()); + Ok(()) + } +} + impl MessageUtils { fn random_pad_len() -> u8 { use rand::RngExt; @@ -126,7 +251,7 @@ impl MessageUtils { pub fn encode_dm_plaintexts( message: &wa::Message, extra_context: Option<&wa::MessageContextInfo>, - destination_jid: &str, + destination_jid: impl DsmDestination, ) -> DmPlaintexts { if message.message_context_info.is_set() { let mut owned = message.clone(); @@ -162,7 +287,7 @@ impl MessageUtils { }); let mut msg_cache = buffa::SizeCache::new(); let content_len = waproto::codec::message_compute_size(message, &mut msg_cache); - let dest = destination_jid.as_bytes(); + let dest_len = destination_jid.encoded_len(); // recipient = content (encoded once) + the extra message_context_info field. // Pre-size for content + the appended mci field + padding so it never @@ -175,7 +300,7 @@ impl MessageUtils { // pre-computed so the spliced content goes straight in, and the buffer is sized // exactly (device_sent_message field + mci field + padding): one allocation, no // reallocation regardless of whether extra_context is present. - let dsm_len = len_delimited_len(TAG_DSM_DESTINATION_JID, dest.len()) + let dsm_len = len_delimited_len(TAG_DSM_DESTINATION_JID, dest_len) + len_delimited_len(TAG_DSM_MESSAGE, content_len); let own_cap = len_delimited_len(TAG_DEVICE_SENT_MESSAGE, dsm_len) + mci_field_len + MAX_PAD; let mut own_devices = Vec::with_capacity(own_cap); @@ -185,7 +310,13 @@ impl MessageUtils { &mut own_devices, ); push_varint(dsm_len as u64, &mut own_devices); // DeviceSentMessage length - push_len_delimited(TAG_DSM_DESTINATION_JID, dest, &mut own_devices); + push_wire_tag( + TAG_DSM_DESTINATION_JID, + buffa::encoding::WireType::LengthDelimited, + &mut own_devices, + ); + push_varint(dest_len as u64, &mut own_devices); + destination_jid.write_into(&mut own_devices); push_len_delimited(TAG_DSM_MESSAGE, &recipient[..content_len], &mut own_devices); if let Some(extra) = extra_context { push_message_field(TAG_MESSAGE_CONTEXT_INFO, extra, &mut own_devices); @@ -214,7 +345,7 @@ impl MessageUtils { pub fn dm_plaintexts_from_encoded( content: &[u8], extra_context: Option<&wa::MessageContextInfo>, - destination_jid: &str, + destination_jid: impl DsmDestination, ) -> DmPlaintexts { const MAX_PAD: usize = 16; @@ -226,12 +357,12 @@ impl MessageUtils { ) }); let content_len = content.len(); - let dest = destination_jid.as_bytes(); + let dest_len = destination_jid.encoded_len(); let mut recipient = Vec::with_capacity(content_len + mci_field_len + MAX_PAD); recipient.extend_from_slice(content); - let dsm_len = len_delimited_len(TAG_DSM_DESTINATION_JID, dest.len()) + let dsm_len = len_delimited_len(TAG_DSM_DESTINATION_JID, dest_len) + len_delimited_len(TAG_DSM_MESSAGE, content_len); let own_cap = len_delimited_len(TAG_DEVICE_SENT_MESSAGE, dsm_len) + mci_field_len + MAX_PAD; let mut own_devices = Vec::with_capacity(own_cap); @@ -241,7 +372,13 @@ impl MessageUtils { &mut own_devices, ); push_varint(dsm_len as u64, &mut own_devices); // DeviceSentMessage length - push_len_delimited(TAG_DSM_DESTINATION_JID, dest, &mut own_devices); + push_wire_tag( + TAG_DSM_DESTINATION_JID, + buffa::encoding::WireType::LengthDelimited, + &mut own_devices, + ); + push_varint(dest_len as u64, &mut own_devices); + destination_jid.write_into(&mut own_devices); push_len_delimited(TAG_DSM_MESSAGE, content, &mut own_devices); if let Some(extra) = extra_context { push_message_field(TAG_MESSAGE_CONTEXT_INFO, extra, &mut own_devices); @@ -261,7 +398,10 @@ impl MessageUtils { /// top-level `message_context_info` that must be hoisted onto the DSM wrapper /// (`wrap_device_sent` semantics). Hoisting requires ownership; the common path in /// [`encode_dm_plaintexts`] borrows instead and never reaches this. - fn encode_dm_plaintexts_owned(mut message: wa::Message, destination_jid: &str) -> DmPlaintexts { + fn encode_dm_plaintexts_owned( + mut message: wa::Message, + destination_jid: impl DsmDestination, + ) -> DmPlaintexts { const MAX_PAD: usize = 16; // Hoist message_context_info onto the wrapper (as wrap_device_sent does) so the @@ -277,12 +417,12 @@ impl MessageUtils { }); let mut msg_cache = buffa::SizeCache::new(); let content_len = waproto::codec::message_compute_size(&message, &mut msg_cache); - let dest = destination_jid.as_bytes(); + let dest_len = destination_jid.encoded_len(); let mut recipient = Vec::with_capacity(content_len + mci_field_len + MAX_PAD); waproto::codec::message_write_to(&message, &mut msg_cache, &mut recipient); - let dsm_len = len_delimited_len(TAG_DSM_DESTINATION_JID, dest.len()) + let dsm_len = len_delimited_len(TAG_DSM_DESTINATION_JID, dest_len) + len_delimited_len(TAG_DSM_MESSAGE, content_len); let own_cap = len_delimited_len(TAG_DEVICE_SENT_MESSAGE, dsm_len) + mci_field_len + MAX_PAD; let mut own_devices = Vec::with_capacity(own_cap); @@ -292,7 +432,13 @@ impl MessageUtils { &mut own_devices, ); push_varint(dsm_len as u64, &mut own_devices); // DeviceSentMessage length - push_len_delimited(TAG_DSM_DESTINATION_JID, dest, &mut own_devices); + push_wire_tag( + TAG_DSM_DESTINATION_JID, + buffa::encoding::WireType::LengthDelimited, + &mut own_devices, + ); + push_varint(dest_len as u64, &mut own_devices); + destination_jid.write_into(&mut own_devices); push_len_delimited(TAG_DSM_MESSAGE, &recipient[..content_len], &mut own_devices); if let Some(mci) = &mci { push_message_field(TAG_MESSAGE_CONTEXT_INFO, mci, &mut own_devices); @@ -1828,6 +1974,162 @@ mod parse_message_info_tests { #[cfg(test)] #[allow(clippy::disallowed_methods)] mod device_sent_tests { + + /// The DSM field writes its length before its bytes, so a `Jid` that + /// measures itself differently than it renders would emit a length prefix + /// that disagrees with the payload: the peer would then read the next field + /// from the wrong offset and the whole message would be garbage. + #[test] + fn a_jid_measures_itself_exactly_as_it_renders() { + use wacore_binary::jid::Jid; + + let cases = [ + "5511987650001@s.whatsapp.net", + "5511987650001:5@s.whatsapp.net", + "5511987650001.2:5@s.whatsapp.net", + "120363021033254949@g.us", + "100000012345678:25@lid", + "867051314767696:0@bot", + "status@broadcast", + "ẞünïcodé-ñ@s.whatsapp.net", + ]; + + for case in cases { + let jid: Jid = case.parse().unwrap_or_else(|e| panic!("parse {case}: {e}")); + let mut written = Vec::new(); + jid.write_into(&mut written); + + assert_eq!( + jid.encoded_len(), + written.len(), + "{case}: the counted length must equal the bytes written" + ); + assert_eq!( + written, + jid.to_string().into_bytes(), + "{case}: writing directly must match rendering through a String" + ); + } + } + + /// The two spellings of the same destination must produce identical wire + /// bytes, since one of them is what actually goes out now. + #[test] + fn naming_the_destination_by_jid_matches_naming_it_by_string() { + use wacore_binary::jid::Jid; + + let jid: Jid = "5511987650001:5@s.whatsapp.net".parse().expect("parse"); + let message = wa::Message { + conversation: Some("destination check".to_string()), + ..Default::default() + }; + let content = waproto::codec::message_to_vec(&message); + + let by_string = + MessageUtils::dm_plaintexts_from_encoded(&content, None, jid.to_string().as_str()); + let by_jid = MessageUtils::dm_plaintexts_from_encoded(&content, None, &jid); + + // Padding is random, so compare the unpadded prefix: everything the DSM + // field contributes lands before it. + let prefix = by_jid.own_devices.len().min(by_string.own_devices.len()) - 16; + assert_eq!( + by_jid.own_devices[..prefix], + by_string.own_devices[..prefix], + "the DSM bytes must not depend on how the destination was named" + ); + } + + /// The generic parameter does not deref-coerce, so every wrapper a caller + /// may hold its destination in has to be named by an implementation. This + /// is a compile-time check as much as a runtime one: a wrapper missing from + /// the list fails to build here rather than in a downstream crate. + #[test] + fn every_string_wrapper_names_the_same_destination() { + use std::borrow::Cow; + use std::rc::Rc; + use std::sync::Arc; + + let message = wa::Message { + conversation: Some("wrapper check".to_string()), + ..Default::default() + }; + let content = waproto::codec::message_to_vec(&message); + let dest = "5511987650001:5@s.whatsapp.net"; + + let mut owned = dest.to_string(); + let boxed: Box = dest.into(); + let rc: Rc = dest.into(); + let arc: Arc = dest.into(); + let cow: Cow<'_, str> = Cow::Borrowed(dest); + let mut owned_mut = dest.to_string(); + + let reference = MessageUtils::dm_plaintexts_from_encoded(&content, None, dest); + let prefix = reference.own_devices.len() - 16; + + // A mutable destination coerced to `&str` before this trait existed, so + // both forms have to keep working; `&mut str` is reached through the + // owned string it borrows from. + let by_mut_string = + MessageUtils::dm_plaintexts_from_encoded(&content, None, &mut owned_mut); + let by_mut_str = + MessageUtils::dm_plaintexts_from_encoded(&content, None, &mut *owned.as_mut_str()); + + // Through `encode_dm_plaintexts` specifically: it is the entry point + // that carried a `Copy` bound, which no `&mut` destination satisfies. + let by_mut_through_encode = + MessageUtils::encode_dm_plaintexts(&message, None, &mut owned_mut); + + // Nested references coerced too, and a finite list of implementations + // could never cover every depth. These pin that the blanket ones do, so + // the extra borrows clippy offers to remove are the whole point here. + #[allow( + clippy::needless_borrows_for_generic_args, + reason = "the nested reference is what this asserts is accepted" + )] + let (by_double, by_triple, by_ref_to_owned) = ( + MessageUtils::dm_plaintexts_from_encoded(&content, None, &dest), + MessageUtils::dm_plaintexts_from_encoded(&content, None, &&dest), + MessageUtils::dm_plaintexts_from_encoded(&content, None, &&owned), + ); + + for (name, produced) in [ + ("&&str", by_double), + ("&&&str", by_triple), + ("&&String", by_ref_to_owned), + ("&mut String", by_mut_string), + ("&mut str", by_mut_str), + ( + "&mut String via encode_dm_plaintexts", + by_mut_through_encode, + ), + ( + "String", + MessageUtils::dm_plaintexts_from_encoded(&content, None, &owned), + ), + ( + "Box", + MessageUtils::dm_plaintexts_from_encoded(&content, None, &boxed), + ), + ( + "Rc", + MessageUtils::dm_plaintexts_from_encoded(&content, None, &rc), + ), + ( + "Arc", + MessageUtils::dm_plaintexts_from_encoded(&content, None, &arc), + ), + ( + "Cow", + MessageUtils::dm_plaintexts_from_encoded(&content, None, &cow), + ), + ] { + assert_eq!( + produced.own_devices[..prefix], + reference.own_devices[..prefix], + "a destination held in {name} must name itself exactly as &str does" + ); + } + } use super::*; fn msg_with_secret(secret: &[u8]) -> wa::Message { diff --git a/wacore/src/reporting_token.rs b/wacore/src/reporting_token.rs index b08710d44..5e7b46ba9 100644 --- a/wacore/src/reporting_token.rs +++ b/wacore/src/reporting_token.rs @@ -519,21 +519,27 @@ impl Piece { } } -/// Extract reporting token content from encoded protobuf message bytes. -/// -/// This function parses raw protobuf bytes and extracts only the fields -/// specified in the whitelist, matching WhatsApp Web's behavior. -/// -/// # Arguments -/// * `data` - Raw protobuf-encoded message bytes -/// * `whitelist` - List of fields to extract +#[cfg(test)] +thread_local! { + /// Counts collector runs so a test can pin that generating a token parses + /// the message once. A nested field owns its bytes, so a second collection + /// is not just a second parse, it is a second set of buffers. + static COLLECT_CALLS: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +/// The whitelisted pieces, in the order the token concatenates them. /// -/// # Returns -/// Concatenated bytes of all extracted fields, or None if no fields match -pub fn extract_reporting_token_content( +/// Split out from [`extract_reporting_token_content`] so the HMAC can consume +/// them directly: the token is `HMAC(key, concat(pieces))` and `Mac::update` is +/// associative over its input, so feeding the pieces one by one gives the same +/// bytes without building the concatenation at all. +fn collect_reporting_token_pieces( data: &[u8], whitelist: &[ReportingField], -) -> Option> { +) -> Option> { + #[cfg(test)] + COLLECT_CALLS.with(|c| c.set(c.get() + 1)); + // The token's bytes are contract: fields are concatenated in ascending // field-number order, ties in wire order (`sort_by_key` is stable). Only // where the bytes come from changed -- a flat field is named by its range @@ -649,11 +655,25 @@ pub fn extract_reporting_token_content( return None; } + // The token's bytes are contract: ascending field number, ties in wire + // order, which `sort_by_key` preserves because it is stable. extracted.sort_by_key(|(num, _)| *num); + Some(extracted) +} + +/// Extract reporting token content from encoded protobuf message bytes. +/// +/// Builds the concatenation. The send path does not need it: it feeds the +/// pieces straight to the HMAC instead. +pub fn extract_reporting_token_content( + data: &[u8], + whitelist: &[ReportingField], +) -> Option> { + let pieces = collect_reporting_token_pieces(data, whitelist)?; - let total_len: usize = extracted.iter().map(|(_, piece)| piece.len()).sum(); + let total_len: usize = pieces.iter().map(|(_, piece)| piece.len()).sum(); let mut result = Vec::with_capacity(total_len); - for (_, piece) in extracted { + for (_, piece) in pieces { match piece { Piece::Borrowed(range) => result.extend_from_slice(&data[range]), Piece::Owned(bytes) => result.extend_from_slice(&bytes), @@ -706,6 +726,44 @@ pub fn calculate_reporting_token( Ok(token) } +/// Same token as [`calculate_reporting_token`] over the concatenated content, +/// without building that concatenation. +/// +/// `Mac::update` is associative over its input, so feeding each whitelisted +/// piece in token order hashes exactly the bytes the concatenation would have +/// held. +/// +/// Takes the pieces rather than collecting them: a nested field owns its +/// re-framed bytes, so collecting twice would materialise those buffers twice +/// and parse the message twice, which costs more than the concatenation this +/// avoids. +fn calculate_reporting_token_over_pieces( + reporting_token_key: &[u8; REPORTING_TOKEN_KEY_SIZE], + data: &[u8], + pieces: &[(u32, Piece)], +) -> Option<[u8; REPORTING_TOKEN_SIZE]> { + // Both callers reach here through `collect_reporting_token_pieces`, which + // already returns `None` for a message with nothing whitelisted. Repeated + // here because the alternative is minting a perfectly valid token over no + // content at all, and that must not depend on a caller remembering to check. + if pieces.is_empty() { + return None; + } + + let mut mac = Hmac::::new_from_slice(reporting_token_key).ok()?; + for (_, piece) in pieces { + match piece { + Piece::Borrowed(range) => mac.update(&data[range.clone()]), + Piece::Owned(bytes) => mac.update(bytes), + } + } + + let result = mac.finalize().into_bytes(); + let mut token = [0u8; REPORTING_TOKEN_SIZE]; + token.copy_from_slice(&result[..REPORTING_TOKEN_SIZE]); + Some(token) +} + /// Result of generating a reporting token for a message #[derive(Debug, Clone)] pub struct ReportingTokenResult { @@ -762,7 +820,7 @@ pub fn generate_reporting_token_from_encoded( if !should_include_reporting_token(message) { return None; } - let content = extract_reporting_token_content(encoded_message, REPORTING_FIELDS)?; + let pieces = collect_reporting_token_pieces(encoded_message, REPORTING_FIELDS)?; let message_secret: [u8; MESSAGE_SECRET_SIZE] = if let Some(secret) = existing_secret { if secret.len() != MESSAGE_SECRET_SIZE { @@ -779,7 +837,7 @@ pub fn generate_reporting_token_from_encoded( derive_reporting_token_key_for_jids(&message_secret, stanza_id, sender_jid, remote_jid) .ok()?; - let token = calculate_reporting_token(&key, &content).ok()?; + let token = calculate_reporting_token_over_pieces(&key, encoded_message, &pieces)?; Some(ReportingTokenResult { message_secret, @@ -841,6 +899,129 @@ pub fn extract_message_secret(message: &wa::Message) -> Option<&[u8]> { #[cfg(test)] mod tests { + + /// Generating a token must parse the message once. Collecting twice would + /// re-materialise every nested field's owned buffer, which costs more than + /// the concatenation this change removes: the streaming path would then be + /// slower than what it replaced for exactly the messages that carry media + /// or quoted text. + #[test] + fn generating_a_token_collects_the_pieces_once() { + let message = wa::Message { + extended_text_message: buffa::MessageField::some(wa::message::ExtendedTextMessage { + text: Some("nested content owns its bytes".to_string()), + ..Default::default() + }), + ..Default::default() + }; + let encoded = waproto::codec::message_to_vec(&message); + let sender: Jid = "5511987650001@s.whatsapp.net".parse().expect("sender"); + let remote: Jid = "5511987650002@s.whatsapp.net".parse().expect("remote"); + + COLLECT_CALLS.with(|c| c.set(0)); + let result = generate_reporting_token_from_encoded( + &message, + &encoded, + "3EB0ABCDEF", + &sender, + &remote, + None, + ); + assert!(result.is_some(), "the message must produce a token"); + // Two: once for the message, once for the nested field, which the + // collector re-frames through the same path. What must not happen is + // four, which is what collecting again to feed the hmac would cost. + assert_eq!( + COLLECT_CALLS.with(|c| c.get()), + 2, + "the pieces must be collected once and reused for the hmac" + ); + } + + /// The streaming HMAC must hash exactly the bytes the concatenation would + /// have held. The token goes on the wire and is verified by the server, so + /// a divergence here is not a performance bug, it is a message the peer + /// rejects. + #[test] + fn streaming_and_concatenating_produce_the_same_token() { + let key = [0x5au8; REPORTING_TOKEN_KEY_SIZE]; + let cases: Vec<(&str, Vec)> = vec![ + ("flat text field", { + let m = wa::Message { + conversation: Some("hello reporting".to_string()), + ..Default::default() + }; + waproto::codec::message_to_vec(&m) + }), + ("nested field", { + let m = wa::Message { + extended_text_message: buffa::MessageField::some( + wa::message::ExtendedTextMessage { + text: Some("nested body".to_string()), + ..Default::default() + }, + ), + ..Default::default() + }; + waproto::codec::message_to_vec(&m) + }), + // Two whitelisted fields, so the pieces have an order to get wrong. + // With a single piece the concatenation is trivially the same + // whatever order it is fed in, and this test would prove nothing. + ("two fields, so order matters", { + let m = wa::Message { + conversation: Some("first by field number".to_string()), + extended_text_message: buffa::MessageField::some( + wa::message::ExtendedTextMessage { + text: Some("sixth by field number".to_string()), + ..Default::default() + }, + ), + ..Default::default() + }; + waproto::codec::message_to_vec(&m) + }), + ("multibyte payload", { + let m = wa::Message { + conversation: Some("olá 🌍 ünïcode".repeat(4)), + ..Default::default() + }; + waproto::codec::message_to_vec(&m) + }), + ]; + + for (name, encoded) in cases { + let concatenated = extract_reporting_token_content(&encoded, REPORTING_FIELDS) + .unwrap_or_else(|| panic!("{name}: the case must extract something")); + let expected = calculate_reporting_token(&key, &concatenated) + .unwrap_or_else(|_| panic!("{name}: hmac over the concatenation")); + let pieces = collect_reporting_token_pieces(&encoded, REPORTING_FIELDS) + .unwrap_or_else(|| panic!("{name}: the case must collect something")); + let streamed = calculate_reporting_token_over_pieces(&key, &encoded, &pieces) + .unwrap_or_else(|| panic!("{name}: hmac over the pieces")); + assert_eq!(streamed, expected, "{name}: the token bytes are contract"); + } + } + + /// Bad path: a message with nothing whitelisted must decline in both + /// spellings rather than hashing an empty input, which would produce a + /// perfectly valid token for no content. + #[test] + fn a_message_with_nothing_to_report_produces_no_token() { + let key = [0x5au8; REPORTING_TOKEN_KEY_SIZE]; + let encoded = waproto::codec::message_to_vec(&wa::Message::default()); + + assert!(extract_reporting_token_content(&encoded, REPORTING_FIELDS).is_none()); + assert!(collect_reporting_token_pieces(&encoded, REPORTING_FIELDS).is_none()); + + // The collector is what stops the HMAC from ever seeing an empty list, + // so the HMAC is asked directly: a token over no content is a token + // that says nothing, and it would still verify. + assert!( + calculate_reporting_token_over_pieces(&key, &encoded, &[]).is_none(), + "hashing no pieces must not mint a token" + ); + } use super::*; #[test] diff --git a/wacore/src/send/dm.rs b/wacore/src/send/dm.rs index b8d7d9748..d97d39b56 100644 --- a/wacore/src/send/dm.rs +++ b/wacore/src/send/dm.rs @@ -166,10 +166,10 @@ pub async fn prepare_dm_stanza( let phash = resolved_devices.phash(); // Splice the shared content into the recipient plaintext and, when present, the - // own-device DeviceSentMessage plaintext. With no own companion devices (e.g. an - // account with nothing else linked), the DSM plaintext — and the destination-jid - // stringify it needs — would be built only to go unused, so encode just the recipient. - // The mci-hoist path re-encodes via `encode_dm_plaintexts` (see `shared_content`). + // own-device DeviceSentMessage plaintext. With no own companion devices (an + // account with nothing else linked), the DSM plaintext would be built only to + // go unused, so encode just the recipient. The mci-hoist path re-encodes via + // `encode_dm_plaintexts` (see `shared_content`). let crate::messages::DmPlaintexts { recipient: recipient_plaintext, own_devices: own_devices_plaintext, @@ -178,14 +178,10 @@ pub async fn prepare_dm_stanza( recipient: MessageUtils::pad_with_context_from_encoded(content, extra_context.as_ref()), own_devices: Vec::new(), }, - Some(content) => MessageUtils::dm_plaintexts_from_encoded( - content, - extra_context.as_ref(), - &to_jid.to_string(), - ), - None => { - MessageUtils::encode_dm_plaintexts(message, extra_context.as_ref(), &to_jid.to_string()) + Some(content) => { + MessageUtils::dm_plaintexts_from_encoded(content, extra_context.as_ref(), to_jid) } + None => MessageUtils::encode_dm_plaintexts(message, extra_context.as_ref(), to_jid), }; let mut participant_nodes = Vec::with_capacity(total_devices);