diff --git a/src/features/signal.rs b/src/features/signal.rs index 2220a0ec5..89c9071c8 100644 --- a/src/features/signal.rs +++ b/src/features/signal.rs @@ -54,15 +54,14 @@ impl<'a> Signal<'a> { drop(_guard); self.client.flush_signal_cache().await?; - match encrypted { - CiphertextMessage::PreKeySignalMessage(msg) => { - Ok((EncType::PreKeyMessage, msg.serialized().to_vec())) - } - CiphertextMessage::SignalMessage(msg) => { - Ok((EncType::Message, msg.serialized().to_vec())) - } - _ => Err(anyhow!("unexpected ciphertext variant")), - } + let (_, is_prekey, bytes) = wacore::send::extract_ciphertext(encrypted) + .ok_or_else(|| anyhow!("unexpected ciphertext variant"))?; + let enc_type = if is_prekey { + EncType::PreKeyMessage + } else { + EncType::Message + }; + Ok((enc_type, bytes.into_vec())) } /// Decrypt a Signal protocol message from a sender. @@ -173,7 +172,7 @@ impl<'a> Signal<'a> { self.client.flush_signal_cache().await?; - Ok((skdm_bytes, ciphertext.serialized().to_vec())) + Ok((skdm_bytes, ciphertext.into_serialized().into_vec())) } /// Decrypt a group (sender-key) message. diff --git a/wacore/libsignal/src/protocol/protocol.rs b/wacore/libsignal/src/protocol/protocol.rs index 34614197b..58104abe8 100644 --- a/wacore/libsignal/src/protocol/protocol.rs +++ b/wacore/libsignal/src/protocol/protocol.rs @@ -133,6 +133,11 @@ impl SignalMessage { &self.serialized } + #[inline] + pub fn into_serialized(self) -> Box<[u8]> { + self.serialized + } + #[inline] pub fn body(&self) -> &[u8] { &self.ciphertext @@ -320,6 +325,11 @@ impl PreKeySignalMessage { pub fn serialized(&self) -> &[u8] { &self.serialized } + + #[inline] + pub fn into_serialized(self) -> Box<[u8]> { + self.serialized + } } impl AsRef<[u8]> for PreKeySignalMessage { @@ -504,6 +514,11 @@ impl SenderKeyMessage { pub fn serialized(&self) -> &[u8] { &self.serialized } + + #[inline] + pub fn into_serialized(self) -> Box<[u8]> { + self.serialized + } } impl AsRef<[u8]> for SenderKeyMessage { @@ -628,6 +643,11 @@ impl SenderKeyDistributionMessage { pub fn serialized(&self) -> &[u8] { &self.serialized } + + #[inline] + pub fn into_serialized(self) -> Box<[u8]> { + self.serialized + } } impl AsRef<[u8]> for SenderKeyDistributionMessage { diff --git a/wacore/src/messages.rs b/wacore/src/messages.rs index 1d4b5a172..b814d22ae 100644 --- a/wacore/src/messages.rs +++ b/wacore/src/messages.rs @@ -7,40 +7,46 @@ use waproto::whatsapp as wa; pub struct MessageUtils; impl MessageUtils { - pub fn pad_message_v2(mut plaintext: Vec) -> Vec { + fn random_pad_len() -> u8 { use rand::RngExt; let mut rng = rand::make_rng::(); + let v = rng.random::() & 0x0F; + if v == 0 { 0x0F } else { v } + } - let mut pad_val = rng.random::() & 0x0F; - if pad_val == 0 { - pad_val = 0x0F; - } - - let padding = vec![pad_val; pad_val as usize]; - plaintext.extend_from_slice(&padding); + pub fn pad_message_v2(mut plaintext: Vec) -> Vec { + let pad = Self::random_pad_len(); + plaintext.resize(plaintext.len() + pad as usize, pad); plaintext } + /// Encode + pad in a single pre-sized allocation. + pub fn encode_and_pad(msg: &wa::Message) -> Vec { + let pad = Self::random_pad_len(); + let mut buf = Vec::with_capacity(msg.encoded_len() + pad as usize); + msg.encode(&mut buf).expect("encode into pre-sized Vec"); + buf.resize(buf.len() + pad as usize, pad); + buf + } + pub fn participant_list_hash(devices: &[wacore_binary::jid::Jid]) -> Result { + // Hash sorted ad_strings incrementally (avoids join() allocation). let mut jids: Vec = devices.iter().map(|j| j.to_ad_string()).collect(); - jids.sort(); - - let concatenated_jids = jids.join(""); + jids.sort_unstable(); - // Use finalize_sha256_array() for zero-allocation hash finalization - let full_hash = { - let mut h = CryptographicHash::new("SHA-256") - .map_err(|e| anyhow!("failed to initialize SHA-256 hasher: {:?}", e))?; - h.update(concatenated_jids.as_bytes()); - h.finalize_sha256_array() - .map_err(|e| anyhow!("failed to finalize hash: {:?}", e))? - }; + let mut h = CryptographicHash::new("SHA-256") + .map_err(|e| anyhow!("failed to initialize SHA-256 hasher: {:?}", e))?; + for jid in &jids { + h.update(jid.as_bytes()); + } - let truncated_hash = &full_hash[..6]; + let full_hash = h + .finalize_sha256_array() + .map_err(|e| anyhow!("failed to finalize hash: {:?}", e))?; Ok(format!( "2:{hash}", - hash = base64::prelude::BASE64_URL_SAFE_NO_PAD.encode(truncated_hash) + hash = base64::prelude::BASE64_URL_SAFE_NO_PAD.encode(&full_hash[..6]) )) } diff --git a/wacore/src/send.rs b/wacore/src/send.rs index ad3e6eff7..44b3f1011 100644 --- a/wacore/src/send.rs +++ b/wacore/src/send.rs @@ -33,6 +33,19 @@ pub(crate) mod stanza { pub const ENC_TYPE_SKMSG: &str = "skmsg"; } +/// Extract (enc_type, is_prekey, serialized) from a CiphertextMessage. +pub fn extract_ciphertext(msg: CiphertextMessage) -> Option<(&'static str, bool, Box<[u8]>)> { + match msg { + CiphertextMessage::SignalMessage(m) => { + Some((stanza::ENC_TYPE_MSG, false, m.into_serialized())) + } + CiphertextMessage::PreKeySignalMessage(m) => { + Some((stanza::ENC_TYPE_PKMSG, true, m.into_serialized())) + } + _ => None, + } +} + /// Unwrap wrapper message types to reach the inner message. /// Matches WA Web's getUnwrappedProtobufMessage (EProtoUtils.js:19-35). fn unwrap_message(msg: &wa::Message) -> &wa::Message { @@ -603,16 +616,12 @@ where .await { Ok(encrypted_payload) => { - let (enc_type, serialized_bytes) = match encrypted_payload { - CiphertextMessage::PreKeySignalMessage(msg) => { - includes_prekey_message = true; - (stanza::ENC_TYPE_PKMSG, msg.serialized().to_vec()) - } - CiphertextMessage::SignalMessage(msg) => { - (stanza::ENC_TYPE_MSG, msg.serialized().to_vec()) - } - _ => continue, + let Some((enc_type, is_prekey, serialized_bytes)) = + extract_ciphertext(encrypted_payload) + else { + continue; }; + includes_prekey_message |= is_prekey; let mut enc_builder = NodeBuilder::new("enc") .attr("v", stanza::ENC_VERSION) @@ -708,7 +717,7 @@ pub async fn prepare_dm_stanza< message.clone() }; - let recipient_plaintext = MessageUtils::pad_message_v2(message_for_encryption.encode_to_vec()); + let recipient_plaintext = MessageUtils::encode_and_pad(&message_for_encryption); let dsm = wa::Message { device_sent_message: Some(Box::new(DeviceSentMessage { @@ -816,19 +825,14 @@ where S: crate::libsignal::protocol::SessionStore, I: crate::libsignal::protocol::IdentityKeyStore, { - let plaintext = MessageUtils::pad_message_v2(message.encode_to_vec()); + let plaintext = MessageUtils::encode_and_pad(message); let signal_address = encryption_jid.to_protocol_address(); let encrypted_message = message_encrypt(&plaintext, &signal_address, session_store, identity_store).await?; - let (enc_type, serialized_bytes) = match encrypted_message { - CiphertextMessage::SignalMessage(msg) => (stanza::ENC_TYPE_MSG, msg.serialized().to_vec()), - CiphertextMessage::PreKeySignalMessage(msg) => { - (stanza::ENC_TYPE_PKMSG, msg.serialized().to_vec()) - } - _ => return Err(anyhow!("Unexpected peer encryption message type")), - }; + let (enc_type, _, serialized_bytes) = extract_ciphertext(encrypted_message) + .ok_or_else(|| anyhow!("Unexpected peer encryption message type"))?; let enc_node = NodeBuilder::new("enc") .attrs([("v", "2"), ("type", enc_type)]) @@ -866,25 +870,14 @@ where S: crate::libsignal::protocol::SessionStore, I: crate::libsignal::protocol::IdentityKeyStore, { - let plaintext = MessageUtils::pad_message_v2(message.encode_to_vec()); + let plaintext = MessageUtils::encode_and_pad(message); let signal_address = encryption_jid.to_protocol_address(); let encrypted = message_encrypt(&plaintext, &signal_address, session_store, identity_store).await?; - let (enc_type, is_prekey, serialized) = match encrypted { - CiphertextMessage::SignalMessage(msg) => { - (stanza::ENC_TYPE_MSG, false, msg.serialized().to_vec()) - } - CiphertextMessage::PreKeySignalMessage(msg) => { - (stanza::ENC_TYPE_PKMSG, true, msg.serialized().to_vec()) - } - _ => { - return Err(anyhow!( - "Unexpected encryption message type for group retry" - )); - } - }; + let (enc_type, is_prekey, serialized) = extract_ciphertext(encrypted) + .ok_or_else(|| anyhow!("Unexpected encryption message type for group retry"))?; // count="N" distinguishes retries from normal sends (MsgCreateDeviceStanza.js:150-153) let mut enc_builder = NodeBuilder::new("enc") @@ -1175,7 +1168,7 @@ pub async fn prepare_group_stanza< } } - let plaintext = MessageUtils::pad_message_v2(message_for_encryption.encode_to_vec()); + let plaintext = MessageUtils::encode_and_pad(&message_for_encryption); let skmsg = encrypt_group_message( stores.sender_key_store, &to_jid, @@ -1185,7 +1178,7 @@ pub async fn prepare_group_stanza< ) .await?; - let skmsg_ciphertext = skmsg.serialized().to_vec(); + let skmsg_ciphertext = skmsg.into_serialized(); let mediatype = media_type_from_message(message); let hide_decrypt_fail = (edit.as_ref().is_some_and(|e| { @@ -1279,7 +1272,7 @@ pub async fn create_sender_key_distribution_message_for_group( ) .await?; - Ok(skdm.serialized().to_vec()) + Ok(skdm.into_serialized().into_vec()) } /// Ensure the status stanza has a `` node listing all recipient