diff --git a/src/client/device_registry.rs b/src/client/device_registry.rs index f158bb8d5..b14966daf 100644 --- a/src/client/device_registry.rs +++ b/src/client/device_registry.rs @@ -665,11 +665,27 @@ impl Client { /// WA Web: `isFromKnownDevice(author)` — local check only, no network. pub(crate) async fn is_from_known_device(&self, sender: &Jid) -> bool { - self.has_device(&sender.user, sender.device).await + self.has_device_for_jid(sender, sender.device).await + } + + /// `has_device` for a caller holding the full `Jid`: the namespace picks + /// the single `lid_pn_cache` probe (see `resolve_lookup_keys_for_jid`). + /// Runs on every successful group decrypt and every retry receipt. + pub(crate) async fn has_device_for_jid(&self, jid: &Jid, device_id: u16) -> bool { + if device_id == 0 { + return true; + } + let lookup = self.resolve_lookup_keys_for_jid(jid).await; + self.has_device_in(&lookup, device_id).await } /// Check if a device exists for a user. /// Returns true for device_id 0 (primary device always exists). + /// + /// Every production caller holds a `Jid` and goes through + /// `has_device_for_jid`; this bare-user form remains for the tests that + /// probe a user under both of its namespaces. + #[cfg(test)] pub(crate) async fn has_device(&self, user: &str, device_id: u16) -> bool { if device_id == 0 { return true; @@ -677,7 +693,10 @@ impl Client { // Borrowed keys avoid allocating the owned lookup variants on this hot path. let lookup = self.resolve_lookup_keys(user).await; + self.has_device_in(&lookup, device_id).await + } + async fn has_device_in(&self, lookup: &UserLookupKeys, device_id: u16) -> bool { for key in lookup.all_keys() { if let Some(record) = self.device_registry_cache.get(key).await { return record.devices.iter().any(|d| d.device_id() == device_id); @@ -1255,7 +1274,24 @@ impl Client { user: &str, ) -> Option { let lookup = self.resolve_lookup_keys(user).await; + self.load_device_record_in(&lookup).await + } + + /// `load_device_record` for a caller holding the full `Jid`, so the known + /// namespace costs one `lid_pn_cache` probe instead of two per user of a + /// device-list response. + pub(crate) async fn load_device_record_for_jid( + &self, + jid: &Jid, + ) -> Option { + let lookup = self.resolve_lookup_keys_for_jid(jid).await; + self.load_device_record_in(&lookup).await + } + async fn load_device_record_in( + &self, + lookup: &UserLookupKeys, + ) -> Option { for key in lookup.all_keys() { if let Some(record) = self.device_registry_cache.get(key).await { // Cold load-modify-persist path: callers mutate the owned record. diff --git a/src/message.rs b/src/message.rs index e5c6158fa..b4fbeb872 100644 --- a/src/message.rs +++ b/src/message.rs @@ -94,10 +94,12 @@ impl EncPayload { ciphertext: bytes::Bytes, enc_node: &NodeRef<'_>, enc_index: usize, + enc_type: EncType, ) -> Option { - let enc_type = EncType::from_wire(enc_node.attrs().optional_string("type")?.as_ref())?; - let padding_version = enc_node.attrs().optional_u64("v").unwrap_or(2) as u8; + // One attribute pass: the caller already classified `type`, so only the + // remaining attributes are read here, through a single parser. let mut attrs = enc_node.attrs(); + let padding_version = attrs.optional_u64("v").unwrap_or(2) as u8; Some(Self { ciphertext, enc_type, @@ -111,25 +113,32 @@ impl EncPayload { } /// Zero-copy extraction from an OwnedNodeRef. + /// + /// `enc_type` is the node's already-parsed `type` attribute; the receive + /// loop validates it before calling, so it is not re-read here. pub(crate) fn from_owned_node( owner: &OwnedNodeRef, enc_node: &NodeRef<'_>, enc_index: usize, + enc_type: EncType, ) -> Option { Self::from_parts( owner.slice_bytes(enc_node.content_bytes()?), enc_node, enc_index, + enc_type, ) } /// Copying extraction from a NodeRef (used in tests where there's no OwnedNodeRef). #[cfg(test)] pub(crate) fn from_node_ref(node: &NodeRef<'_>, enc_index: usize) -> Option { + let enc_type = EncType::from_wire(node.attrs().optional_string("type")?.as_ref())?; Self::from_parts( bytes::Bytes::copy_from_slice(node.content_bytes()?), node, enc_index, + enc_type, ) } } diff --git a/src/message/receive.rs b/src/message/receive.rs index b26e5743d..35e5659c3 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -331,7 +331,7 @@ impl Client { // `had_unknown_enc` means "produced no usable payload": either the // type is unrecognized or it's known but the body is empty. // Either way the stanza needs the fallback ack or the server replays. - if EncType::from_wire(enc_type.as_ref()).is_none() { + let Some(parsed_enc_type) = EncType::from_wire(enc_type.as_ref()) else { log::warn!("Enc node has unknown type: {enc_type}"); self.report_raw_enc_decrypt_failure( &info, @@ -341,23 +341,24 @@ impl Client { ); had_unknown_enc = true; continue; - } - - let payload = match EncPayload::from_owned_node(node, enc_node, enc_index) { - Some(p) => p, - None => { - log::warn!("Enc node {enc_type} has no content"); - self.report_raw_enc_decrypt_failure( - &info, - enc_index, - Some(enc_type.as_ref()), - EncDecryptFailureReason::MalformedNode, - ); - had_unknown_enc = true; - continue; - } }; + let payload = + match EncPayload::from_owned_node(node, enc_node, enc_index, parsed_enc_type) { + Some(p) => p, + None => { + log::warn!("Enc node {enc_type} has no content"); + self.report_raw_enc_decrypt_failure( + &info, + enc_index, + Some(enc_type.as_ref()), + EncDecryptFailureReason::MalformedNode, + ); + had_unknown_enc = true; + continue; + } + }; + let bucket = if payload.enc_type.is_bot_secret() { &mut bot_payloads } else if payload.enc_type.is_session() { diff --git a/src/portable_cache.rs b/src/portable_cache.rs index 7fa74a563..0da029297 100644 --- a/src/portable_cache.rs +++ b/src/portable_cache.rs @@ -96,7 +96,14 @@ where } } - fn remove_key(&mut self, key: &K) -> Option> { + /// Borrowed removal: the `order` side is keyed by the entry's own `seq`, + /// so nothing here ever needs an owned `K`, and callers with a `&str` or + /// `&Jid` need not clone the key just to delete it. + fn remove_key(&mut self, key: &Q) -> Option> + where + K: Borrow, + Q: Hash + Eq + ?Sized, + { let entry = self.map.remove(key)?; self.order.remove(&entry.seq); Some(entry) @@ -385,14 +392,6 @@ where false } - fn find_key(inner: &CacheInner, key: &Q) -> Option - where - K: Borrow, - Q: Hash + Eq + ?Sized, - { - inner.map.get_key_value(key).map(|(k, _)| k.clone()) - } - /// Whether `entry`'s access stamp has aged past [`TTI_RENEWAL_DIVISOR`]'s /// tolerance and is worth pushing forward under the write lock. A cache /// without TTI never renews. @@ -420,14 +419,13 @@ where // place. Removing without it could drop a replacement written // while the guard was down and judge it by a stale `now`. let observed = (entry.seq, entry.inserted_at); - let owned_key = Self::find_key(&guard, key)?; drop(guard); let mut wguard = self.inner.write().await; if let Some(e) = wguard.map.get(key) && (e.seq, e.inserted_at) == observed && self.is_expired(e, now) { - wguard.remove_key(&owned_key); + wguard.remove_key(key); } return None; } @@ -495,9 +493,8 @@ where .map .get(key) .is_some_and(|entry| self.is_expired(entry, now)) - && let Some(owned_key) = Self::find_key(&guard, key) { - guard.remove_key(&owned_key); + guard.remove_key(key); } let (next, result) = update(guard.map.get(key).map(|entry| &entry.value)); @@ -550,8 +547,7 @@ where Q: Hash + Eq + ?Sized, { let mut guard = self.inner.write().await; - let owned_key = Self::find_key(&guard, key)?; - let entry = guard.remove_key(&owned_key)?; + let entry = guard.remove_key(key)?; // Nothing to date until an entry is actually in hand. let now = self.entry_time(); if self.is_expired(&entry, now) { @@ -567,9 +563,7 @@ where Q: Hash + Eq + ?Sized, { let mut guard = self.inner.write().await; - if let Some(owned_key) = Self::find_key(&guard, key) { - guard.remove_key(&owned_key); - } + guard.remove_key(key); } /// Reliably remove all entries, awaiting the write lock. Prefer this in diff --git a/src/receipt.rs b/src/receipt.rs index fd0f65b21..8b43e5138 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -1,5 +1,5 @@ use crate::client::Client; -use crate::types::events::{Event, Receipt}; +use crate::types::events::{Event, EventKind, Receipt}; use crate::types::message::MessageInfo; use crate::types::presence::ReceiptType; use log::debug; @@ -628,6 +628,12 @@ impl Client { from.observe(), users.len() ); + // Pure event production from here on: `dispatch` would drop every + // one of these on the floor without a subscriber, so skip building + // N receipts (each with a `Jid` clone and a `Vec`) up front. + if !self.core.event_bus.has_handler_for(EventKind::Receipt) { + return; + } for user in users { // Missing `` means the server didn't disambiguate the // per-user time; fall back to the stanza-level `t`. diff --git a/src/retry.rs b/src/retry.rs index 9b716187c..a035fb473 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -524,7 +524,7 @@ impl Client { // evicted retry still triggers it. let sender_device_id = info.requester.device(); let device_known = self - .has_device(&info.requester.user, sender_device_id) + .has_device_for_jid(&info.requester, sender_device_id) .await; if !device_known { // Parity with WA Web's MdRetryFromUnknownDevice WAM (id 2178), which diff --git a/src/usync.rs b/src/usync.rs index 619b18d01..6ab5a7b86 100644 --- a/src/usync.rs +++ b/src/usync.rs @@ -288,7 +288,7 @@ impl Client { // Preserve key_index values from existing records (set via account_sync) // Use alias-aware lookup (resolves LID ↔ PN) to find // existing record regardless of which key it was stored under - let mut existing_record = self.load_device_record(&user_list.user.user).await; + let mut existing_record = self.load_device_record_for_jid(&user_list.user).await; // Decode key-index-list if present (WA Web: handleKeyIndexResult) let decoded_key_index = user_list @@ -450,7 +450,7 @@ impl Client { for own in device_snapshot.pn.iter().chain(device_snapshot.lid.iter()) { let bare = own.to_non_ad(); // Carry the cached device_hash so an unchanged list is skipped server-side. - if let Some(record) = self.load_device_record(&bare.user).await + if let Some(record) = self.load_device_record_for_jid(&bare).await && let Some(phash) = record.phash { hashes.insert(bare.clone(), (String::from(phash), record.timestamp)); diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index ec91b41a5..5a9fe5dc3 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -1218,8 +1218,9 @@ impl SqliteStore { ) -> Result<()> { let pool = self.pool.clone(); let db_semaphore = self.db_semaphore.clone(); - let address_owned = address.to_string(); - let key_vec = key.to_vec(); + // The key is a `Copy` array and the address is refcount-shared, so an + // attempt costs no heap allocation beyond the closure itself. + let address_owned: Arc = Arc::from(address); const MAX_RETRIES: u32 = 5; @@ -1232,7 +1233,6 @@ impl SqliteStore { let pool_clone = pool.clone(); let address_clone = address_owned.clone(); - let key_clone = key_vec.clone(); let result = crate::pool::spawn_blocking(move || -> std::result::Result<(), DieselOrStore> { @@ -1241,13 +1241,13 @@ impl SqliteStore { .map_err(|e| DieselOrStore::Store(StoreError::Connection(Box::new(e))))?; diesel::insert_into(identities::table) .values(( - identities::address.eq(address_clone), - identities::key.eq(&key_clone[..]), + identities::address.eq(address_clone.as_ref()), + identities::key.eq(&key[..]), identities::device_id.eq(device_id), )) .on_conflict((identities::address, identities::device_id)) .do_update() - .set(identities::key.eq(&key_clone[..])) + .set(identities::key.eq(&key[..])) .execute(&mut *conn) .map_err(DieselOrStore::Diesel)?; Ok(()) @@ -1350,8 +1350,11 @@ impl SqliteStore { ) -> Result<()> { let pool = self.pool.clone(); let db_semaphore = self.db_semaphore.clone(); - let address_owned = address.to_string(); - let session_vec = session.to_vec(); + // Copied once, then refcount-shared across attempts: this runs after + // every Signal encrypt/decrypt, and a session record is several KiB, + // so a per-attempt `Vec` clone was a memcpy on the happy path too. + let address_owned: Arc = Arc::from(address); + let session_bytes = Bytes::copy_from_slice(session); const MAX_RETRIES: u32 = 5; @@ -1364,7 +1367,7 @@ impl SqliteStore { let pool_clone = pool.clone(); let address_clone = address_owned.clone(); - let session_clone = session_vec.clone(); + let session_clone = session_bytes.clone(); let result = crate::pool::spawn_blocking(move || -> std::result::Result<(), DieselOrStore> { @@ -1373,13 +1376,13 @@ impl SqliteStore { .map_err(|e| DieselOrStore::Store(StoreError::Connection(Box::new(e))))?; diesel::insert_into(sessions::table) .values(( - sessions::address.eq(address_clone), - sessions::record.eq(&session_clone), + sessions::address.eq(address_clone.as_ref()), + sessions::record.eq(session_clone.as_ref()), sessions::device_id.eq(device_id), )) .on_conflict((sessions::address, sessions::device_id)) .do_update() - .set(sessions::record.eq(&session_clone)) + .set(sessions::record.eq(session_clone.as_ref())) .execute(&mut *conn) .map_err(DieselOrStore::Diesel)?; Ok(()) @@ -2046,7 +2049,8 @@ impl SignalStore for SqliteStore { let pool = self.pool.clone(); let db_semaphore = self.db_semaphore.clone(); let device_id = self.device_id; - let record = record.to_vec(); + // One copy, then refcount clones per attempt (see put_session_for_device). + let record = Bytes::copy_from_slice(record); const MAX_RETRIES: u32 = 5; @@ -2068,14 +2072,14 @@ impl SignalStore for SqliteStore { diesel::insert_into(prekeys::table) .values(( prekeys::id.eq(id as i32), - prekeys::key.eq(&record_clone), + prekeys::key.eq(record_clone.as_ref()), prekeys::uploaded.eq(uploaded), prekeys::device_id.eq(device_id), )) .on_conflict((prekeys::id, prekeys::device_id)) .do_update() .set(( - prekeys::key.eq(&record_clone), + prekeys::key.eq(record_clone.as_ref()), prekeys::uploaded.eq(uploaded), )) .execute(&mut *conn) @@ -2315,7 +2319,8 @@ impl SignalStore for SqliteStore { let pool = self.pool.clone(); let db_semaphore = self.db_semaphore.clone(); let device_id = self.device_id; - let record = record.to_vec(); + // One copy, then refcount clones per attempt (see put_session_for_device). + let record = Bytes::copy_from_slice(record); const MAX_RETRIES: u32 = 5; @@ -2337,12 +2342,12 @@ impl SignalStore for SqliteStore { diesel::insert_into(signed_prekeys::table) .values(( signed_prekeys::id.eq(id as i32), - signed_prekeys::record.eq(&record_clone), + signed_prekeys::record.eq(record_clone.as_ref()), signed_prekeys::device_id.eq(device_id), )) .on_conflict((signed_prekeys::id, signed_prekeys::device_id)) .do_update() - .set(signed_prekeys::record.eq(&record_clone)) + .set(signed_prekeys::record.eq(record_clone.as_ref())) .execute(&mut *conn) .map_err(DieselOrStore::Diesel)?; Ok(()) diff --git a/wacore/appstate/src/hash.rs b/wacore/appstate/src/hash.rs index a40916b50..60f9bcacd 100644 --- a/wacore/appstate/src/hash.rs +++ b/wacore/appstate/src/hash.rs @@ -1,5 +1,8 @@ +use hmac::digest::KeyInit; +use hmac::{Hmac, Mac}; use serde::{Deserialize, Serialize}; use serde_big_array::BigArray; +use sha2::{Sha256, Sha512}; use std::collections::HashMap; use wacore_libsignal::crypto::CryptographicMac; use waproto::whatsapp as wa; @@ -480,25 +483,28 @@ pub fn generate_content_mac( // We mirror that exactly so the HMAC input is bytewise identical. let mut key_data_length = [0u8; 8]; key_data_length[7] = ((key_id.len() + 1) & 0xff) as u8; - let mut mac = - CryptographicMac::new("HmacSha512", key).expect("HmacSha512 is a valid algorithm"); + // Typed HMACs rather than `CryptographicMac::new("...")`: these two run per + // decoded record, and the string-dispatched enum costs a name compare chain + // plus a Sha512-sized stack object on every call. Output is byte-identical. + let mut mac = Hmac::::new_from_slice(key).expect("HMAC accepts any key length"); mac.update(&op_byte); mac.update(key_id); mac.update(data); mac.update(&key_data_length); - let mut out = [0u8; 64]; - mac.finalize_into(&mut out) - .expect("64 bytes is enough for HmacSha512"); + let out = mac.finalize().into_bytes(); let mut result = [0u8; 32]; result.copy_from_slice(&out[..32]); result } -pub fn generate_index_mac(index_json_bytes: &[u8], key: &[u8; 32]) -> Vec { - let mut mac = - CryptographicMac::new("HmacSha256", key).expect("HmacSha256 is a valid algorithm"); +fn index_mac_array(index_json_bytes: &[u8], key: &[u8; 32]) -> [u8; 32] { + let mut mac = Hmac::::new_from_slice(key).expect("HMAC accepts any key length"); mac.update(index_json_bytes); - mac.finalize() + mac.finalize().into_bytes().into() +} + +pub fn generate_index_mac(index_json_bytes: &[u8], key: &[u8; 32]) -> Vec { + index_mac_array(index_json_bytes, key).to_vec() } pub fn validate_index_mac( @@ -506,7 +512,8 @@ pub fn validate_index_mac( expected_mac: &[u8], key: &[u8; 32], ) -> Result<(), AppStateError> { - if generate_index_mac(index_json_bytes, key).as_slice() != expected_mac { + // Compare against the stack array: no heap allocation per validated record. + if index_mac_array(index_json_bytes, key) != expected_mac { Err(AppStateError::MismatchingIndexMAC) } else { Ok(()) diff --git a/wacore/appstate/src/processor.rs b/wacore/appstate/src/processor.rs index 3813ef299..5cbbe4394 100644 --- a/wacore/appstate/src/processor.rs +++ b/wacore/appstate/src/processor.rs @@ -409,9 +409,21 @@ where } // Decode all mutations and collect MACs in a single pass + // SET and REMOVE are disjoint, and a patch is almost always all one or the + // other, so sizing both lists to the full count wasted one allocation. + let sets = patch + .mutations + .iter() + .filter(|m| { + matches!( + known_op(m.operation), + Ok(wa::syncd_mutation::SyncdOperation::SET) + ) + }) + .count(); let mut mutations = Vec::with_capacity(patch.mutations.len()); - let mut added_macs = Vec::with_capacity(patch.mutations.len()); - let mut removed_index_macs = Vec::with_capacity(patch.mutations.len()); + let mut added_macs = Vec::with_capacity(sets); + let mut removed_index_macs = Vec::with_capacity(patch.mutations.len() - sets); for m in &patch.mutations { if m.record.is_set() { diff --git a/wacore/libsignal/src/crypto/aes_gcm.rs b/wacore/libsignal/src/crypto/aes_gcm.rs index 356719aae..72220536b 100644 --- a/wacore/libsignal/src/crypto/aes_gcm.rs +++ b/wacore/libsignal/src/crypto/aes_gcm.rs @@ -73,11 +73,10 @@ impl GcmGhash { let leftover = msg.len() - 16 * full_blocks; assert!(leftover < TAG_SIZE); - let (chunks, _) = msg[..16 * full_blocks].as_chunks::<16>(); - for chunk in chunks { - let block: ghash::Block = (*chunk).into(); - self.ghash.update(std::slice::from_ref(&block)); - } + // One call for the whole run: `update_padded` on a block-multiple + // slice is exactly `update(blocks)` with no padding, and lets the + // carryless-multiply backend batch instead of taking one block per call. + self.ghash.update_padded(&msg[..16 * full_blocks]); self.msg_buf[0..leftover].copy_from_slice(&msg[full_blocks * 16..]); self.msg_buf_offset = leftover; diff --git a/wacore/libsignal/src/crypto/provider.rs b/wacore/libsignal/src/crypto/provider.rs index e37240a04..faf28694a 100644 --- a/wacore/libsignal/src/crypto/provider.rs +++ b/wacore/libsignal/src/crypto/provider.rs @@ -329,8 +329,11 @@ impl SignalCryptoProvider for RustCryptoProvider { let encrypted_size = plaintext.len() + padding; let start = out.len(); + // Append the plaintext first and zero only the padding tail: a + // `resize` over the whole length memsets bytes the next line overwrites. + out.reserve(encrypted_size); + out.extend_from_slice(plaintext); out.resize(start + encrypted_size, 0); - out[start..start + plaintext.len()].copy_from_slice(plaintext); let encryptor = cbc::Encryptor::::new(key.into(), iv.into()); let written = encryptor @@ -410,16 +413,18 @@ impl SignalCryptoProvider for RustCryptoProvider { } let (ct, tag) = ciphertext_with_tag.split_at(ciphertext_with_tag.len() - TAG); - // Decrypt into a scratch, verify tag; only commit to `out` on success - // so failures leave it untouched. - let mut scratch = ct.to_vec(); + // Decrypt in place on `out`'s own tail, then verify the tag. A failure + // truncates back to the original length, so callers still observe + // `out` untouched, without a scratch allocation and a second copy. + let start = out.len(); let mut dec = Aes256GcmDecryption::new(key, nonce, aad).map_err(|_| CryptoProviderError::BadInput)?; - dec.decrypt(&mut scratch); - dec.verify_tag(tag) - .map_err(|_| CryptoProviderError::AuthFailed)?; - - out.extend_from_slice(&scratch); + out.extend_from_slice(ct); + dec.decrypt(&mut out[start..]); + if dec.verify_tag(tag).is_err() { + out.truncate(start); + return Err(CryptoProviderError::AuthFailed); + } Ok(()) } diff --git a/wacore/noise/src/state.rs b/wacore/noise/src/state.rs index a1b1e2965..84d62b27f 100644 --- a/wacore/noise/src/state.rs +++ b/wacore/noise/src/state.rs @@ -118,7 +118,12 @@ impl NoiseState { /// Per Noise spec § 5.2: when `protocol_name` is ≤ HASHLEN bytes, append /// zero bytes to make HASHLEN; otherwise hash with SHA256. pub fn new(pattern: impl AsRef<[u8]>, prologue: &[u8]) -> Result { - let pattern = pattern.as_ref(); + Self::new_inner(pattern.as_ref(), prologue) + } + + // The generic shell above is instantiated once per argument type at the + // call sites; the body lives here so it is compiled once. + fn new_inner(pattern: &[u8], prologue: &[u8]) -> Result { let h: [u8; 32] = if pattern.len() <= 32 { let mut h = [0u8; 32]; h[..pattern.len()].copy_from_slice(pattern);