Skip to content

perf: zero-copy ciphertext serialization and DRY send pipeline #502

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jlucaso1 merged 1 commit into from
Apr 7, 2026
Merged

perf: zero-copy ciphertext serialization and DRY send pipeline #502

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 9 additions & 10 deletions src/features/signal.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
20 changes: 20 additions & 0 deletions wacore/libsignal/src/protocol/protocol.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down
48 changes: 27 additions & 21 deletions wacore/src/messages.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,40 +7,46 @@ use waproto::whatsapp as wa;
pub struct MessageUtils;

impl MessageUtils {
pub fn pad_message_v2(mut plaintext: Vec<u8>) -> Vec<u8> {
fn random_pad_len() -> u8 {
use rand::RngExt;
let mut rng = rand::make_rng::<rand::rngs::StdRng>();
let v = rng.random::<u8>() & 0x0F;
if v == 0 { 0x0F } else { v }
}

let mut pad_val = rng.random::<u8>() & 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<u8>) -> Vec<u8> {
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<u8> {
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<String> {
// Hash sorted ad_strings incrementally (avoids join() allocation).
let mut jids: Vec<String> = 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])
))
}

Expand Down
63 changes: 28 additions & 35 deletions wacore/src/send.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
};
Comment on lines +619 to 623

@coderabbitai coderabbitai Bot Apr 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Add diagnostics before skipping unsupported ciphertext variants.

This branch silently drops a device encryption attempt. A warning log here would make partial fanout failures observable during debugging/incident response.

Suggested patch 
-                let Some((enc_type, is_prekey, serialized_bytes)) =
-                    extract_ciphertext(encrypted_payload)
-                else {
-                    continue;
-                };
+                let Some((enc_type, is_prekey, serialized_bytes)) =
+                    extract_ciphertext(encrypted_payload)
+                else {
+                    log::warn!(
+                        "Unsupported ciphertext variant for device {}; skipping encryption output",
+                        signal_address
+                    );
+                    continue;
+                };
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
let Some((enc_type, is_prekey, serialized_bytes)) =
extract_ciphertext(encrypted_payload)
else {
continue;
};
let Some((enc_type, is_prekey, serialized_bytes)) =
extract_ciphertext(encrypted_payload)
else {
log::warn!(
"Unsupported ciphertext variant for device {}; skipping encryption output",
signal_address
);
continue;
};
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@wacore/src/send.rs` around lines 619 - 623, The code silently skips
unsupported ciphertexts when extract_ciphertext(encrypted_payload) returns None;
add a warning log immediately before the continue to surface partial fanout
failures. At the extract_ciphertext call site, emit a warning (using the
existing logging facility, e.g. warn! or log::warn!) that includes identifying
context such as a short debug/hex of encrypted_payload and the device/recipient
identifiers available in the surrounding scope, then continue; this makes
unsupported ciphertext variants observable without changing control flow.

includes_prekey_message |= is_prekey;

let mut enc_builder = NodeBuilder::new("enc")
.attr("v", stanza::ENC_VERSION)
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)])
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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,
Expand All @@ -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| {
Expand Down Expand Up @@ -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 `<participants>` node listing all recipient
Expand Down
Loading