diff --git a/Cargo.lock b/Cargo.lock index ab72306bd..e7402908c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3959,6 +3959,7 @@ dependencies = [ "serde_json", "sha1 0.11.0", "sha2 0.11.0", + "smallvec", "smoothutf8", "subtle", "thiserror 2.0.19", diff --git a/Cargo.toml b/Cargo.toml index a7e147ed8..6491f9b3a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -74,28 +74,26 @@ unexpected_cfgs = { level = "warn", check-cfg = ["cfg(docsrs)"] } unused_qualifications = "warn" [workspace.dependencies] -# Shared dependencies aes = "0.9.1" aes-gcm = { version = "0.11.0", default-features = false, features = ["aes", "alloc", "bytes"] } anyhow = { version = "1.0", default-features = false } async-channel = { version = "2.5.0", default-features = false, features = ["std"] } -zerocopy = { version = "0.8.55", default-features = false, features = ["derive"] } async-lock = { version = "3", default-features = false } async-trait = "0.1.91" base64 = { version = "0.22.1", default-features = false, features = ["alloc"] } +bon = { version = "3.9.3", default-features = false, features = ["std"] } buffa = { version = "0.9.1", default-features = false, features = ["json"] } buffa-build = { version = "0.9.1", default-features = false } buffa-descriptor = { version = "0.9.1", default-features = false } bytemuck = { version = "1.25", default-features = false } bytes = { version = "1.12", default-features = false } -bon = { version = "3.9.3", default-features = false, features = ["std"] } cbc = { version = "0.2", features = ["alloc"] } chrono = { version = "0.4", default-features = false } compact_str = { version = "0.10", default-features = false } ctr = { version = "0.10", default-features = false } -divan = { package = "codspeed-divan-compat", version = "5.0.1" } diesel = { version = "2.3.11", default-features = false, features = ["sqlite", "r2d2", "32-column-tables"] } diesel_migrations = { version = "2.3.2", default-features = false, features = ["sqlite"] } +divan = { package = "codspeed-divan-compat", version = "5.0.1" } env_logger = { version = "0.11", default-features = false } event-listener = { version = "5", default-features = false } flate2 = { version = "1.1.9", default-features = false, features = ["zlib-rs"] } @@ -116,6 +114,8 @@ serde-big-array = "0.5" serde_json = { version = "1.0", default-features = false } sha1 = { version = "0.11.0", default-features = false } sha2 = { version = "0.11.0", default-features = false } +# Shared dependencies +smallvec = "1.15" smoothutf8 = { version = "0.2.3", default-features = false } subtle = { version = "2.6", default-features = false } thiserror = "2.0.19" @@ -132,6 +132,7 @@ wacore-libsignal = { path = "./wacore/libsignal", version = "0.6.0" } wacore-noise = { path = "./wacore/noise", version = "0.6.0" } waproto = { path = "./waproto", version = "0.6.0" } yoke = { version = "0.8", features = ["derive"] } +zerocopy = { version = "0.8.55", default-features = false, features = ["derive"] } zlib-rs = { version = "0.6.6", default-features = false, features = ["std", "rust-allocator"] } [features] @@ -202,9 +203,9 @@ async-channel = { workspace = true } async-lock = { workspace = true } async-trait = { workspace = true } base64 = { workspace = true } +bon = { workspace = true, optional = true } buffa = { workspace = true } bytes = { workspace = true } -bon = { workspace = true, optional = true } chrono = { workspace = true, features = ["clock"] } event-listener = { workspace = true } futures = { workspace = true, features = ["std"] } diff --git a/src/client/app_state.rs b/src/client/app_state.rs index 69600900e..cc20c9a31 100644 --- a/src/client/app_state.rs +++ b/src/client/app_state.rs @@ -1110,11 +1110,12 @@ impl Client { let result = async { self.ensure_e2e_sessions(std::slice::from_ref(&peer)) .await?; + let request_id = self.generate_message_id(); self.send_message_impl( peer, msg, crate::send::SendPipelineOptions { - request_id: Some(self.generate_message_id()), + request_id: Some(&request_id), peer: true, ..Default::default() }, diff --git a/src/client/context_impl.rs b/src/client/context_impl.rs index b04218a88..c90fcbf12 100644 --- a/src/client/context_impl.rs +++ b/src/client/context_impl.rs @@ -70,11 +70,6 @@ impl SendContextResolver for Client { } // Reuse the DM path's helpers so both lock the identical per-device mutexes. let keys = self.build_session_lock_keys(device_jids).await; - let mutexes = self.session_mutexes_for(&keys).await; - let mut guards = Vec::with_capacity(mutexes.len()); - for mutex in &mutexes { - guards.push(mutex.lock_arc().await); - } - SessionLockGuard::hold(Box::new(guards)) + SessionLockGuard::hold(Box::new(self.session_guards_for(&keys).await)) } } diff --git a/src/client/messaging.rs b/src/client/messaging.rs index a27749848..3f2a0c39a 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -197,7 +197,7 @@ impl Client { &edit_container_message, crate::send::SendPipelineOptions { edit: Some(crate::types::message::EditAttribute::MessageEdit), - request_id, + request_id: request_id.as_deref(), borrowed_message_id, ..Default::default() }, diff --git a/src/client/tests.rs b/src/client/tests.rs index 60cf8977c..93c47afd0 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -1478,7 +1478,7 @@ async fn cleanup_connection_state_flushes_dirty_signal_state() { let client = create_offline_sync_test_client().await; // A dirty identity lives only in the write-back cache until flushed. - let addr = ProtocolAddress::new("5550001000@s.whatsapp.net".to_string(), 1u32.into()); + let addr = ProtocolAddress::new("5550001000@s.whatsapp.net", 1u32.into()); client.signal_cache.put_identity(&addr, &[7u8; 32]).await; client.cleanup_connection_state().await; @@ -1579,7 +1579,7 @@ async fn cleanup_connection_state_keeps_state_when_flush_fails() { // A malformed identity (not 32 bytes) makes flush() error out, standing // in for a transient backend write failure during cleanup. - let bad = ProtocolAddress::new("5550002000@s.whatsapp.net".to_string(), 1u32.into()); + let bad = ProtocolAddress::new("5550002000@s.whatsapp.net", 1u32.into()); client.signal_cache.put_identity(&bad, &[0u8; 16]).await; // A valid dirty sender key that must not be dropped when the flush fails. diff --git a/src/features/groups.rs b/src/features/groups.rs index 51a53a480..88285e5f9 100644 --- a/src/features/groups.rs +++ b/src/features/groups.rs @@ -1400,7 +1400,7 @@ impl<'a> Groups<'a> { group_jid.clone(), &msg, crate::send::SendPipelineOptions { - request_id: Some(message_id.clone()), + request_id: Some(&message_id), extra_stanza_nodes: meta.into_iter().collect(), ..Default::default() }, diff --git a/src/features/signal.rs b/src/features/signal.rs index 157cf83e7..86a2b0c93 100644 --- a/src/features/signal.rs +++ b/src/features/signal.rs @@ -649,11 +649,7 @@ impl<'a> Signal<'a> { // Acquire per-device session locks before encrypting (matches DM send path) let lock_jids = self.client.build_session_lock_keys(&device_jids).await; - let session_mutexes = self.client.session_mutexes_for(&lock_jids).await; - let mut _session_guards = Vec::with_capacity(session_mutexes.len()); - for mutex in &session_mutexes { - _session_guards.push(mutex.lock().await); - } + let _session_guards = self.client.session_guards_for(&lock_jids).await; let plaintext = MessageUtils::encode_and_pad(message); let mut adapter = self.client.signal_adapter().await; diff --git a/src/message/special.rs b/src/message/special.rs index 3cefed74b..34bb2be34 100644 --- a/src/message/special.rs +++ b/src/message/special.rs @@ -354,7 +354,7 @@ impl Client { requester.clone(), message, crate::send::SendPipelineOptions { - request_id: Some(message_id.to_owned()), + request_id: Some(message_id), peer: true, ..Default::default() }, diff --git a/src/pdo.rs b/src/pdo.rs index 42b0da3cc..be0798681 100644 --- a/src/pdo.rs +++ b/src/pdo.rs @@ -273,7 +273,7 @@ impl Client { to, msg, crate::send::SendPipelineOptions { - request_id: Some(msg_id.clone()), + request_id: Some(&msg_id), peer: true, ..Default::default() }, diff --git a/src/send/mod.rs b/src/send/mod.rs index 5c423f33b..596fed549 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -169,7 +169,7 @@ struct SendBranchOutput { struct GroupBranchRequest<'a> { to: Jid, message: &'a wa::Message, - request_id: String, + request_id: &'a str, force_key_distribution: bool, edit: Option, extra_stanza_nodes: &'a [Node], @@ -181,7 +181,7 @@ struct GroupBranchRequest<'a> { struct DmBranchRequest<'a> { to: Jid, message: &'a wa::Message, - request_id: String, + request_id: &'a str, sent_at: SendInstant, edit: Option, extra_stanza_nodes: Vec, @@ -396,11 +396,14 @@ impl SendInstant { } #[derive(Default)] -pub(crate) struct SendPipelineOptions { +pub(crate) struct SendPipelineOptions<'a> { /// Instant this operation is stamped with, when the caller already sampled /// one. `None` makes [`Client::send_message_impl`] sample its own. pub(crate) sent_at: Option, - pub(crate) request_id: Option, + /// Borrowed on purpose: the caller that already owns an id (because it + /// returns it, or stamped state with it) lends it for the whole send + /// instead of handing over a copy. + pub(crate) request_id: Option<&'a str>, pub(crate) peer: bool, pub(crate) force_key_distribution: bool, pub(crate) edit: Option, @@ -918,11 +921,10 @@ impl Client { Some(id) => id, None => self.generate_message_id_at(sent_at.unix_secs_u64()), }; - // Both paths below consume `to` and `request_id`, so save copies for the result. - let result = SendResult { - message_id: request_id.clone(), - to: to.clone(), - }; + // Both paths below consume `to`, so save a copy for the result. The id + // is not copied: it is lent to the pipeline as `&str` and moved into + // the result once the send returns. + let result_to = to.clone(); // Newsletters are not E2E encrypted — send as plaintext via SMAX stanza. // Matches WA Web's OutMessagePublishNewsletterRequest + ContentType mixins. @@ -949,7 +951,10 @@ impl Client { .children(children) .build(); self.send_node(stanza).await?; - return Ok(result); + return Ok(SendResult { + message_id: request_id, + to: result_to, + }); } let (edit, inferred_meta) = infer_stanza_metadata(&message); @@ -966,7 +971,7 @@ impl Client { &message, SendPipelineOptions { sent_at: Some(sent_at), - request_id: Some(request_id), + request_id: Some(&request_id), edit, extra_stanza_nodes: extra_nodes, stanza_type: stanza_type_override, @@ -977,7 +982,10 @@ impl Client { ) .await .map_err(SendError::from_anyhow)?; - Ok(result) + Ok(SendResult { + message_id: request_id, + to: result_to, + }) } /// Send a status/story update using sender-key encryption. @@ -1617,7 +1625,7 @@ impl Client { &self, to: Jid, message: &wa::Message, - options: SendPipelineOptions, + options: SendPipelineOptions<'_>, ) -> Result<(), anyhow::Error> { let SendPipelineOptions { sent_at, @@ -1635,7 +1643,7 @@ impl Client { // rest sample here so the pipeline below still has exactly one. let sent_at = sent_at.unwrap_or_else(SendInstant::now); validate_extra_stanza_nodes(&extra_stanza_nodes)?; - if request_id_override.as_ref().is_some_and(String::is_empty) { + if request_id_override.is_some_and(str::is_empty) { return Err(SendError::InvalidRequest("message ID must not be empty".into()).into()); } // Newsletters are plaintext channels and never use the E2E path. Text @@ -1675,16 +1683,18 @@ impl Client { (to, false) }; - // Generate request ID early (doesn't need lock) - let request_id = match request_id_override { + // Generate request ID early (doesn't need lock). This frame owns the + // only copy for the whole send: the branch builders, the phash waiter + // and the messageSecret persistence all borrow it, so a send names its + // message exactly once no matter how many stages read that name. + let generated_request_id; + let request_id: &str = match request_id_override { Some(id) => id, - None => self.generate_message_id_at(sent_at.unix_secs_u64()), + None => { + generated_request_id = self.generate_message_id_at(sent_at.unix_secs_u64()); + &generated_request_id + } }; - // `request_id` is moved into the branch-specific stanza builders below; - // keep a copy for the post-send messageSecret persistence (the secret - // itself is generated inside prepare_dm/group_stanza, not on `message`, - // so it's threaded back out via PreparedStanza.message_secret below). - let outbound_id_clone = request_id.clone(); let tc_issue_target = to.clone(); // Dispatch to a concrete boxed future per branch: this function's own @@ -1745,18 +1755,25 @@ impl Client { // Registered before the stanza goes out: the ack can arrive while // send_node is still returning, and a waiter installed afterwards would // miss it. - let ack_message_id = if !borrowed_message_id - && let Some(phash) = dm_phash - && let Some(msg_id) = stanza_to_send - .attrs() - .optional_string("id") - .map(|s| s.into_owned()) - { + // Keying the waiter off `request_id` rather than re-reading the stanza + // is only sound while every branch stamps the id it was handed; assert + // that instead of paying an owned copy of an attribute we already have. + debug_assert_eq!( + stanza_to_send.attrs().optional_string("id").as_deref(), + Some(request_id), + "branch stanza must carry the id this send was named with" + ); + let ack_message_id = if !borrowed_message_id && let Some(phash) = dm_phash { // Group sends also invalidate group cache on mismatch: the server's // participant set diverged, so the next send needs a fresh query. let invalidate_group = tc_issue_target.is_group(); - self.register_phash_waiter(&msg_id, phash, tc_issue_target.clone(), invalidate_group); - Some(msg_id) + self.register_phash_waiter( + request_id, + phash, + tc_issue_target.clone(), + invalidate_group, + ); + Some(request_id) } else { None }; @@ -1772,7 +1789,7 @@ impl Client { } if let Err(e) = self.send_node(stanza_to_send).await { - if let Some(ref msg_id) = ack_message_id { + if let Some(msg_id) = ack_message_id { self.response_waiters_guard().remove(msg_id); } return Err(e.into()); @@ -1791,7 +1808,7 @@ impl Client { self.persist_outbound_msg_secret( &tc_issue_target, &sender, - &outbound_id_clone, + request_id, secret, class, sent_at, @@ -1834,7 +1851,7 @@ impl Client { &self, to: Jid, message: &wa::Message, - request_id: String, + request_id: &str, ) -> Result { let node = { // Peer messages are only valid for individual users, not groups @@ -1913,7 +1930,7 @@ impl Client { // the id is borrowed: it would replace the original message's // retry-cache entry, so a retry receipt for it returns this edit. if !borrowed_message_id { - self.add_recent_message(&to, &request_id, message, shared_content.clone()) + self.add_recent_message(&to, request_id, message, shared_content.clone()) .await; } @@ -2123,7 +2140,7 @@ impl Client { account: account_info.as_deref(), to: &to, message, - message_id: &request_id, + message_id: request_id, force_distribution: force_skdm, distribution_targets: skdm_target_devices, distribution_policy: wacore::send::SenderKeyDistributionPolicy::BestEffort, @@ -2203,7 +2220,7 @@ impl Client { account: account_info.as_deref(), to: &to, message, - message_id: &request_id, + message_id: request_id, force_distribution: retry_force, distribution_targets: retry_targets, distribution_policy: @@ -2276,13 +2293,13 @@ impl Client { if is_status_addon { self.add_recent_message( &Jid::status_broadcast(), - &request_id, + request_id, message, shared_content.clone(), ) .await; } else { - self.add_recent_message(&to, &request_id, message, shared_content.clone()) + self.add_recent_message(&to, request_id, message, shared_content.clone()) .await; } } @@ -2361,11 +2378,7 @@ impl Client { } let lock_jids = self.build_session_lock_keys(dm_devices.devices()).await; - let _session_mutexes = self.session_mutexes_for(&lock_jids).await; - let mut _session_guards = Vec::with_capacity(_session_mutexes.len()); - for mutex in &_session_mutexes { - _session_guards.push(mutex.lock().await); - } + let _session_guards = self.session_guards_for(&lock_jids).await; let mut store_adapter = self.signal_adapter().await; @@ -2380,7 +2393,7 @@ impl Client { account: device_snapshot.account.as_deref(), to: &stanza_to, message, - message_id: &request_id, + message_id: request_id, edit: edit.as_ref(), extra_nodes: &extra_stanza_nodes, devices: &dm_devices, @@ -2465,16 +2478,57 @@ impl Client { keys } - /// Fetch per-device session mutexes in deadlock-free order. + /// Take every per-device session lock, in `jids` order. + /// + /// INVARIANT: acquisition order IS `jids` order, and callers pass keys from + /// [`Self::build_session_lock_keys`], which sorts them. That single order is + /// what keeps two sends overlapping on a device from deadlocking, so a + /// change here has to preserve it. + /// + /// Each mutex is locked as it is resolved rather than resolving the whole + /// set first: the handles exist only to be locked, so the vector holding + /// them was pure staging. The guards themselves must still be collected — + /// they are what keeps the locks held for the caller's scope. + pub(crate) async fn session_guards_for( + &self, + jids: &[Jid], + ) -> Vec> { + // A duplicate key would have this loop await a lock it already holds, + // which is a silent self-deadlock rather than a panic: the send just + // never returns. Every caller goes through `build_session_lock_keys`, + // which sorts and dedups, so this only fires if a future path forgets + // to. + debug_assert!( + jids.windows(2).all(|pair| pair[0] != pair[1]), + "session lock keys must be deduped before acquisition, or the loop deadlocks on itself" + ); + + let mut guards = Vec::with_capacity(jids.len()); + // A `ProtocolAddress` IS the "{name}.0" string the lock map is keyed by, + // and it holds it inline, so the whole loop names its keys without + // allocating a formatting buffer. + let mut addr = wacore::types::jid::make_reusable_protocol_address(); + for jid in jids { + jid.reset_protocol_address(&mut addr); + let mutex = self.session_lock_for(addr.as_str()).await; + guards.push(mutex.lock_arc().await); + } + guards + } + + /// The mutexes [`Self::session_guards_for`] would take, without taking + /// them. Only tests need this: production code always wants the guards, and + /// resolving handles it does not lock is what this commit removed. + #[cfg(test)] pub(crate) async fn session_mutexes_for( &self, jids: &[Jid], ) -> Vec>> { let mut mutexes = Vec::with_capacity(jids.len()); - let mut buf = wacore::types::jid::make_address_buffer(); + let mut addr = wacore::types::jid::make_reusable_protocol_address(); for jid in jids { - wacore::types::jid::write_protocol_address_to(jid, &mut buf); - mutexes.push(self.session_lock_for(&buf).await); + jid.reset_protocol_address(&mut addr); + mutexes.push(self.session_lock_for(addr.as_str()).await); } mutexes } @@ -4541,6 +4595,88 @@ mod tests { "100000012345678@lid.0" ); } + + /// Every key handed in ends up locked, and every one is released when + /// the guards are dropped. The device counts are the three the DM path + /// actually produces: none (a fan-out that resolved to nothing), one + /// (a steady 1:1) and several (companion devices in play). + #[tokio::test] + async fn taking_guards_locks_every_key_and_releasing_them_frees_every_key() { + let client = crate::test_utils::create_test_client_with_name("guards_cover").await; + + for count in [0usize, 1, 3] { + let devices: Vec = (0..count) + .map(|i| Jid::from_str(&format!("10000001234567{i}:5@lid")).unwrap()) + .collect(); + let keys = client.build_session_lock_keys(&devices).await; + assert_eq!(keys.len(), count, "one key per device at count {count}"); + let mutexes = client.session_mutexes_for(&keys).await; + + let guards = client.session_guards_for(&keys).await; + assert_eq!(guards.len(), count, "one guard per key at count {count}"); + for (i, mutex) in mutexes.iter().enumerate() { + assert!( + mutex.try_lock().is_none(), + "key {i} of {count} must be held while the guards live" + ); + } + + drop(guards); + for (i, mutex) in mutexes.iter().enumerate() { + assert!( + mutex.try_lock().is_some(), + "key {i} of {count} must be free once the guards are dropped" + ); + } + } + } + + /// The keys are locked in the order given, which is the sorted order + /// `build_session_lock_keys` produces. That single global order is the + /// only thing keeping two sends that overlap on a device from + /// deadlocking, so acquiring out of order must be observable. + /// + /// Blocking the SECOND key and then waiting for the FIRST to become + /// contended is what pins the order down: a taker that went second-first + /// would park on the blocked key and never touch the first one. + #[tokio::test] + async fn keys_are_locked_in_the_order_they_are_given() { + let client = crate::test_utils::create_test_client_with_name("guards_order").await; + + let devices: Vec = ["100000012345670:5@lid", "100000012345671:5@lid"] + .iter() + .map(|s| Jid::from_str(s).unwrap()) + .collect(); + let keys = client.build_session_lock_keys(&devices).await; + assert_eq!(keys.len(), 2); + let mutexes = client.session_mutexes_for(&keys).await; + + let blocker = mutexes[1].lock_arc().await; + + let mut taker = tokio::spawn({ + let client = client.clone(); + let keys = keys.clone(); + async move { client.session_guards_for(&keys).await.len() } + }); + + // Bounded work, not a deadline: yield until the first key is taken. + let mut polls = 0; + while mutexes[0].try_lock().is_some() { + polls += 1; + assert!( + polls < 10_000, + "the first key was never taken, so acquisition did not start there" + ); + tokio::task::yield_now().await; + } + assert!( + futures::poll!(&mut taker).is_pending(), + "the taker must still be parked on the second key" + ); + + drop(blocker); + assert_eq!(taker.await.expect("taker finishes"), 2); + } } // ---- outbound messageSecret capture --------------------------------- @@ -4657,7 +4793,7 @@ mod tests { peer, &msg, SendPipelineOptions { - request_id: Some(request_id.to_string()), + request_id: Some(request_id), peer: true, ..Default::default() }, @@ -4710,7 +4846,7 @@ mod tests { peer, &msg, SendPipelineOptions { - request_id: Some(request_id.to_string()), + request_id: Some(request_id), peer: true, stanza_type: Some(StanzaType::Poll), ..Default::default() @@ -4875,7 +5011,7 @@ mod tests { peer_pn, &msg, SendPipelineOptions { - request_id: Some(request_id.to_string()), + request_id: Some(request_id), ..Default::default() }, ) @@ -4951,7 +5087,7 @@ mod tests { peer_pn.clone(), &msg, SendPipelineOptions { - request_id: Some(request_id.to_string()), + request_id: Some(request_id), ..Default::default() }, ) @@ -5410,6 +5546,204 @@ mod tests { }; assert_eq!(prepared.message_secret.as_ref().unwrap().len(), 32); } + + /// A send names its message once and every downstream stage reads that same + /// name: the wire stanza, the phash ack-waiter, the outbound messageSecret + /// and the returned `SendResult`. A non-ASCII id is used on purpose — a + /// truncating or byte-indexing copy anywhere in that chain would show up + /// here and nowhere else. + #[tokio::test] + async fn one_id_names_the_stanza_the_waiter_the_secret_and_the_result() { + let (client, _transport) = crate::test_utils::create_iq_test_client().await; + let (peer_pn, _peer_lid) = seed_dm_wire_namespace_state(&client).await; + + let message_id = "ID_ünïcødé_✅_ONE"; + let result = client + .send_message_with_options( + peer_pn.clone(), + wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }, + SendOptions::default().with_message_id(message_id), + ) + .await + .expect("connected test client should complete the send"); + + assert_eq!( + result.message_id, message_id, + "result carries the caller id" + ); + assert_eq!(result.to, peer_pn, "result carries the caller target"); + + let waiters = client.response_waiters_guard(); + assert!( + waiters.contains_key(message_id), + "the phash ack-waiter must be keyed by the send's own id" + ); + assert_eq!(waiters.len(), 1, "no second entry under another spelling"); + drop(waiters); + + let secret = client.msg_secret_buffer.lookup( + &peer_pn.to_non_ad_string(), + &client.pn().expect("own pn").to_non_ad_string(), + message_id, + ); + assert!( + secret.is_some(), + "the outbound messageSecret must be bound to the same id" + ); + } + + /// The waiter is installed before the stanza reaches the socket (a fast ack + /// can land while `send_node` is still returning), so a send that fails on + /// the wire has to take it back out — under the id it registered. Removing + /// under anything else leaks an entry that a later ack could resolve. + #[tokio::test] + async fn a_failed_send_takes_its_phash_waiter_back_out() { + let client = crate::test_utils::create_test_client_with_name("phash_waiter_rollback").await; + let (peer_pn, _peer_lid) = seed_dm_wire_namespace_state(&client).await; + + let message_id = "ID_ünïcødé_✅_ROLLBACK"; + let result = client + .send_message_impl( + peer_pn, + &wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }, + SendPipelineOptions { + request_id: Some(message_id), + ..Default::default() + }, + ) + .await; + assert!(result.is_err(), "no socket: the send must fail on the wire"); + assert_eq!( + client.response_waiters_guard().len(), + 0, + "a failed send must leave no waiter behind, under any key" + ); + } + + /// A borrowed id belongs to another message: registering a waiter under it + /// would overwrite the original send's waiter, and binding a secret under it + /// would overwrite the original's secret. + #[tokio::test] + async fn a_borrowed_id_registers_no_waiter_and_binds_no_secret() { + let (client, _transport) = crate::test_utils::create_iq_test_client().await; + let (peer_pn, _peer_lid) = seed_dm_wire_namespace_state(&client).await; + + let message_id = "ID_BORROWED_1"; + client + .send_message_impl( + peer_pn.clone(), + &wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }, + SendPipelineOptions { + request_id: Some(message_id), + borrowed_message_id: true, + ..Default::default() + }, + ) + .await + .expect("connected test client should complete the send"); + + assert_eq!( + client.response_waiters_guard().len(), + 0, + "a borrowed id must not claim the waiter slot" + ); + assert!( + client + .msg_secret_buffer + .lookup( + &peer_pn.to_non_ad_string(), + &client.pn().expect("own pn").to_non_ad_string(), + message_id, + ) + .is_none(), + "a borrowed id must not claim the secret slot" + ); + } + + /// An empty id would name nothing: it must be refused at both entry points + /// before any state is stamped with it. + #[tokio::test] + async fn an_empty_id_is_refused_at_both_entry_points() { + let client = crate::test_utils::create_test_client_with_name("empty_send_id").await; + let peer: Jid = "100000000000777@s.whatsapp.net".parse().unwrap(); + let msg = wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }; + + let public = client + .send_message_with_options( + peer.clone(), + msg.clone(), + SendOptions::default().with_message_id(""), + ) + .await; + assert!( + matches!(public, Err(SendError::InvalidRequest(_))), + "public send must reject an empty id, got {public:?}" + ); + + let internal = client + .send_message_impl( + peer, + &msg, + SendPipelineOptions { + request_id: Some(""), + ..Default::default() + }, + ) + .await; + let internal = internal.expect_err("internal send must reject an empty id"); + assert!( + internal + .to_string() + .contains("message ID must not be empty"), + "unexpected error: {internal}" + ); + } + + /// The plaintext newsletter branch returns before the E2E pipeline, so it + /// builds its own result; it must still hand back the id it stamped and the + /// channel it addressed. + #[tokio::test] + async fn the_newsletter_branch_returns_the_id_and_target_it_stamped() { + let (client, _transport) = crate::test_utils::create_iq_test_client().await; + let channel: Jid = "123456789@newsletter".parse().unwrap(); + + let message_id = "ID_ünïcødé_✅_NEWS"; + let waiter = client + .wait_for_sent_node(crate::client::NodeFilter::tag("message").attr("id", message_id)); + let result = client + .send_message_with_options( + channel.clone(), + wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }, + SendOptions::default().with_message_id(message_id), + ) + .await + .expect("newsletter send is plaintext and needs no session"); + + assert_eq!(result.message_id, message_id); + assert_eq!(result.to, channel); + + let node = waiter.await.expect("the stanza should be captured"); + assert_eq!( + node.attrs().optional_string("id").as_deref(), + Some(message_id), + "the wire id must be the same one the result reports" + ); + } } #[cfg(test)] diff --git a/src/signal_flush.rs b/src/signal_flush.rs index f92b2152d..fe8f8c94f 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -266,7 +266,7 @@ mod tests { } fn dirty_session(client: &Arc, user: &str) -> ProtocolAddress { - let addr = ProtocolAddress::new(user.to_string(), 1.into()); + let addr = ProtocolAddress::new(user, 1.into()); assert!( client .signal_cache @@ -306,7 +306,7 @@ mod tests { wait_for_backend_session(&client, &covered).await; // A raised lease must be durable when the call returns. - let addr = ProtocolAddress::new("15550003002".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550003002", 1.into()); let mut record = SessionRecord::new_fresh(); record.reserve_sender_chain_counters(0); assert!(client.signal_cache.try_put_session(&addr, record).is_ok()); @@ -504,7 +504,7 @@ mod tests { tokio::time::sleep(Duration::from_millis(2)).await; } // And B still makes progress: its dirty entry lands. - let addr = ProtocolAddress::new("15550007002".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550007002", 1.into()); wait_for_backend_session(&client, &addr).await; } diff --git a/src/store/signal.rs b/src/store/signal.rs index 55d5378d9..a39164943 100644 --- a/src/store/signal.rs +++ b/src/store/signal.rs @@ -572,7 +572,7 @@ mod tests { async fn direct_session_store_preserves_clean_reload_and_recovery_ceiling() { let backend = crate::test_utils::create_test_backend().await; let device = Device::new(backend.clone()); - let address = ProtocolAddress::new("15550001001".to_string(), 1.into()); + let address = ProtocolAddress::new("15550001001", 1.into()); SessionStore::store_session(&device, &address, &leased_session()) .await diff --git a/src/store/signal_adapter.rs b/src/store/signal_adapter.rs index 1992cd151..90d820254 100644 --- a/src/store/signal_adapter.rs +++ b/src/store/signal_adapter.rs @@ -512,7 +512,7 @@ mod tests { let cache = Arc::new(SignalStoreCache::new()); let adapter = SignalProtocolStoreAdapter::new(device, cache.clone()); - let addr = ProtocolAddress::new("bob".to_string(), 1.into()); + let addr = ProtocolAddress::new("bob", 1.into()); // The real path stores the promoted session before buffering the prekey. cache.put_session(&addr, SessionRecord::new_fresh()).await; adapter @@ -635,7 +635,7 @@ mod tests { let backend: Arc = Arc::new(InMemoryBackend::new()); let cache = Arc::new(SignalStoreCache::new()); - let address = ProtocolAddress::new("15550006666".to_string(), 1.into()); + let address = ProtocolAddress::new("15550006666", 1.into()); let (record, identity_pair) = outbound_session(); let expected = record.serialize().expect("serialize session"); cache.put_session(&address, record).await; @@ -717,7 +717,7 @@ mod tests { let cache = Arc::new(SignalStoreCache::new()); let device = Arc::new(RwLock::new(Device::new(backend))); let mut adapter = SignalProtocolStoreAdapter::new(device, cache.clone()); - let address = ProtocolAddress::new("15550005555".to_string(), 1.into()); + let address = ProtocolAddress::new("15550005555", 1.into()); cache .put_session(&address, SessionRecord::new_fresh()) .await; @@ -743,7 +743,7 @@ mod tests { async fn session_store_fast_paths_round_trip() { use wacore::libsignal::protocol::SessionStore as _; let mut adapter = test_adapter(); - let addr = ProtocolAddress::new("15550002222".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550002222", 1.into()); // Cold cache: goes through the async fallback (backend consult). assert!(!adapter.session_store.has_session(&addr).await.unwrap()); @@ -783,7 +783,7 @@ mod tests { async fn identity_fast_paths_keep_change_semantics() { use wacore::libsignal::protocol::{IdentityKeyPair, IdentityKeyStore as _}; let mut adapter = test_adapter(); - let addr = ProtocolAddress::new("15550003333".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550003333", 1.into()); let mut rng = rand::make_rng::(); let first = *IdentityKeyPair::generate(&mut rng).identity_key(); @@ -861,7 +861,7 @@ mod hook_alloc_tests { #[test] fn the_trusted_identity_hook_answers_without_allocating() { let adapter = adapter_for_test(); - let address = ProtocolAddress::new("bob@s.whatsapp.net".to_string(), 1.into()); + let address = ProtocolAddress::new("bob@s.whatsapp.net", 1.into()); let identity = some_identity(); // Through the resolver, not the hook directly: the point is that the @@ -888,7 +888,7 @@ mod hook_alloc_tests { #[test] fn the_save_identity_hook_declines_when_nothing_is_cached() { let mut adapter = adapter_for_test(); - let address = ProtocolAddress::new("never-seen@s.whatsapp.net".to_string(), 1.into()); + let address = ProtocolAddress::new("never-seen@s.whatsapp.net", 1.into()); let identity = some_identity(); assert!( @@ -906,7 +906,7 @@ mod hook_alloc_tests { #[tokio::test] async fn the_save_identity_hook_answers_once_the_entry_is_cached() { let mut adapter = adapter_for_test(); - let address = ProtocolAddress::new("bob@s.whatsapp.net".to_string(), 1.into()); + let address = ProtocolAddress::new("bob@s.whatsapp.net", 1.into()); let first = some_identity(); let second = some_identity(); @@ -939,7 +939,7 @@ mod hook_alloc_tests { #[test] fn the_session_hook_declines_when_the_cache_cannot_answer() { let adapter = adapter_for_test(); - let address = ProtocolAddress::new("never-seen@s.whatsapp.net".to_string(), 1.into()); + let address = ProtocolAddress::new("never-seen@s.whatsapp.net", 1.into()); assert!( adapter.session_store.try_has_session(&address).is_none(), diff --git a/src/voip/facade.rs b/src/voip/facade.rs index 90957714e..3ad471917 100644 --- a/src/voip/facade.rs +++ b/src/voip/facade.rs @@ -429,11 +429,7 @@ impl<'a> OutgoingCall<'a> { // in the shared pre-flight); hold the per-device session locks place_call's encrypt also // takes, so it can't clobber a concurrent send advancing the same session. let lock_jids = self.client.build_session_lock_keys(&devices).await; - let session_mutexes = self.client.session_mutexes_for(&lock_jids).await; - let mut session_guards = Vec::with_capacity(session_mutexes.len()); - for mutex in &session_mutexes { - session_guards.push(mutex.lock().await); - } + let session_guards = self.client.session_guards_for(&lock_jids).await; let mut would_pkmsg = Vec::with_capacity(devices.len()); for d in &devices { @@ -598,11 +594,7 @@ async fn place_call( // lock; concurrent ratchet mutations would corrupt session state. let raw = { let lock_jids = client.build_session_lock_keys(devices).await; - let session_mutexes = client.session_mutexes_for(&lock_jids).await; - let mut _session_guards = Vec::with_capacity(session_mutexes.len()); - for mutex in &session_mutexes { - _session_guards.push(mutex.lock().await); - } + let _session_guards = client.session_guards_for(&lock_jids).await; // Sessions were asserted upstream (`OutgoingCall::start`), so skip the network session-ensure and // encrypt against the existing sessions directly: a device whose session is somehow still diff --git a/tests/bench-integration/Cargo.toml b/tests/bench-integration/Cargo.toml index 6c2a58e7f..c393f6024 100644 --- a/tests/bench-integration/Cargo.toml +++ b/tests/bench-integration/Cargo.toml @@ -11,11 +11,11 @@ e2e-tests = { path = "../e2e" } env_logger = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread", "sync", "time"] } whatsapp-rust = { path = "../..", default-features = false, features = [ - "danger-skip-tls-verify", - "simd", - "tokio-runtime", - "tokio-native", - "signal", + "danger-skip-tls-verify", + "simd", + "tokio-runtime", + "tokio-native", + "signal", ] } [[bench]] diff --git a/tests/e2e/Cargo.toml b/tests/e2e/Cargo.toml index 4ec3e6569..523597e00 100644 --- a/tests/e2e/Cargo.toml +++ b/tests/e2e/Cargo.toml @@ -21,18 +21,18 @@ uuid = { workspace = true, features = ["v4"] } wacore = { path = "../../wacore", features = ["test-util"] } wacore-binary = { path = "../../wacore/binary" } whatsapp-rust = { path = "../..", default-features = false, features = [ - "danger-skip-cert-chain-verify", - "danger-skip-tls-verify", - "simd", - "tokio-runtime", - "tokio-native", - "signal", + "danger-skip-cert-chain-verify", + "danger-skip-tls-verify", + "simd", + "tokio-runtime", + "tokio-native", + "signal", ] } whatsapp-rust-tokio-transport = { path = "../../transports/tokio-transport", features = [ - "danger-skip-tls-verify", + "danger-skip-tls-verify", ] } whatsapp-rust-ureq-http-client = { path = "../../http_clients/ureq-client", features = [ - "danger-skip-tls-verify", + "danger-skip-tls-verify", ] } [dev-dependencies] diff --git a/tests/signal_durability_sqlite.rs b/tests/signal_durability_sqlite.rs index a7f494e68..245f4a1c7 100644 --- a/tests/signal_durability_sqlite.rs +++ b/tests/signal_durability_sqlite.rs @@ -56,7 +56,7 @@ impl SenderKeyStore for CachedSenderKeyStore<'_> { } fn dm_address() -> wacore::libsignal::protocol::ProtocolAddress { - wacore::libsignal::protocol::ProtocolAddress::new("15550008001".to_string(), 1.into()) + wacore::libsignal::protocol::ProtocolAddress::new("15550008001", 1.into()) } fn group_name() -> SenderKeyName { diff --git a/transports/tokio-transport/Cargo.toml b/transports/tokio-transport/Cargo.toml index 5b8b45544..bb7ca5343 100644 --- a/transports/tokio-transport/Cargo.toml +++ b/transports/tokio-transport/Cargo.toml @@ -25,10 +25,10 @@ rustls = { version = "0.23", default-features = false, features = ["ring"] } tokio = { workspace = true, features = ["macros", "rt", "sync"] } tokio-rustls = { version = "0.26.4", default-features = false, features = ["ring"] } tokio-websockets = { version = "0.13.3", features = [ - "client", - "rustls-bring-your-own-connector", - "rand", - "ring", + "client", + "rustls-bring-your-own-connector", + "rand", + "ring", ] } wacore = { workspace = true } webpki-roots = "1.0.9" diff --git a/wacore/Cargo.toml b/wacore/Cargo.toml index 2bd0ea59c..204910a19 100644 --- a/wacore/Cargo.toml +++ b/wacore/Cargo.toml @@ -49,7 +49,6 @@ dhat-heap = ["dep:dhat"] [dependencies] aes = { workspace = true } aes-gcm = { workspace = true, optional = true } -zerocopy = { workspace = true, optional = true } anyhow = { workspace = true } async-channel = { workspace = true } async-lock = { workspace = true } @@ -81,6 +80,7 @@ serde-big-array = { workspace = true } serde_json = { workspace = true, features = ["std"] } sha1 = { workspace = true } sha2 = { workspace = true } +smallvec = { workspace = true } smoothutf8 = { workspace = true } subtle = { workspace = true } thiserror = { workspace = true } @@ -92,6 +92,7 @@ wacore-derive = { workspace = true } wacore-libsignal = { workspace = true } wacore-noise = { workspace = true } waproto = { workspace = true } +zerocopy = { workspace = true, optional = true } [build-dependencies] buffa-build = { workspace = true } diff --git a/wacore/benches/send_receive_benchmark.rs b/wacore/benches/send_receive_benchmark.rs index 957ed5e6c..37a617a17 100644 --- a/wacore/benches/send_receive_benchmark.rs +++ b/wacore/benches/send_receive_benchmark.rs @@ -784,7 +784,7 @@ fn bench_dm_send(bencher: divan::Bencher) { d.bob_jid.clone(), &signal_addr, &d.msg, - "b-001".into(), + "b-001", None, )) .unwrap(); diff --git a/wacore/binary/Cargo.toml b/wacore/binary/Cargo.toml index d559d59e1..2e20ff587 100644 --- a/wacore/binary/Cargo.toml +++ b/wacore/binary/Cargo.toml @@ -24,7 +24,7 @@ bytes = { workspace = true } compact_str = { workspace = true } itoa = { workspace = true } serde = { workspace = true, optional = true } -smallvec = "1.15" +smallvec = { workspace = true } smoothutf8 = { workspace = true } stable_deref_trait = "1.2.1" yoke = { workspace = true } diff --git a/wacore/libsignal/benches/libsignal_benchmark.rs b/wacore/libsignal/benches/libsignal_benchmark.rs index 2a361432b..5d02507d2 100644 --- a/wacore/libsignal/benches/libsignal_benchmark.rs +++ b/wacore/libsignal/benches/libsignal_benchmark.rs @@ -305,7 +305,7 @@ impl User { .unwrap(); }); - let address = ProtocolAddress::new(name.to_string(), device_id.into()); + let address = ProtocolAddress::new(name, device_id.into()); Self { address, diff --git a/wacore/libsignal/src/core/address.rs b/wacore/libsignal/src/core/address.rs index 78c8fdd08..fde3195a9 100644 --- a/wacore/libsignal/src/core/address.rs +++ b/wacore/libsignal/src/core/address.rs @@ -172,15 +172,8 @@ impl fmt::Display for DeviceId { } } -const fn digit_count(n: u32) -> usize { - if n == 0 { - return 1; - } - n.ilog10() as usize + 1 -} - #[inline] -fn append_device_suffix(buf: &mut String, device_id: DeviceId) { +fn append_device_suffix(buf: &mut AddressBuf, device_id: DeviceId) { let id = u32::from(device_id); if id == 0 { buf.push_str(".0"); @@ -190,34 +183,162 @@ fn append_device_suffix(buf: &mut String, device_id: DeviceId) { } } +/// Longest address held without touching the heap. +/// +/// A real address is short: `"5511987650001:5@c.us.0"` is 22 bytes and the +/// longest server the JID enum can render adds ten more, so this covers every +/// address the protocol produces with room for drift. Anything longer still +/// works -- it spills to a `String`, which is where every address used to live +/// unconditionally. +const INLINE_CAPACITY: usize = 47; + +/// A protocol address's characters: inline while they fit, heap when they do +/// not. +/// +/// Which arm holds them is not part of the value. `ProtocolAddress` is a +/// `HashMap` key in the session cache, so equality, ordering and hashing all go +/// through [`AddressBuf::as_str`] and never look at the representation: the +/// same characters answer identically whether they were built inline or spilled. +#[derive(Clone)] +pub struct AddressBuf(AddressRepr); + +/// Prints what the buffer currently holds, never the array behind it. +/// +/// `clear` only rewinds the length, so the tail of an inline buffer still +/// holds the previous address's bytes. A derived `Debug` would print all of +/// the inline capacity and leak an unrelated peer's JID into any log line that +/// formats an address, or an error that embeds one. +impl fmt::Debug for AddressBuf { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Debug::fmt(self.as_str(), f) + } +} + +#[derive(Clone)] +enum AddressRepr { + Inline { + bytes: [u8; INLINE_CAPACITY], + len: u8, + }, + Heap(String), +} + +impl Default for AddressBuf { + #[inline] + fn default() -> Self { + Self::empty() + } +} + +impl AddressBuf { + /// An empty buffer, holding nothing on the heap. + #[inline] + pub fn empty() -> Self { + Self(AddressRepr::Inline { + bytes: [0; INLINE_CAPACITY], + len: 0, + }) + } + + #[inline] + pub fn as_str(&self) -> &str { + match &self.0 { + // Only whole `&str` fragments are ever appended, never split, so + // the inline bytes are valid UTF-8 by construction. + AddressRepr::Inline { bytes, len } => std::str::from_utf8(&bytes[..usize::from(*len)]) + .expect("address is built from whole str fragments"), + AddressRepr::Heap(buf) => buf, + } + } + + #[inline] + fn len(&self) -> usize { + match &self.0 { + AddressRepr::Inline { len, .. } => usize::from(*len), + AddressRepr::Heap(buf) => buf.len(), + } + } + + /// Empty the buffer, keeping whatever room it already has: an address that + /// once needed the heap is usually rewritten to something just as long. + #[inline] + pub fn clear(&mut self) { + match &mut self.0 { + AddressRepr::Inline { len, .. } => *len = 0, + AddressRepr::Heap(buf) => buf.clear(), + } + } + + pub fn push_str(&mut self, s: &str) { + match &mut self.0 { + AddressRepr::Inline { bytes, len } => { + let start = usize::from(*len); + let end = start + s.len(); + if end <= INLINE_CAPACITY { + bytes[start..end].copy_from_slice(s.as_bytes()); + *len = end as u8; + return; + } + let mut heap = String::with_capacity(end); + // Valid UTF-8 for the same reason as in `as_str`. + heap.push_str( + std::str::from_utf8(&bytes[..start]) + .expect("address is built from whole str fragments"), + ); + heap.push_str(s); + self.0 = AddressRepr::Heap(heap); + } + AddressRepr::Heap(buf) => buf.push_str(s), + } + } + + #[inline] + pub fn push(&mut self, c: char) { + self.push_str(c.encode_utf8(&mut [0u8; 4])); + } +} + +impl fmt::Write for AddressBuf { + #[inline] + fn write_str(&mut self, s: &str) -> fmt::Result { + self.push_str(s); + Ok(()) + } +} + /// Single-buffer protocol address. The buffer stores `"{name}.{device_id}"` and /// `name_len` marks where the name ends, so `name()` and `as_str()` are both -/// zero-cost slices. One String instead of two — halves allocation count for -/// one-shot construction and eliminates the copy in `reset_with()`. +/// zero-cost slices. One buffer instead of two — halves allocation count for +/// one-shot construction and eliminates the copy in `reset_with()`. The buffer +/// itself is inline up to 47 bytes, so the common address is a value +/// with nothing behind it. #[derive(Clone, Debug)] pub struct ProtocolAddress { - buf: String, + buf: AddressBuf, name_len: usize, device_id: DeviceId, } impl ProtocolAddress { - pub fn new(name: String, device_id: DeviceId) -> Self { - let name_len = name.len(); - let mut buf = name; - append_device_suffix(&mut buf, device_id); - Self { - buf, - name_len, - device_id, - } + pub fn new(name: &str, device_id: DeviceId) -> Self { + let mut address = Self::empty(device_id); + address.reset_with(|buf| buf.push_str(name)); + address } - /// Pre-allocated empty address. Call `reset_with()` to fill. - pub fn with_capacity(capacity: usize, device_id: DeviceId) -> Self { - let suffix_len = 1 + digit_count(u32::from(device_id)); + /// An address with no name yet, ready for [`Self::reset_with`]. No capacity + /// argument: the buffer starts inline and grows only if an address ever + /// exceeds it. + /// + /// The device suffix is written immediately even though the name is empty. + /// `Hash`, `Eq` and `Ord` all read the rendered string, so leaving the + /// buffer truly empty would make every unnamed address compare equal + /// regardless of its device, and two of them would collide as map keys. + pub fn empty(device_id: DeviceId) -> Self { + let mut buf = AddressBuf::empty(); + append_device_suffix(&mut buf, device_id); Self { - buf: String::with_capacity(capacity + suffix_len), + buf, name_len: 0, device_id, } @@ -225,7 +346,7 @@ impl ProtocolAddress { /// Write the name via closure, then append the device_id suffix. /// Single write pass — no intermediate copy. - pub fn reset_with(&mut self, write_name: impl FnOnce(&mut String)) { + pub fn reset_with(&mut self, write_name: impl FnOnce(&mut AddressBuf)) { self.buf.clear(); write_name(&mut self.buf); self.name_len = self.buf.len(); @@ -234,7 +355,7 @@ impl ProtocolAddress { #[inline] pub fn name(&self) -> &str { - &self.buf[..self.name_len] + &self.buf.as_str()[..self.name_len] } #[inline] @@ -244,13 +365,13 @@ impl ProtocolAddress { #[inline] pub fn as_str(&self) -> &str { - &self.buf + self.buf.as_str() } } impl PartialEq for ProtocolAddress { fn eq(&self, other: &Self) -> bool { - self.buf == other.buf + self.as_str() == other.as_str() } } @@ -258,7 +379,7 @@ impl Eq for ProtocolAddress {} impl Hash for ProtocolAddress { fn hash(&self, state: &mut H) { - self.buf.hash(state); + self.as_str().hash(state); } } @@ -270,12 +391,267 @@ impl PartialOrd for ProtocolAddress { impl Ord for ProtocolAddress { fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.buf.cmp(&other.buf) + self.as_str().cmp(other.as_str()) } } impl fmt::Display for ProtocolAddress { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - f.write_str(&self.buf) + f.write_str(self.as_str()) + } +} + +#[cfg(test)] +mod address_buffer_tests { + + /// Every comparison reads the rendered string, so an address with no name + /// yet still has to carry its device. Two unnamed addresses that differ + /// only by device would otherwise be the same map key, and the session + /// cache would serve one device's entry to another. + #[test] + fn unnamed_addresses_still_differ_by_device() { + let first = ProtocolAddress::empty(DeviceId::new(1)); + let second = ProtocolAddress::empty(DeviceId::new(2)); + + assert_ne!(first, second, "the device must keep them apart"); + assert_ne!( + hash_of(&first), + hash_of(&second), + "equal hashes would collide them in the session cache" + ); + assert_eq!(first.name(), "", "neither has a name yet"); + assert_eq!(first.device_id(), DeviceId::new(1)); + } + + /// Reusing an address for a shorter name leaves the previous one's bytes in + /// the inline tail, because the reset only rewinds the length. Anything + /// that prints the backing array rather than the live slice therefore leaks + /// the peer we were talking to before, into any log line that formats an + /// address or an error carrying one. + #[test] + fn debug_never_shows_the_address_that_was_there_before() { + let mut address = ProtocolAddress::new("5511999998888@s.whatsapp.net", 1.into()); + assert!( + address.buf.is_inline(), + "precondition: this name fits inline" + ); + let previous_tail = "8888@s.whatsapp.net"; + assert!( + format!("{:?}", address.buf).contains(previous_tail), + "precondition: the long name is what the buffer holds" + ); + + address.reset_with(|buf| buf.push_str("55119@lid")); + + let shown = format!("{:?}", address.buf); + assert!( + !shown.contains(previous_tail), + "the previous address must not survive into the output: {shown}" + ); + assert!( + !shown.contains('['), + "printing the backing array is what leaks the tail: {shown}" + ); + } + + /// The same guarantee one level up, since this is the type that actually + /// reaches logs and error messages. + #[test] + fn a_protocol_address_prints_only_its_own_name() { + let mut addr = ProtocolAddress::new("5511999998888@s.whatsapp.net", 1.into()); + addr.reset_with(|buf| buf.push_str("55119@lid")); + + let shown = format!("{addr:?}"); + assert!( + !shown.contains("8888"), + "a reused address must not print the previous peer: {shown}" + ); + assert!( + shown.contains("55119@lid"), + "it must still print what it holds: {shown}" + ); + } + use super::*; + use std::collections::hash_map::DefaultHasher; + + impl AddressBuf { + /// Test-only view of the representation. Production code must never + /// branch on this: the whole contract is that the two arms are + /// indistinguishable as values. + fn is_inline(&self) -> bool { + matches!(self.0, AddressRepr::Inline { .. }) + } + } + + fn hash_of(address: &ProtocolAddress) -> u64 { + let mut hasher = DefaultHasher::new(); + address.hash(&mut hasher); + hasher.finish() + } + + /// A name that fits, plus its device suffix, never reaches the heap, and + /// still splits into name and full string at the right place. + #[test] + fn a_name_that_fits_stays_inline_and_reads_back_whole() { + let address = ProtocolAddress::new("5511987650001:5@c.us", DeviceId::new(0)); + assert!(address.buf.is_inline()); + assert_eq!(address.name(), "5511987650001:5@c.us"); + assert_eq!(address.as_str(), "5511987650001:5@c.us.0"); + assert_eq!(address.device_id(), DeviceId::new(0)); + } + + /// The exact boundary in both directions: the last name that fits with its + /// suffix, and the first one that does not. + #[test] + fn the_inline_boundary_holds_on_both_sides() { + let suffix_len = ".0".len(); + let longest_fitting = "a".repeat(INLINE_CAPACITY - suffix_len); + let fits = ProtocolAddress::new(&longest_fitting, DeviceId::new(0)); + assert!( + fits.buf.is_inline(), + "the longest fitting name must stay inline" + ); + assert_eq!(fits.as_str().len(), INLINE_CAPACITY); + assert_eq!(fits.name(), longest_fitting); + + let one_too_long = "a".repeat(INLINE_CAPACITY - suffix_len + 1); + let spills = ProtocolAddress::new(&one_too_long, DeviceId::new(0)); + assert!(!spills.buf.is_inline(), "one byte over must spill"); + assert_eq!(spills.name(), one_too_long); + assert_eq!(spills.as_str(), format!("{one_too_long}.0")); + } + + /// A multi-digit device id takes a different path than `".0"`: the suffix + /// is written through `write!` rather than one `push_str`, so the spill can + /// land *between* the dot and the digits. Getting that wrong would leave a + /// half-written suffix or a `name_len` measured after it, which is why the + /// case is worth pinning separately from the fast path above. + #[test] + fn a_multi_digit_suffix_spills_without_splitting_itself() { + let suffix = ".123"; + // Sized so the name alone fits and the suffix is what pushes it over. + let name = "b".repeat(INLINE_CAPACITY - suffix.len() + 1); + let address = ProtocolAddress::new(&name, DeviceId::new(123)); + + assert!( + !address.buf.is_inline(), + "the suffix is what must have pushed this past the inline capacity" + ); + assert_eq!( + address.as_str(), + format!("{name}{suffix}"), + "a fragment written across the spill would corrupt the address" + ); + assert_eq!( + address.name(), + name, + "name_len must still point at the end of the name, not into the suffix" + ); + assert_eq!(address.device_id(), DeviceId::new(123)); + } + + /// The representation is not part of the value. `ProtocolAddress` is a + /// `HashMap` key, so an inline value and a heap value holding the same + /// characters must compare, order and hash the same. + #[test] + fn the_same_characters_compare_and_hash_alike_from_either_representation() { + let name = "5511987650001:5@c.us"; + let inline = ProtocolAddress::new(name, DeviceId::new(0)); + + // Force the heap arm, then rewrite it to the short name: a cleared + // heap buffer keeps its allocation, so this really is the same + // characters in the other representation. + let mut spilled = ProtocolAddress::new(&"z".repeat(INLINE_CAPACITY * 2), DeviceId::new(0)); + assert!(!spilled.buf.is_inline()); + spilled.reset_with(|buf| buf.push_str(name)); + assert!( + !spilled.buf.is_inline(), + "the fixture must still be on the heap, or it proves nothing" + ); + assert!(inline.buf.is_inline()); + + assert_eq!(inline, spilled); + assert_eq!(inline.cmp(&spilled), std::cmp::Ordering::Equal); + assert_eq!(hash_of(&inline), hash_of(&spilled)); + assert_eq!(inline.name(), spilled.name()); + assert_eq!(inline.as_str(), spilled.as_str()); + assert_eq!(inline.to_string(), spilled.to_string()); + + let mut map = std::collections::HashMap::new(); + map.insert(inline, "value"); + assert_eq!( + map.get(&spilled).copied(), + Some("value"), + "a heap-built key must find the inline-built entry" + ); + } + + /// Different characters must still differ, whichever arm holds them -- + /// otherwise the equality above would be trivially true. + #[test] + fn different_characters_stay_different_across_representations() { + let inline = ProtocolAddress::new("alice@c.us", DeviceId::new(0)); + let mut spilled = ProtocolAddress::new(&"z".repeat(INLINE_CAPACITY * 2), DeviceId::new(0)); + spilled.reset_with(|buf| buf.push_str("bob@c.us")); + assert!(!spilled.buf.is_inline()); + + assert_ne!(inline, spilled); + assert_eq!(inline.cmp(&spilled), std::cmp::Ordering::Less); + } + + /// Multi-byte characters survive both the inline copy and the spill. The + /// spill copies whole appended fragments, so it can never land inside a + /// character. + #[test] + fn multibyte_names_survive_inline_and_spilled() { + let short = "héllo✅@c.us"; + let inline = ProtocolAddress::new(short, DeviceId::new(0)); + assert!(inline.buf.is_inline()); + assert_eq!(inline.name(), short); + + // Fill to one byte short of the limit, then append a two-byte char: + // the append cannot fit, so the whole fragment moves to the heap. + let filler = "a".repeat(INLINE_CAPACITY - 1); + let mut spilled = ProtocolAddress::empty(DeviceId::new(0)); + spilled.reset_with(|buf| { + buf.push_str(&filler); + buf.push('é'); + }); + assert!(!spilled.buf.is_inline()); + assert_eq!(spilled.name(), format!("{filler}é")); + assert_eq!(spilled.as_str(), format!("{filler}é.0")); + } + + /// An empty name is a real state (a freshly reset address), and the + /// suffix still lands. + #[test] + fn an_empty_name_still_carries_its_device_suffix() { + let address = ProtocolAddress::new("", DeviceId::new(0)); + assert!(address.buf.is_inline()); + assert_eq!(address.name(), ""); + assert_eq!(address.as_str(), ".0"); + } + + /// A non-zero device id is rendered, and the name boundary still excludes + /// the suffix however many digits it takes. + #[test] + fn a_multi_digit_device_id_is_appended_outside_the_name() { + let address = ProtocolAddress::new("alice@c.us", DeviceId::new(123)); + assert_eq!(address.name(), "alice@c.us"); + assert_eq!(address.as_str(), "alice@c.us.123"); + assert_eq!(address.device_id(), DeviceId::new(123)); + } + + /// Rewriting an address must replace its content, not append to it. + #[test] + fn resetting_replaces_the_previous_name() { + let mut address = ProtocolAddress::new("alice@c.us", DeviceId::new(0)); + address.reset_with(|buf| buf.push_str("bob@c.us")); + assert_eq!(address.name(), "bob@c.us"); + assert_eq!(address.as_str(), "bob@c.us.0"); + + let mut spilled = ProtocolAddress::new(&"z".repeat(INLINE_CAPACITY * 2), DeviceId::new(0)); + spilled.reset_with(|buf| buf.push_str("bob@c.us")); + assert_eq!(spilled.as_str(), "bob@c.us.0"); } } diff --git a/wacore/libsignal/src/core/mod.rs b/wacore/libsignal/src/core/mod.rs index d1479e4a9..48a531fe8 100644 --- a/wacore/libsignal/src/core/mod.rs +++ b/wacore/libsignal/src/core/mod.rs @@ -8,6 +8,6 @@ mod address; pub mod curve; pub use address::{ - Aci, DeviceId, Pni, ProtocolAddress, ServiceId, ServiceIdFixedWidthBinaryBytes, ServiceIdKind, - WrongKindOfServiceIdError, + Aci, AddressBuf, DeviceId, Pni, ProtocolAddress, ServiceId, ServiceIdFixedWidthBinaryBytes, + ServiceIdKind, WrongKindOfServiceIdError, }; diff --git a/wacore/libsignal/src/protocol/mod.rs b/wacore/libsignal/src/protocol/mod.rs index 6ff39826a..6c44a5b35 100644 --- a/wacore/libsignal/src/protocol/mod.rs +++ b/wacore/libsignal/src/protocol/mod.rs @@ -38,7 +38,8 @@ mod stores; mod timestamp; pub use crate::core::curve::{CurveError, KeyPair, PreparedVerifyingKey, PrivateKey, PublicKey}; pub use crate::core::{ - Aci, DeviceId, Pni, ProtocolAddress, ServiceId, ServiceIdFixedWidthBinaryBytes, ServiceIdKind, + Aci, AddressBuf, DeviceId, Pni, ProtocolAddress, ServiceId, ServiceIdFixedWidthBinaryBytes, + ServiceIdKind, }; pub use crate::protocol::protocol::SENDERKEY_MESSAGE_CURRENT_VERSION; pub use crate::protocol::sender_keys::InvalidSenderKeySessionError; diff --git a/wacore/libsignal/src/protocol/session_cipher.rs b/wacore/libsignal/src/protocol/session_cipher.rs index dc6df1e27..d06f281cc 100644 --- a/wacore/libsignal/src/protocol/session_cipher.rs +++ b/wacore/libsignal/src/protocol/session_cipher.rs @@ -1887,12 +1887,12 @@ mod tests { fn setup_established_session() -> TestPair { let mut rng = rand::make_rng::(); - let alice_addr = ProtocolAddress::new("alice".to_string(), 1.into()); + let alice_addr = ProtocolAddress::new("alice", 1.into()); let alice_id = IdentityKeyPair::generate(&mut rng); let mut alice_sessions = MemSessionStore::new(); let mut alice_identity = MemIdentityStore::new(alice_id, 1); - let bob_addr = ProtocolAddress::new("bob".to_string(), 1.into()); + let bob_addr = ProtocolAddress::new("bob", 1.into()); let bob_id = IdentityKeyPair::generate(&mut rng); let bob_identity_key = *bob_id.identity_key(); @@ -2439,7 +2439,7 @@ mod tests { MemSignedPreKeyStore, PreKeyBundle, ) { - let bob_addr = ProtocolAddress::new("bob".to_string(), 1.into()); + let bob_addr = ProtocolAddress::new("bob", 1.into()); let bob_id = IdentityKeyPair::generate(rng); let bob_identity_key = *bob_id.identity_key(); @@ -2550,7 +2550,7 @@ mod tests { (false, IdentityChange::NewOrUnchanged), (true, IdentityChange::ReplacedExisting), ] { - let alice_addr = ProtocolAddress::new("alice".to_string(), 1.into()); + let alice_addr = ProtocolAddress::new("alice", 1.into()); let alice_id = IdentityKeyPair::generate(&mut rng); let mut alice_sessions = MemSessionStore::new(); let mut alice_identity = MemIdentityStore::new(alice_id, 1); @@ -2613,7 +2613,7 @@ mod tests { #[test] fn process_prekey_signals_reuse_for_established_session() { let mut rng = rand::make_rng::(); - let alice_addr = ProtocolAddress::new("alice".to_string(), 1.into()); + let alice_addr = ProtocolAddress::new("alice", 1.into()); let alice_id = IdentityKeyPair::generate(&mut rng); let mut alice_sessions = MemSessionStore::new(); let mut alice_identity = MemIdentityStore::new(alice_id, 1); @@ -2902,7 +2902,7 @@ mod tests { fn decrypt_with_empty_session_returns_session_not_found() { let mut rng = rand::make_rng::(); - let alice_addr = ProtocolAddress::new("alice".to_string(), 1.into()); + let alice_addr = ProtocolAddress::new("alice", 1.into()); let alice_id = IdentityKeyPair::generate(&mut rng); let bob_id = IdentityKeyPair::generate(&mut rng); let alice_identity_key = *alice_id.identity_key(); diff --git a/wacore/libsignal/src/protocol/storage/traits_hook_tests.rs b/wacore/libsignal/src/protocol/storage/traits_hook_tests.rs index 8bd9a8563..73b4a2587 100644 --- a/wacore/libsignal/src/protocol/storage/traits_hook_tests.rs +++ b/wacore/libsignal/src/protocol/storage/traits_hook_tests.rs @@ -12,7 +12,7 @@ use crate::protocol::{IdentityKey, IdentityKeyPair, ProtocolAddress, SignalProto use futures::executor::block_on; fn test_address() -> ProtocolAddress { - ProtocolAddress::new("5511987650001@s.whatsapp.net".to_string(), 1.into()) + ProtocolAddress::new("5511987650001@s.whatsapp.net", 1.into()) } fn test_identity() -> IdentityKey { diff --git a/wacore/libsignal/tests/counter_lease.rs b/wacore/libsignal/tests/counter_lease.rs index 1769469ec..5fc90b78e 100644 --- a/wacore/libsignal/tests/counter_lease.rs +++ b/wacore/libsignal/tests/counter_lease.rs @@ -208,7 +208,7 @@ impl Peer { .expect("valid bundle"); let peer = Self { - address: ProtocolAddress::new(name.to_string(), 1u32.into()), + address: ProtocolAddress::new(name, 1u32.into()), identity_store: InMemoryIdentityKeyStore { identity_key_pair, registration_id, diff --git a/wacore/libsignal/tests/session_divergence.rs b/wacore/libsignal/tests/session_divergence.rs index ef89dc575..b0882c765 100644 --- a/wacore/libsignal/tests/session_divergence.rs +++ b/wacore/libsignal/tests/session_divergence.rs @@ -213,7 +213,7 @@ impl Peer { }); Self { - address: ProtocolAddress::new(name.to_string(), device_id.into()), + address: ProtocolAddress::new(name, device_id.into()), identity_store, prekey_store, signed_prekey_store, @@ -502,8 +502,8 @@ fn alice_delete_then_rebuild_loses_old_chain() { #[test] fn pkmsg_reset_does_not_fix_peer_outbound_if_delivered_to_wrong_store_key() { - let bob_lid = ProtocolAddress::new("100000000000001@lid".to_string(), 0.into()); - let bob_pn = ProtocolAddress::new("15550001000@c.us".to_string(), 0.into()); + let bob_lid = ProtocolAddress::new("100000000000001@lid", 0.into()); + let bob_pn = ProtocolAddress::new("15550001000@c.us", 0.into()); let mut alice = Peer::new("alice", 1); let mut bob = Peer::new("100000000000001@lid", 0); diff --git a/wacore/src/reporting_token.rs b/wacore/src/reporting_token.rs index 03ce08162..b08710d44 100644 --- a/wacore/src/reporting_token.rs +++ b/wacore/src/reporting_token.rs @@ -17,6 +17,7 @@ //! without seeing the full message content. Only specific fields are extracted //! based on a predefined whitelist matching WhatsApp Web behavior. +use smallvec::SmallVec; use std::{fmt, sync::LazyLock}; use anyhow::{Result, anyhow}; @@ -499,6 +500,25 @@ fn encode_varint(value: u64) -> Vec { buf[..len].to_vec() } +/// One extracted field's bytes on the way to the token. +/// +/// A field the token copies verbatim is named by its range in the input, so it +/// is copied once, into the result. Only a nested field, whose content is +/// re-framed under a fresh tag and length, has bytes of its own. +enum Piece { + Borrowed(core::ops::Range), + Owned(Vec), +} + +impl Piece { + fn len(&self) -> usize { + match self { + Piece::Borrowed(range) => range.len(), + Piece::Owned(bytes) => bytes.len(), + } + } +} + /// Extract reporting token content from encoded protobuf message bytes. /// /// This function parses raw protobuf bytes and extracts only the fields @@ -514,8 +534,12 @@ pub fn extract_reporting_token_content( data: &[u8], whitelist: &[ReportingField], ) -> Option> { - // Pre-size: most messages have 1-3 fields - let mut extracted: Vec<(u32, Vec)> = Vec::with_capacity(4); + // 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 + // in the input, and the staging list is inline for the field counts a real + // message has. + let mut extracted: SmallVec<[(u32, Piece); 4]> = SmallVec::new(); let mut pos = 0; while pos < data.len() { @@ -539,7 +563,7 @@ pub fn extract_reporting_token_content( let (_, val_len) = decode_varint(&data[pos..])?; pos += val_len; if entry.is_some() { - extracted.push((field_number, data[field_start..pos].to_vec())); + extracted.push((field_number, Piece::Borrowed(field_start..pos))); } } wire_type::FIXED64 => { @@ -548,7 +572,7 @@ pub fn extract_reporting_token_content( } pos += 8; if entry.is_some() { - extracted.push((field_number, data[field_start..pos].to_vec())); + extracted.push((field_number, Piece::Borrowed(field_start..pos))); } } wire_type::FIXED32 => { @@ -557,7 +581,7 @@ pub fn extract_reporting_token_content( } pos += 4; if entry.is_some() { - extracted.push((field_number, data[field_start..pos].to_vec())); + extracted.push((field_number, Piece::Borrowed(field_start..pos))); } } wire_type::LENGTH_DELIMITED => { @@ -588,7 +612,7 @@ pub fn extract_reporting_token_content( field_bytes.extend_from_slice(&tag_buf[..tag_len]); field_bytes.extend_from_slice(&len_buf[..len_len]); field_bytes.extend(nested); - extracted.push((field_number, field_bytes)); + extracted.push((field_number, Piece::Owned(field_bytes))); } } else if let Some(subfields) = entry.subfields { if let Some(nested) = extract_reporting_token_content( @@ -607,10 +631,10 @@ pub fn extract_reporting_token_content( field_bytes.extend_from_slice(&tag_buf[..tag_len]); field_bytes.extend_from_slice(&len_buf[..len_len]); field_bytes.extend(nested); - extracted.push((field_number, field_bytes)); + extracted.push((field_number, Piece::Owned(field_bytes))); } } else { - extracted.push((field_number, data[field_start..pos].to_vec())); + extracted.push((field_number, Piece::Borrowed(field_start..pos))); } } } @@ -627,11 +651,23 @@ pub fn extract_reporting_token_content( extracted.sort_by_key(|(num, _)| *num); - let total_len: usize = extracted.iter().map(|(_, v)| v.len()).sum(); + let total_len: usize = extracted.iter().map(|(_, piece)| piece.len()).sum(); let mut result = Vec::with_capacity(total_len); - for (_, bytes) in extracted { - result.extend(bytes); + for (_, piece) in extracted { + match piece { + Piece::Borrowed(range) => result.extend_from_slice(&data[range]), + Piece::Owned(bytes) => result.extend_from_slice(&bytes), + } } + // The reservation is only worth making if it matches what was written; an + // understated `Piece::len` would reallocate here and go unnoticed + // otherwise. Checked rather than asserted on capacity, which the allocator + // is free to round up. + debug_assert_eq!( + result.len(), + total_len, + "the reservation disagreed with the bytes written" + ); Some(result) } @@ -1629,6 +1665,197 @@ mod tests { ); } + /// Nothing whitelisted matched: the token has no content at all, which is + /// a distinct answer from "content that happens to be empty". + #[test] + fn no_matching_field_yields_no_content() { + let data = vec![ + 0x08, 0x96, 0x01, // field 1: varint 150 + 0x12, 0x05, b'h', b'e', b'l', b'l', b'o', // field 2: "hello" + ]; + let whitelist = &[ReportingField::new(9)]; + assert!(extract_reporting_token_content(&data, whitelist).is_none()); + assert!(extract_reporting_token_content(&[], whitelist).is_none()); + } + + /// The concatenation order is by ascending field number, not wire order: + /// a message that puts a higher-numbered field first must still hash the + /// same as one that does not. + #[test] + fn fields_are_concatenated_in_field_number_order_not_wire_order() { + let low = [0x0a, 0x01, b'a']; // field 1: "a" + let high = [0x12, 0x01, b'b']; // field 2: "b" + let whitelist = &[ReportingField::new(1), ReportingField::new(2)]; + + let mut wire_ascending = Vec::new(); + wire_ascending.extend_from_slice(&low); + wire_ascending.extend_from_slice(&high); + let mut wire_descending = Vec::new(); + wire_descending.extend_from_slice(&high); + wire_descending.extend_from_slice(&low); + + let expected = [0x0a, 0x01, b'a', 0x12, 0x01, b'b']; + assert_eq!( + extract_reporting_token_content(&wire_ascending, whitelist).unwrap(), + expected + ); + assert_eq!( + extract_reporting_token_content(&wire_descending, whitelist).unwrap(), + expected + ); + } + + /// Repeats of the same field keep their wire order relative to each other, + /// which is what makes the sort's stability part of the wire contract. + #[test] + fn repeats_of_one_field_keep_their_wire_order() { + let data = vec![ + 0x0a, 0x01, b'x', // field 1: "x" + 0x12, 0x01, b'm', // field 2: "m" + 0x0a, 0x01, b'y', // field 1: "y" + ]; + let whitelist = &[ReportingField::new(1), ReportingField::new(2)]; + + assert_eq!( + extract_reporting_token_content(&data, whitelist).unwrap(), + vec![0x0a, 0x01, b'x', 0x0a, 0x01, b'y', 0x12, 0x01, b'm'], + ); + } + + /// More fields than the inline staging list holds: spilling must not + /// reorder or drop anything. + #[test] + fn more_fields_than_the_inline_list_holds_still_concatenate_in_order() { + let mut data = Vec::new(); + let mut whitelist = Vec::new(); + // Emit fields 8..=1 (descending on the wire) so the sort has work to do. + for field in (1u8..=8).rev() { + data.push((field << 3) | 2); + data.push(1); + data.push(b'a' + field); + whitelist.push(ReportingField::new(u32::from(field))); + } + + let extracted = extract_reporting_token_content(&data, &whitelist) + .expect("eight whitelisted fields extract"); + + let mut expected = Vec::new(); + for field in 1u8..=8 { + expected.push((field << 3) | 2); + expected.push(1); + expected.push(b'a' + field); + } + assert_eq!(extracted, expected); + } + + /// A multi-byte payload must be copied byte for byte: a range-based piece + /// that got its bounds wrong would slice a UTF-8 sequence in half. + #[test] + fn a_multibyte_payload_survives_the_copy_byte_for_byte() { + let text = "héllo ✅ 日本"; + let bytes = text.as_bytes(); + let mut data = vec![0x0a, u8::try_from(bytes.len()).unwrap()]; + data.extend_from_slice(bytes); + // A second, unwhitelisted field so the extraction is not the whole input. + data.extend_from_slice(&[0x12, 0x02, 0xff, 0xfe]); + + let whitelist = &[ReportingField::new(1)]; + let extracted = extract_reporting_token_content(&data, whitelist) + .expect("the whitelisted field extracts"); + + assert_eq!(extracted, data[..2 + bytes.len()]); + assert_eq!(std::str::from_utf8(&extracted[2..]).unwrap(), text); + } + + /// A nested field is re-framed under a fresh length, so it is the one kind + /// of piece that cannot be a range into the input. + #[test] + fn a_nested_field_is_rebuilt_and_still_ordered_with_the_flat_ones() { + let inner = vec![ + 0x0a, 0x01, b'a', // field 1: "a" (kept) + 0x12, 0x01, b'b', // field 2: "b" (dropped) + ]; + // field 1 (flat, kept) then field 6 (nested), emitted in that order. + let mut data = vec![0x0a, 0x01, b'z']; + data.push(0x32); + data.push(u8::try_from(inner.len()).unwrap()); + data.extend_from_slice(&inner); + + static TEST_SUBFIELDS: &[ReportingField] = &[ReportingField::new(1)]; + let whitelist = &[ + ReportingField::new(1), + ReportingField::with_subfields(6, TEST_SUBFIELDS), + ]; + + let extracted = extract_reporting_token_content(&data, whitelist) + .expect("flat + nested fields extract"); + + assert_eq!( + extracted, + vec![ + 0x0a, 0x01, b'z', // field 1 verbatim + 0x32, 0x03, 0x0a, 0x01, b'a', // field 6 re-framed around field 1 only + ], + ); + } + + /// A nested field whose every subfield is filtered out contributes + /// nothing, rather than an empty re-framed wrapper. + #[test] + fn a_nested_field_with_nothing_kept_contributes_nothing() { + let inner = vec![0x12, 0x01, b'b']; // field 2 only + let mut data = vec![0x32, u8::try_from(inner.len()).unwrap()]; + data.extend_from_slice(&inner); + + static TEST_SUBFIELDS: &[ReportingField] = &[ReportingField::new(1)]; + let whitelist = &[ReportingField::with_subfields(6, TEST_SUBFIELDS)]; + + assert!(extract_reporting_token_content(&data, whitelist).is_none()); + } + + /// The result is allocated once, from a length summed over every piece. + /// A piece kind left out of that sum makes the buffer grow mid-write, + /// which is exactly the allocation this staging model exists to avoid. + #[test] + fn the_result_is_allocated_once_for_the_exact_length_it_holds() { + // Sizes chosen so an understated reservation cannot land back on the + // final length by accident: Vec grows to at least twice its capacity. + let inner = vec![0x0a, 0x01, b'a']; + let mut nested_and_flat = vec![0x0a, 0x0a]; + nested_and_flat.extend_from_slice(b"0123456789"); + nested_and_flat.push(0x32); + nested_and_flat.push(u8::try_from(inner.len()).unwrap()); + nested_and_flat.extend_from_slice(&inner); + + static TEST_SUBFIELDS: &[ReportingField] = &[ReportingField::new(1)]; + let cases: [(&[u8], &[ReportingField]); 2] = [ + (&[0x0a, 0x01, b'z'], &[ReportingField::new(1)]), + ( + &nested_and_flat, + &[ + ReportingField::new(1), + ReportingField::with_subfields(6, TEST_SUBFIELDS), + ], + ), + ]; + + for (data, whitelist) in cases { + let extracted = + extract_reporting_token_content(data, whitelist).expect("content extracts"); + // The property is that the length reserved from `Piece::len` is the + // length actually written; an understated reservation reallocates + // silently. That is checked by the `debug_assert` inside the + // extractor, which every call here exercises, so this case exists + // to feed it shapes: flat fields, a nested field, and both mixed. + // Asserting `capacity()` instead would be testing the allocator, + // which is free to round a request up. + assert!( + !extracted.is_empty(), + "each case must actually extract something, or it feeds the check nothing" + ); + } + } + #[test] fn test_excluded_message_types() { // Verify all excluded message types return None/false diff --git a/wacore/src/send/dm.rs b/wacore/src/send/dm.rs index 365cce242..b8d7d9748 100644 --- a/wacore/src/send/dm.rs +++ b/wacore/src/send/dm.rs @@ -201,8 +201,10 @@ pub async fn prepare_dm_stanza( // a bare-enc mode would require refactoring the encryption layer. // The form is accepted by the server regardless. + // Both fan-outs append into the vector already sized for the whole + // participant set, so neither stages a node list of its own. if !recipient_devices.is_empty() { - let result = encrypt_for_devices( + let summary = encrypt_for_devices_into( runtime, stores, resolver, @@ -210,14 +212,14 @@ pub async fn prepare_dm_stanza( &recipient_plaintext, hide_decrypt_fail, mediatype, + &mut participant_nodes, ) .await?; - participant_nodes.extend(result.participant_nodes); - includes_prekey_message = includes_prekey_message || result.includes_prekey_message; + includes_prekey_message = includes_prekey_message || summary.includes_prekey_message; } if !own_other_devices.is_empty() { - let result = encrypt_for_devices( + let summary = encrypt_for_devices_into( runtime, stores, resolver, @@ -225,10 +227,10 @@ pub async fn prepare_dm_stanza( &own_devices_plaintext, hide_decrypt_fail, mediatype, + &mut participant_nodes, ) .await?; - participant_nodes.extend(result.participant_nodes); - includes_prekey_message = includes_prekey_message || result.includes_prekey_message; + includes_prekey_message = includes_prekey_message || summary.includes_prekey_message; } // All per-device encrypts failed: an empty would silently diff --git a/wacore/src/send/encrypt.rs b/wacore/src/send/encrypt.rs index 21f8f87d7..b1fe8be92 100644 --- a/wacore/src/send/encrypt.rs +++ b/wacore/src/send/encrypt.rs @@ -371,6 +371,61 @@ pub async fn encrypt_for_devices( .await } +/// What a fan-out reports back when its `` nodes went straight into +/// the caller's stanza buffer instead of a per-fan-out [`EncryptResult`]. +pub struct EncryptFanoutSummary { + pub includes_prekey_message: bool, + /// True if any device returned 406 (unregistered) during prekey fetch. + pub had_unregistered_device: bool, +} + +/// [`encrypt_for_devices`] for a caller that already owns the buffer the nodes +/// belong in. +/// +/// [`EncryptResult`] is shaped for the group path, which needs the encrypted +/// device list to tell a partial SKDM distribution from a complete one. A DM +/// never asks that question and knows up front how many participants it can +/// have, so it sizes one vector and lets each fan-out append into it: the +/// per-fan-out node vector and the device list it would otherwise carry are +/// both work done only to be moved and dropped. +#[allow(clippy::too_many_arguments)] +pub async fn encrypt_for_devices_into( + runtime: &dyn Runtime, + stores: &mut SignalStores<'_>, + resolver: &dyn SendContextResolver, + devices: &[Jid], + plaintext_to_encrypt: &[u8], + hide_decrypt_fail: bool, + mediatype: Option<&str>, + participant_nodes: &mut Vec, +) -> Result { + let plan = ensure_sessions_for_devices(runtime, stores, resolver, devices).await?; + // `first_error` is dropped here exactly as `encrypt_for_devices` drops it: + // a DM reports failure through the empty-participants check, not per device. + let RawEncryptAttempt { result: raw, .. } = encrypt_for_devices_with_sessions_raw_detailed( + runtime, + stores, + devices, + plaintext_to_encrypt, + plan, + ) + .await?; + + participant_nodes.reserve(raw.devices.len()); + for one in raw.devices { + participant_nodes.push(encrypted_device_to_participant_node( + one, + mediatype, + hide_decrypt_fail, + )); + } + + Ok(EncryptFanoutSummary { + includes_prekey_message: raw.includes_prekey_message, + had_unregistered_device: raw.had_unregistered_device, + }) +} + /// Session material prepared for one encrypt fan-out: per-index LID /// encryption overrides (mirroring the `devices` slice it was built from) /// plus whether any device 406'd during prekey fetch. Produced only by diff --git a/wacore/src/send/peer.rs b/wacore/src/send/peer.rs index 720ab856c..ff7d5a2ab 100644 --- a/wacore/src/send/peer.rs +++ b/wacore/src/send/peer.rs @@ -9,7 +9,7 @@ pub async fn prepare_peer_stanza( transport_jid: Jid, signal_address: &ProtocolAddress, message: &wa::Message, - request_id: String, + request_id: &str, account: Option<&wa::ADVSignedDeviceIdentity>, ) -> Result where @@ -38,7 +38,7 @@ pub async fn prepare_peer_stanza_with_options( transport_jid: Jid, signal_address: &ProtocolAddress, message: &wa::Message, - request_id: String, + request_id: &str, account: Option<&wa::ADVSignedDeviceIdentity>, options: PeerMessageOptions, ) -> Result diff --git a/wacore/src/send/tests.rs b/wacore/src/send/tests.rs index c5cc5f6bd..d513b8eb3 100644 --- a/wacore/src/send/tests.rs +++ b/wacore/src/send/tests.rs @@ -2086,7 +2086,7 @@ mod group_retry { jid.clone(), &addr, &wa::Message::default(), - "peer-test-1".into(), + "peer-test-1", account, options, ) @@ -2191,7 +2191,7 @@ mod group_retry { jid.clone(), &addr, &wa::Message::default(), - "peer-test-no-account".into(), + "peer-test-no-account", None, ) .await; @@ -2242,7 +2242,7 @@ mod group_retry { jid.clone(), &addr, &wa::Message::default(), - "peer-preflight-1".into(), + "peer-preflight-1", None, ) .await; @@ -2409,7 +2409,7 @@ mod group_retry { jid.clone(), &addr, &wa::Message::default(), - "preflight-take-bail".into(), + "preflight-take-bail", None, ) .await; @@ -2428,7 +2428,7 @@ mod group_retry { jid.clone(), &addr, &wa::Message::default(), - "preflight-take-pass".into(), + "preflight-take-pass", Some(&account), ) .await; @@ -4351,4 +4351,619 @@ mod local_identity_change_on_send { "replaced identity on the send path must be reported via the resolver" ); } + + /// The DM fan-out writes its `` nodes into the stanza's own + /// participant vector instead of staging one per half. + mod dm_fanout_sink { + use super::*; + + fn sentinel() -> Node { + NodeBuilder::new("sentinel").build() + } + + fn participant_jids(nodes: &[Node]) -> Vec { + nodes + .iter() + .map(|n| { + n.attrs() + .optional_string("jid") + .expect("participant node carries a jid") + .into_owned() + }) + .collect() + } + + async fn fan_out_into( + devices: &[Jid], + resolver: &MockSendContextResolver, + nodes: &mut Vec, + ) -> EncryptFanoutSummary { + let (mut session_store, mut identity_store) = stores_with_sessions(devices).await; + let mut prekey_store = UnusedPreKeyStore; + let signed_prekey_store = UnusedSignedPreKeyStore; + let mut sender_key_store = MemSenderKeyStore::default(); + let mut stores = raw_fanout_stores( + &mut sender_key_store, + &mut session_store, + &mut identity_store, + &mut prekey_store, + &signed_prekey_store, + ); + encrypt_for_devices_into( + &TokioTestRuntime, + &mut stores, + resolver, + devices, + b"payload", + false, + None, + nodes, + ) + .await + .expect("fan-out into the caller's buffer") + } + + /// A half with no devices must contribute nothing at all: it may not + /// clear, replace, or grow the buffer it was handed. + #[tokio::test] + async fn an_empty_half_leaves_the_buffer_exactly_as_it_found_it() { + let mut nodes = vec![sentinel()]; + let resolver = MockSendContextResolver::new(); + + let summary = fan_out_into(&[], &resolver, &mut nodes).await; + + assert_eq!(nodes.len(), 1, "an empty half must append nothing"); + assert_eq!(nodes[0].tag.as_ref(), "sentinel", "and remove nothing"); + assert!(!summary.includes_prekey_message); + assert!(!summary.had_unregistered_device); + } + + /// The single-device DM, which is the whole fan-out on a steady 1:1 + /// chat: one node, appended after whatever the caller already had. + #[tokio::test] + async fn one_device_appends_one_node_after_the_existing_content() { + let device: Jid = "5511900000001:0@s.whatsapp.net".parse().unwrap(); + let mut nodes = vec![sentinel()]; + let resolver = MockSendContextResolver::new(); + + let summary = fan_out_into(std::slice::from_ref(&device), &resolver, &mut nodes).await; + + assert_eq!(nodes.len(), 2); + assert_eq!( + nodes[0].tag.as_ref(), + "sentinel", + "the sink appends; it does not overwrite" + ); + assert_eq!(participant_jids(&nodes[1..]), vec![device.to_string()]); + assert!( + summary.includes_prekey_message, + "a session whose pre-key is still unacked emits pkmsg" + ); + } + + /// Several devices, appended in fan-out order after the existing + /// content, so two halves in a row concatenate rather than interleave. + #[tokio::test] + async fn many_devices_append_in_order_after_the_existing_content() { + let first: Vec = (0..3u16) + .map(|i| format!("5511900000002:{i}@s.whatsapp.net").parse().unwrap()) + .collect(); + let second: Vec = vec!["5511900000003:1@s.whatsapp.net".parse().unwrap()]; + let mut nodes = vec![sentinel()]; + let resolver = MockSendContextResolver::new(); + + fan_out_into(&first, &resolver, &mut nodes).await; + fan_out_into(&second, &resolver, &mut nodes).await; + + assert_eq!(nodes[0].tag.as_ref(), "sentinel"); + let mut expected: Vec = first.iter().map(Jid::to_string).collect(); + expected.extend(second.iter().map(Jid::to_string)); + assert_eq!( + participant_jids(&nodes[1..]), + expected, + "each half appends its own devices, in order, after the last" + ); + } + + /// Skip-on-fail: a device with neither a session nor a bundle drops out + /// of the fan-out, and the surviving devices still land in the buffer. + #[tokio::test] + async fn a_device_that_cannot_encrypt_contributes_no_node() { + let good: Jid = "5511900000004:0@s.whatsapp.net".parse().unwrap(); + let sessionless: Jid = "5511900000005:0@s.whatsapp.net".parse().unwrap(); + + // Only `good` gets a session; the resolver offers no bundle for the + // other, so its encrypt has nothing to work with. + let (mut session_store, mut identity_store) = + stores_with_sessions(std::slice::from_ref(&good)).await; + let mut prekey_store = UnusedPreKeyStore; + let signed_prekey_store = UnusedSignedPreKeyStore; + let mut sender_key_store = MemSenderKeyStore::default(); + let mut stores = raw_fanout_stores( + &mut sender_key_store, + &mut session_store, + &mut identity_store, + &mut prekey_store, + &signed_prekey_store, + ); + let resolver = MockSendContextResolver::new().with_missing_bundle(sessionless.clone()); + + let mut nodes = vec![sentinel()]; + encrypt_for_devices_into( + &TokioTestRuntime, + &mut stores, + &resolver, + &[good.clone(), sessionless], + b"payload", + false, + None, + &mut nodes, + ) + .await + .expect("one bad device must not abort the fan-out"); + + assert_eq!( + participant_jids(&nodes[1..]), + vec![good.to_string()], + "only the device that could encrypt is in the participant list" + ); + } + + /// End to end through `prepare_dm_stanza`: recipient devices and own + /// companion devices are two separate fan-outs but one participant + /// list, recipients first. + #[tokio::test] + async fn a_dm_stanza_carries_both_halves_in_one_participants_node() { + let own_jid: Jid = "5511900000010:0@s.whatsapp.net".parse().unwrap(); + let recipient_a: Jid = "5511900000011:0@s.whatsapp.net".parse().unwrap(); + let recipient_b: Jid = "5511900000011:1@s.whatsapp.net".parse().unwrap(); + let own_companion: Jid = "5511900000010:2@s.whatsapp.net".parse().unwrap(); + let all = vec![ + recipient_a.clone(), + recipient_b.clone(), + own_companion.clone(), + own_jid.clone(), + ]; + + let (mut session_store, mut identity_store) = stores_with_sessions(&[ + recipient_a.clone(), + recipient_b.clone(), + own_companion.clone(), + ]) + .await; + let mut prekey_store = UnusedPreKeyStore; + let signed_prekey_store = UnusedSignedPreKeyStore; + let mut sender_key_store = MemSenderKeyStore::default(); + let mut stores = raw_fanout_stores( + &mut sender_key_store, + &mut session_store, + &mut identity_store, + &mut prekey_store, + &signed_prekey_store, + ); + let resolver = MockSendContextResolver::new(); + let devices = ResolvedDmDevices::new(all, &own_jid, None); + let to = recipient_a.to_non_ad(); + let message = wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }; + + let prepared = prepare_dm_stanza( + &TokioTestRuntime, + &mut stores, + &resolver, + DmStanzaRequest { + own_jid: &own_jid, + account: None, + to: &to, + message: &message, + message_id: "DM_SINK_1", + edit: None, + extra_nodes: &[], + devices: &devices, + pre_encoded: None, + }, + ) + .await + .expect("dm stanza"); + + let participants = prepared + .node + .get_optional_child("participants") + .expect("stanza has a participants node"); + let entries = participants.children().expect("participants has children"); + // Same reasoning as the sink tests: the recipient half drains a + // FuturesUnordered, so which of its devices lands first is not + // promised. The boundary between the halves is, because they are + // sequential awaits, and that is what this test is about. + let written = participant_jids(entries); + assert_eq!( + written.len(), + 3, + "each device contributes exactly one participant node" + ); + let (recipients, own) = written.split_at(2); + assert_eq!( + recipients + .iter() + .cloned() + .collect::>(), + [recipient_a.to_string(), recipient_b.to_string()] + .into_iter() + .collect::>(), + "both recipient devices belong to the first half" + ); + assert_eq!( + own, + [own_companion.to_string()], + "the own-device half lands after the recipient half, in one list" + ); + } + + /// The empty-participants guard still fires when every device drops + /// out: an empty `` would silently drop the message. + #[tokio::test] + async fn a_dm_whose_every_device_fails_is_refused() { + let own_jid: Jid = "5511900000020:0@s.whatsapp.net".parse().unwrap(); + let recipient: Jid = "5511900000021:0@s.whatsapp.net".parse().unwrap(); + + let (mut session_store, mut identity_store) = stores_with_sessions(&[]).await; + let mut prekey_store = UnusedPreKeyStore; + let signed_prekey_store = UnusedSignedPreKeyStore; + let mut sender_key_store = MemSenderKeyStore::default(); + let mut stores = raw_fanout_stores( + &mut sender_key_store, + &mut session_store, + &mut identity_store, + &mut prekey_store, + &signed_prekey_store, + ); + let resolver = MockSendContextResolver::new().with_missing_bundle(recipient.clone()); + let devices = + ResolvedDmDevices::new(vec![recipient.clone(), own_jid.clone()], &own_jid, None); + let to = recipient.to_non_ad(); + let message = wa::Message { + conversation: Some("hi".into()), + ..Default::default() + }; + + let err = prepare_dm_stanza( + &TokioTestRuntime, + &mut stores, + &resolver, + DmStanzaRequest { + own_jid: &own_jid, + account: None, + to: &to, + message: &message, + message_id: "DM_SINK_2", + edit: None, + extra_nodes: &[], + devices: &devices, + pre_encoded: None, + }, + ) + .await + .err() + .expect("a stanza with no participants must not be built"); + assert!( + err.to_string().contains("encryption failed for all"), + "unexpected error: {err}" + ); + } + } + + /// The session phase hands its scratch address on to the single-device + /// encrypt instead of both phases building one for the same device. + mod reused_protocol_address { + use super::*; + + async fn fan_out_one_plan( + devices: &[Jid], + resolver: &MockSendContextResolver, + session_store: &mut MemSessionStore, + identity_store: &mut MemIdentityStore, + plan: Option, + ) -> EncryptForDevicesRaw { + let mut prekey_store = UnusedPreKeyStore; + let signed_prekey_store = UnusedSignedPreKeyStore; + let mut sender_key_store = MemSenderKeyStore::default(); + let mut stores = raw_fanout_stores( + &mut sender_key_store, + session_store, + identity_store, + &mut prekey_store, + &signed_prekey_store, + ); + let plan = match plan { + Some(plan) => plan, + None => { + ensure_sessions_for_devices(&TokioTestRuntime, &mut stores, resolver, devices) + .await + .expect("session phase") + } + }; + encrypt_for_devices_with_sessions_raw( + &TokioTestRuntime, + &mut stores, + devices, + b"payload", + plan, + ) + .await + .expect("encrypt fan-out") + } + + /// The reused buffer must name the same device the session phase just + /// interrogated: a stale or mis-written name resolves no session and + /// the device silently drops out of the fan-out. + #[tokio::test] + async fn the_reused_address_still_names_the_device_it_was_checked_for() { + let device: Jid = "5511900000030:0@s.whatsapp.net".parse().unwrap(); + let (mut session_store, mut identity_store) = + stores_with_sessions(std::slice::from_ref(&device)).await; + let resolver = MockSendContextResolver::new(); + + let raw = fan_out_one_plan( + std::slice::from_ref(&device), + &resolver, + &mut session_store, + &mut identity_store, + None, + ) + .await; + + assert_eq!(raw.devices.len(), 1, "the one device must encrypt"); + assert_eq!(raw.devices[0].device_jid, device); + } + + /// A PN device whose session lives under its LID address: the address + /// the encrypt uses is the overridden (LID) one, not the device's own. + /// Only the LID address has a session, so getting this wrong drops the + /// device. + #[tokio::test] + async fn a_lid_upgraded_device_encrypts_against_its_lid_address() { + let pn: Jid = "5511900000031:0@s.whatsapp.net".parse().unwrap(); + let lid: Jid = "100000000000031:0@lid".parse().unwrap(); + let (mut session_store, mut identity_store) = + stores_with_sessions(std::slice::from_ref(&lid)).await; + let resolver = MockSendContextResolver::new() + .with_phone_to_lid(pn.user.as_str(), lid.user.as_str()); + + let raw = fan_out_one_plan( + std::slice::from_ref(&pn), + &resolver, + &mut session_store, + &mut identity_store, + None, + ) + .await; + + assert_eq!( + raw.devices.len(), + 1, + "only the LID address has a session; the PN address would find none" + ); + assert_eq!( + raw.devices[0].device_jid, pn, + "the wire still names the device, only the Signal address is upgraded" + ); + } + + /// Session state that survives `clone_box`. The session-establishment + /// tasks each get their own clone of the store, so a per-value map + /// would carry their writes away with them and the encrypt that + /// follows would find nothing. + #[derive(Clone, Default)] + struct SharedSessionStore( + std::sync::Arc>>>, + ); + + #[async_trait::async_trait] + impl SessionStore for SharedSessionStore { + async fn load_session(&self, a: &ProtocolAddress) -> SigResult> { + Ok(self + .0 + .lock() + .unwrap() + .get(a) + .and_then(|b| SessionRecord::deserialize(b).ok())) + } + async fn has_session(&self, a: &ProtocolAddress) -> SigResult { + Ok(self.0.lock().unwrap().contains_key(a)) + } + async fn store_session( + &mut self, + a: &ProtocolAddress, + r: SessionRecord, + ) -> SigResult<()> { + self.0.lock().unwrap().insert(a.clone(), r.serialize()?); + Ok(()) + } + } + + /// See [`SharedSessionStore`]. + #[derive(Clone)] + struct SharedIdentityStore { + pair: IdentityKeyPair, + known: std::sync::Arc>>, + } + + #[async_trait::async_trait] + impl IdentityKeyStore for SharedIdentityStore { + async fn get_identity_key_pair(&self) -> SigResult { + Ok(self.pair.clone()) + } + async fn get_local_registration_id(&self) -> SigResult { + Ok(42) + } + async fn save_identity( + &mut self, + a: &ProtocolAddress, + id: &IdentityKey, + ) -> SigResult { + let mut known = self.known.lock().unwrap(); + let changed = known.get(a).is_some_and(|k| k != id); + known.insert(a.clone(), *id); + Ok(IdentityChange::from_changed(changed)) + } + async fn is_trusted_identity( + &self, + _: &ProtocolAddress, + _: &IdentityKey, + _: Direction, + ) -> SigResult { + Ok(true) + } + async fn get_identity(&self, a: &ProtocolAddress) -> SigResult> { + Ok(self.known.lock().unwrap().get(a).copied()) + } + } + + /// The case the reuse must not get wrong: a cold PN device that the + /// session phase upgraded to LID and established a session for. The + /// buffer is left holding the PN name (the last thing the session loop + /// wrote for it), while the encrypt has to address the LID session that + /// was just created. Only rewriting the buffer gets that right. + #[tokio::test] + async fn a_cold_pn_device_upgraded_to_lid_encrypts_against_the_new_lid_session() { + let pn: Jid = "5511900000033:0@s.whatsapp.net".parse().unwrap(); + let lid: Jid = "100000000000033:0@lid".parse().unwrap(); + let mut rng = rand::make_rng::(); + + // No session anywhere yet: the session phase has to create one, and + // it creates it under the LID address. + let mut session_store = SharedSessionStore::default(); + let mut identity_store = SharedIdentityStore { + pair: IdentityKeyPair::generate(&mut rng), + known: Default::default(), + }; + let sessions = session_store.0.clone(); + let mut prekey_store = UnusedPreKeyStore; + let signed_prekey_store = UnusedSignedPreKeyStore; + let mut sender_key_store = MemSenderKeyStore::default(); + let mut stores = SignalStores { + sender_key_store: &mut sender_key_store, + session_store: &mut session_store, + identity_store: &mut identity_store, + prekey_store: &mut prekey_store, + signed_prekey_store: &signed_prekey_store, + }; + let resolver = MockSendContextResolver::new() + .with_phone_to_lid(pn.user.as_str(), lid.user.as_str()) + .with_bundle(pn.normalize_for_prekey_bundle(), signed_prekey_bundle()); + + let plan = ensure_sessions_for_devices( + &TokioTestRuntime, + &mut stores, + &resolver, + std::slice::from_ref(&pn), + ) + .await + .expect("session phase"); + + { + let sessions = sessions.lock().unwrap(); + assert!( + sessions.contains_key(&lid.to_protocol_address()), + "the session phase must have created the LID session" + ); + assert!( + !sessions.contains_key(&pn.to_protocol_address()), + "and nothing under the PN address the buffer was left holding" + ); + } + + let raw = encrypt_for_devices_with_sessions_raw( + &TokioTestRuntime, + &mut stores, + std::slice::from_ref(&pn), + b"payload", + plan, + ) + .await + .expect("encrypt fan-out"); + + assert_eq!( + raw.devices.len(), + 1, + "the freshly established LID session must be the one encrypted against" + ); + assert!( + raw.includes_prekey_message, + "a brand new session emits pkmsg" + ); + } + + /// A plan that never ran a session phase carries no buffer, so the + /// encrypt has to build its own address as before. + #[tokio::test] + async fn a_plan_with_no_session_phase_builds_its_own_address() { + let device: Jid = "5511900000032:0@s.whatsapp.net".parse().unwrap(); + let (mut session_store, mut identity_store) = + stores_with_sessions(std::slice::from_ref(&device)).await; + let resolver = MockSendContextResolver::new(); + + let raw = fan_out_one_plan( + std::slice::from_ref(&device), + &resolver, + &mut session_store, + &mut identity_store, + Some(SessionPlan::assume_ready(1)), + ) + .await; + + assert_eq!(raw.devices.len(), 1); + assert_eq!(raw.devices[0].device_jid, device); + } + + /// The multi-device branch gives every job its own address and must be + /// untouched by the buffer the plan now carries. + #[tokio::test] + async fn several_devices_each_get_their_own_address() { + let devices: Vec = (0..3u16) + .map(|i| format!("551190000004{i}:0@s.whatsapp.net").parse().unwrap()) + .collect(); + let (mut session_store, mut identity_store) = stores_with_sessions(&devices).await; + let resolver = MockSendContextResolver::new(); + + let raw = fan_out_one_plan( + &devices, + &resolver, + &mut session_store, + &mut identity_store, + None, + ) + .await; + + let mut encrypted: Vec = + raw.devices.iter().map(|d| d.device_jid.clone()).collect(); + encrypted.sort_by_key(Jid::to_string); + assert_eq!( + encrypted, devices, + "every device gets its own session address" + ); + } + + /// An empty device list has nothing to name: the plan still carries a + /// buffer and neither branch may touch it. + #[tokio::test] + async fn an_empty_device_list_names_nothing() { + let (mut session_store, mut identity_store) = stores_with_sessions(&[]).await; + let resolver = MockSendContextResolver::new(); + + let raw = fan_out_one_plan( + &[], + &resolver, + &mut session_store, + &mut identity_store, + None, + ) + .await; + + assert!(raw.devices.is_empty()); + assert!(!raw.includes_prekey_message); + } + } } diff --git a/wacore/src/store/signal_cache.rs b/wacore/src/store/signal_cache.rs index ce1af0248..360fec23a 100644 --- a/wacore/src/store/signal_cache.rs +++ b/wacore/src/store/signal_cache.rs @@ -1708,7 +1708,7 @@ mod sender_key_lock_tests { async fn try_put_session_marks_dirty_and_flushes() { let cache = SignalStoreCache::new(); let backend = crate::store::in_memory::InMemoryBackend::new(); - let addr = ProtocolAddress::new("15550009999".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550009999", 1.into()); assert!( cache @@ -1731,7 +1731,7 @@ mod sender_key_lock_tests { #[tokio::test] async fn try_session_paths_fall_back_under_contention() { let cache = SignalStoreCache::new(); - let addr = ProtocolAddress::new("15550009999".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550009999", 1.into()); let guard = cache.sessions.lock().await; assert!( @@ -1770,7 +1770,7 @@ mod sender_key_lock_tests { async fn cancelled_checkout_queues_under_contention_and_remains_flushable() { let cache = SignalStoreCache::new(); let backend = crate::store::in_memory::InMemoryBackend::new(); - let addr = ProtocolAddress::new("15550008888".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550008888", 1.into()); cache.put_session(&addr, SessionRecord::new_fresh()).await; let (record, generation) = cache.checkout_session(&addr, &backend).await.unwrap(); @@ -1801,7 +1801,7 @@ mod sender_key_lock_tests { async fn lossy_clear_rejects_an_older_checkout_generation() { let cache = SignalStoreCache::new(); let backend = crate::store::in_memory::InMemoryBackend::new(); - let addr = ProtocolAddress::new("15550007777".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550007777", 1.into()); cache.put_session(&addr, SessionRecord::new_fresh()).await; let (record, generation) = cache.checkout_session(&addr, &backend).await.unwrap(); @@ -1822,7 +1822,7 @@ mod sender_key_lock_tests { async fn lossy_clear_invalidates_checkouts_before_waiting_for_the_cache() { let cache = Arc::new(SignalStoreCache::new()); let backend = crate::store::in_memory::InMemoryBackend::new(); - let addr = ProtocolAddress::new("15550007776".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550007776", 1.into()); cache.put_session(&addr, SessionRecord::new_fresh()).await; let (record, checkout) = cache.checkout_session(&addr, &backend).await.unwrap(); @@ -1859,7 +1859,7 @@ mod sender_key_lock_tests { #[tokio::test] async fn stale_checkout_cannot_overwrite_a_new_owner() { let cache = SignalStoreCache::new(); - let addr = ProtocolAddress::new("15550007775".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550007775", 1.into()); cache.put_session(&addr, SessionRecord::new_fresh()).await; let (old_record, old_checkout) = cache @@ -1896,7 +1896,7 @@ mod sender_key_lock_tests { #[tokio::test] async fn checkout_rejects_a_competing_owner() { let cache = SignalStoreCache::new(); - let addr = ProtocolAddress::new("15550007770".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550007770", 1.into()); cache.put_session(&addr, SessionRecord::new_fresh()).await; let (record, generation) = cache @@ -1926,7 +1926,7 @@ mod sender_key_lock_tests { async fn restore_does_not_resurrect_a_deleted_slot() { let cache = SignalStoreCache::new(); let backend = crate::store::in_memory::InMemoryBackend::new(); - let addr = ProtocolAddress::new("15550007771".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550007771", 1.into()); cache.put_session(&addr, SessionRecord::new_fresh()).await; let (record, generation) = cache.checkout_session(&addr, &backend).await.unwrap(); @@ -1947,7 +1947,7 @@ mod sender_key_lock_tests { async fn queued_restore_does_not_overwrite_a_delete() { let cache = SignalStoreCache::new(); let backend = crate::store::in_memory::InMemoryBackend::new(); - let addr = ProtocolAddress::new("15550007772".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550007772", 1.into()); cache.put_session(&addr, SessionRecord::new_fresh()).await; let (record, generation) = cache.checkout_session(&addr, &backend).await.unwrap(); @@ -1972,7 +1972,7 @@ mod sender_key_lock_tests { async fn empty_checkout_reserves_and_releases_its_slot() { let cache = SignalStoreCache::new(); let backend = crate::store::in_memory::InMemoryBackend::new(); - let addr = ProtocolAddress::new("15550007773".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550007773", 1.into()); let (record, generation) = cache.checkout_session(&addr, &backend).await.unwrap(); assert!(record.is_none()); @@ -1995,7 +1995,7 @@ mod sender_key_lock_tests { async fn peek_prefers_a_cache_write_that_wins_the_backend_race() { let cache = Arc::new(SignalStoreCache::new()); let backend = Arc::new(BlockingSessionLookup::new()); - let addr = ProtocolAddress::new("15550007774".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550007774", 1.into()); let peek = tokio::spawn({ let cache = cache.clone(); @@ -2014,7 +2014,7 @@ mod sender_key_lock_tests { async fn existence_prefers_a_cache_write_that_wins_the_backend_race() { let cache = Arc::new(SignalStoreCache::new()); let backend = Arc::new(BlockingSessionLookup::new()); - let addr = ProtocolAddress::new("15550007772".to_string(), 2.into()); + let addr = ProtocolAddress::new("15550007772", 2.into()); let exists = tokio::spawn({ let cache = cache.clone(); @@ -2032,7 +2032,7 @@ mod sender_key_lock_tests { #[tokio::test] async fn try_has_session_reports_known_absent() { let cache = SignalStoreCache::new(); - let addr = ProtocolAddress::new("15550009999".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550009999", 1.into()); cache.delete_session(&addr).await; assert_eq!( @@ -2045,7 +2045,7 @@ mod sender_key_lock_tests { #[tokio::test] async fn try_identity_paths_cover_hit_miss_and_contention() { let cache = SignalStoreCache::new(); - let addr = ProtocolAddress::new("15550009999".to_string(), 1.into()); + let addr = ProtocolAddress::new("15550009999", 1.into()); let key_bytes = [7u8; 32]; assert_eq!( @@ -2089,7 +2089,7 @@ mod consumed_prekey_atomicity_tests { .store_prekey(PREKEY_ID, b"durable-prekey", false) .await .unwrap(); - ProtocolAddress::new("bob".to_string(), 1.into()) + ProtocolAddress::new("bob", 1.into()) } /// The inbound pkmsg decrypt promotes the session into the volatile cache and @@ -2194,8 +2194,8 @@ mod consumed_prekey_atomicity_tests { backend.store_prekey(PREKEY_A, b"a", false).await.unwrap(); backend.store_prekey(PREKEY_B, b"b", false).await.unwrap(); - let addr_a = ProtocolAddress::new("alice".to_string(), 1.into()); - let addr_b = ProtocolAddress::new("bob".to_string(), 1.into()); + let addr_a = ProtocolAddress::new("alice", 1.into()); + let addr_b = ProtocolAddress::new("bob", 1.into()); // Both decrypts promote their session (dirty) and buffer their prekey. cache.put_session(&addr_a, SessionRecord::new_fresh()).await; @@ -2650,8 +2650,8 @@ mod consumed_prekey_atomicity_tests { .await .unwrap(); - let addr_a = ProtocolAddress::new("alice".to_string(), 1.into()); - let addr_b = ProtocolAddress::new("bob".to_string(), 1.into()); + let addr_a = ProtocolAddress::new("alice", 1.into()); + let addr_b = ProtocolAddress::new("bob", 1.into()); let cache = StdArc::new(SignalStoreCache::new()); let violation = StdArc::new(AtomicBool::new(false)); @@ -2756,7 +2756,7 @@ mod eviction_tests { use crate::store::in_memory::InMemoryBackend; fn addr(i: usize) -> ProtocolAddress { - ProtocolAddress::new(format!("user{i}@s.whatsapp.net"), DeviceId::new(0)) + ProtocolAddress::new(&format!("user{i}@s.whatsapp.net"), DeviceId::new(0)) } #[test] @@ -2983,8 +2983,8 @@ mod lease_reload_tests { async fn post_flush_clear_preserves_only_live_checkouts() { let backend = InMemoryBackend::new(); let cache = SignalStoreCache::new(); - let active = ProtocolAddress::new("15550001007".to_string(), 1.into()); - let idle = ProtocolAddress::new("15550001008".to_string(), 1.into()); + let active = ProtocolAddress::new("15550001007", 1.into()); + let idle = ProtocolAddress::new("15550001008", 1.into()); cache.put_session(&active, leased_session()).await; cache.put_session(&idle, leased_session()).await; cache.flush(&backend).await.expect("flush"); @@ -3020,7 +3020,7 @@ mod lease_reload_tests { DEFAULT_MAX_CACHE_ENTRIES, [0xA1; 16], ); - let address = ProtocolAddress::new("15550001001".to_string(), 1.into()); + let address = ProtocolAddress::new("15550001001", 1.into()); cache.put_session(&address, leased_session()).await; cache.flush(&backend).await.expect("flush"); cache.clear_after_flush().await; @@ -3054,7 +3054,7 @@ mod lease_reload_tests { DEFAULT_MAX_CACHE_ENTRIES, [0xA1; 16], ); - let address = ProtocolAddress::new("15550001002".to_string(), 1.into()); + let address = ProtocolAddress::new("15550001002", 1.into()); cache.put_session(&address, leased_session()).await; cache.flush(&backend).await.expect("initial flush"); @@ -3271,7 +3271,7 @@ mod pre_wire_gate_tests { use async_lock::Barrier; fn addr(user: &str) -> ProtocolAddress { - ProtocolAddress::new(user.to_string(), 1.into()) + ProtocolAddress::new(user, 1.into()) } fn leased_record() -> SessionRecord { diff --git a/wacore/src/store/signal_cache_durability_chaos.rs b/wacore/src/store/signal_cache_durability_chaos.rs index b228a7bdc..108f93566 100644 --- a/wacore/src/store/signal_cache_durability_chaos.rs +++ b/wacore/src/store/signal_cache_durability_chaos.rs @@ -135,7 +135,7 @@ impl ChaosHarness { 32, incarnation(seed ^ 0xA5A5_A5A5_A5A5_A5A5, 0), ), - dm_address: ProtocolAddress::new("15550007001".to_string(), 1.into()), + dm_address: ProtocolAddress::new("15550007001", 1.into()), group_name: SenderKeyName::from_parts( "120363000000070001@g.us", "15550007002@s.whatsapp.net:0", diff --git a/wacore/src/types/jid.rs b/wacore/src/types/jid.rs index 3a75afdfb..eb847d2e5 100644 --- a/wacore/src/types/jid.rs +++ b/wacore/src/types/jid.rs @@ -1,4 +1,4 @@ -use crate::libsignal::protocol::{DeviceId, ProtocolAddress}; +use crate::libsignal::protocol::{AddressBuf, DeviceId, ProtocolAddress}; use crate::libsignal::store::sender_key_name::SenderKeyName; use wacore_binary::{DEFAULT_USER_SERVER, Jid, LEGACY_USER_SERVER}; @@ -26,15 +26,57 @@ pub fn make_address_buffer() -> String { String::with_capacity(SIGNAL_ADDRESS_CAPACITY) } -/// Create a pre-allocated `ProtocolAddress` for hot loops. +/// Create a reusable `ProtocolAddress` for hot loops. /// Call `reset_protocol_address` to fill without allocation. pub fn make_reusable_protocol_address() -> ProtocolAddress { - ProtocolAddress::with_capacity(SIGNAL_ADDRESS_CAPACITY, SIGNAL_DEVICE_ID) + ProtocolAddress::empty(SIGNAL_DEVICE_ID) +} + +/// Somewhere an address name can be written. +/// +/// The address format lives in exactly one function, and that function has to +/// serve both a plain `String` and the buffer inside a `ProtocolAddress` (which +/// is inline, not a `String`). This is what lets it do that without the format +/// existing in two places. +pub trait AddressSink { + fn clear(&mut self); + fn push_str(&mut self, s: &str); + fn push(&mut self, c: char); +} + +impl AddressSink for String { + #[inline] + fn clear(&mut self) { + String::clear(self); + } + #[inline] + fn push_str(&mut self, s: &str) { + String::push_str(self, s); + } + #[inline] + fn push(&mut self, c: char) { + String::push(self, c); + } +} + +impl AddressSink for AddressBuf { + #[inline] + fn clear(&mut self) { + AddressBuf::clear(self); + } + #[inline] + fn push_str(&mut self, s: &str) { + AddressBuf::push_str(self, s); + } + #[inline] + fn push(&mut self, c: char) { + AddressBuf::push(self, c); + } } /// Write the signal address name (`{user}[:device]@{server}`) into `buf`, /// clearing it first. All other address helpers delegate to this. -pub fn write_signal_address_to(jid: &Jid, buf: &mut String) { +pub fn write_signal_address_to(jid: &Jid, buf: &mut W) { buf.clear(); let server = mapped_server(jid.server.as_str()); buf.push_str(&jid.user); @@ -47,7 +89,7 @@ pub fn write_signal_address_to(jid: &Jid, buf: &mut String) { } /// Write the full protocol address (`{signal_address}.0`) into `buf`. -pub fn write_protocol_address_to(jid: &Jid, buf: &mut String) { +pub fn write_protocol_address_to(jid: &Jid, buf: &mut W) { write_signal_address_to(jid, buf); buf.push_str(".0"); } @@ -111,7 +153,11 @@ impl JidExt for Jid { } fn to_protocol_address(&self) -> ProtocolAddress { - ProtocolAddress::new(self.to_signal_address_string(), SIGNAL_DEVICE_ID) + // Written straight into the address: the intermediate `String` this + // used to build was allocated only to be copied in and dropped. + let mut addr = make_reusable_protocol_address(); + self.reset_protocol_address(&mut addr); + addr } fn to_protocol_address_string(&self) -> String { @@ -217,6 +263,56 @@ mod tests { } } + /// The one writer must produce the same bytes into either sink, or the + /// heap path and the inline path would name the same device differently. + #[test] + fn both_sinks_receive_the_same_address() { + let cases = [ + "123456789@lid", + "123456789:33@lid", + "100000000000001.1:75@lid", + "15550000001@s.whatsapp.net", + "15550000001:33@s.whatsapp.net", + "120363000000000001@g.us", + "999999999999999999@newsletter", + ]; + for jid_str in cases { + let jid = Jid::from_str(jid_str).unwrap(); + + let mut string_sink = String::new(); + write_protocol_address_to(&jid, &mut string_sink); + + let mut address = make_reusable_protocol_address(); + jid.reset_protocol_address(&mut address); + + assert_eq!( + address.as_str(), + string_sink, + "the two sinks disagree for {jid_str}" + ); + } + } + + /// Reusing one buffer across JIDs must leave no trace of the previous one, + /// including when the previous name was longer. + #[test] + fn a_reused_address_keeps_nothing_from_the_previous_jid() { + let long = Jid::from_str("100000000000001.1:75@lid").unwrap(); + let short = Jid::from_str("1@lid").unwrap(); + + let mut address = make_reusable_protocol_address(); + address.reset_with(|buf| buf.push_str(&"z".repeat(200))); + jid_reset(&mut address, &long); + assert_eq!(address.as_str(), "100000000000001.1:75@lid.0"); + jid_reset(&mut address, &short); + assert_eq!(address.as_str(), "1@lid.0"); + assert_eq!(address.name(), "1@lid"); + } + + fn jid_reset(address: &mut ProtocolAddress, jid: &Jid) { + jid.reset_protocol_address(address); + } + #[test] fn test_write_functions_dry() { let jid = Jid::from_str("15550000001@s.whatsapp.net").unwrap(); @@ -228,4 +324,19 @@ mod tests { write_protocol_address_to(&jid, &mut buf); assert_eq!(buf, "15550000001@c.us.0"); } + + /// The writer's "clears it first" contract holds for the inline sink too: + /// a reused buffer must be overwritten, not appended to. + #[test] + fn the_inline_sink_is_cleared_before_each_write() { + let first = Jid::from_str("15550000001@s.whatsapp.net").unwrap(); + let second = Jid::from_str("123456789:33@lid").unwrap(); + + let mut buf = AddressBuf::empty(); + write_signal_address_to(&first, &mut buf); + assert_eq!(buf.as_str(), "15550000001@c.us"); + + write_protocol_address_to(&second, &mut buf); + assert_eq!(buf.as_str(), "123456789:33@lid.0"); + } }