Skip to content

perf!: reduce allocations in upload, send, and Jid paths #471

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
Mar 31, 2026
Merged

perf!: reduce allocations in upload, send, and Jid paths #471

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
24 changes: 12 additions & 12 deletions src/features/status.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,19 +73,19 @@ impl<'a> Status<'a> {
/// the `UploadResponse`, JPEG thumbnail bytes, and optional caption.
pub async fn send_image(
&self,
upload: &UploadResponse,
upload: UploadResponse,
thumbnail: Vec<u8>,
caption: Option<&str>,
recipients: &[Jid],
options: StatusSendOptions,
) -> Result<SendResult, anyhow::Error> {
let message = wa::Message {
image_message: Some(Box::new(wa::message::ImageMessage {
url: Some(upload.url.clone()),
direct_path: Some(upload.direct_path.clone()),
media_key: Some(upload.media_key.clone()),
file_sha256: Some(upload.file_sha256.clone()),
file_enc_sha256: Some(upload.file_enc_sha256.clone()),
url: Some(upload.url),
direct_path: Some(upload.direct_path),
media_key: Some(upload.media_key.to_vec()),
file_sha256: Some(upload.file_sha256.to_vec()),
file_enc_sha256: Some(upload.file_enc_sha256.to_vec()),
file_length: Some(upload.file_length),
mimetype: Some("image/jpeg".to_string()),
jpeg_thumbnail: Some(thumbnail),
Expand All @@ -106,7 +106,7 @@ impl<'a> Status<'a> {
/// the `UploadResponse`, JPEG thumbnail bytes, duration in seconds, and optional caption.
pub async fn send_video(
&self,
upload: &UploadResponse,
upload: UploadResponse,
thumbnail: Vec<u8>,
duration_seconds: u32,
caption: Option<&str>,
Expand All @@ -115,11 +115,11 @@ impl<'a> Status<'a> {
) -> Result<SendResult, anyhow::Error> {
let message = wa::Message {
video_message: Some(Box::new(wa::message::VideoMessage {
url: Some(upload.url.clone()),
direct_path: Some(upload.direct_path.clone()),
media_key: Some(upload.media_key.clone()),
file_sha256: Some(upload.file_sha256.clone()),
file_enc_sha256: Some(upload.file_enc_sha256.clone()),
url: Some(upload.url),
direct_path: Some(upload.direct_path),
media_key: Some(upload.media_key.to_vec()),
file_sha256: Some(upload.file_sha256.to_vec()),
file_enc_sha256: Some(upload.file_enc_sha256.to_vec()),
file_length: Some(upload.file_length),
mimetype: Some("video/mp4".to_string()),
jpeg_thumbnail: Some(thumbnail),
Expand Down
44 changes: 25 additions & 19 deletions src/upload.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,9 +156,9 @@ where
return Ok(UploadResponse {
url,
direct_path,
media_key: enc.media_key.to_vec(),
file_enc_sha256: enc.file_enc_sha256.to_vec(),
file_sha256: enc.file_sha256.to_vec(),
media_key: enc.media_key,
file_enc_sha256: enc.file_enc_sha256,
file_sha256: enc.file_sha256,
file_length,
media_key_timestamp,
});
Expand Down Expand Up @@ -196,9 +196,9 @@ where
return Ok(UploadResponse {
url: raw.url,
direct_path: raw.direct_path,
media_key: enc.media_key.to_vec(),
file_enc_sha256: enc.file_enc_sha256.to_vec(),
file_sha256: enc.file_sha256.to_vec(),
media_key: enc.media_key,
file_enc_sha256: enc.file_enc_sha256,
file_sha256: enc.file_sha256,
file_length,
media_key_timestamp,
});
Expand Down Expand Up @@ -233,9 +233,9 @@ where
pub struct UploadResponse {
pub url: String,
pub direct_path: String,
pub media_key: Vec<u8>,
pub file_enc_sha256: Vec<u8>,
pub file_sha256: Vec<u8>,
pub media_key: [u8; 32],
pub file_enc_sha256: [u8; 32],
pub file_sha256: [u8; 32],
pub file_length: u64,
/// Unix timestamp (seconds) when the media key was generated.
pub media_key_timestamp: i64,
Expand All @@ -254,6 +254,19 @@ impl From<UploadResponse> for wacore::sticker_pack::MediaUploadInfo {
}
}

impl UploadResponse {
/// Convert crypto fields to `Vec<u8>` for protobuf message construction.
pub fn media_key_vec(&self) -> Vec<u8> {
self.media_key.to_vec()
}
pub fn file_sha256_vec(&self) -> Vec<u8> {
self.file_sha256.to_vec()
}
pub fn file_enc_sha256_vec(&self) -> Vec<u8> {
self.file_enc_sha256.to_vec()
}
}

#[derive(Deserialize)]
struct RawUploadResponse {
url: String,
Expand All @@ -264,7 +277,7 @@ struct RawUploadResponse {
#[derive(Default, Clone)]
pub struct UploadOptions {
/// Reuse an existing media key instead of generating a fresh one.
pub media_key: Option<Vec<u8>>,
pub media_key: Option<[u8; 32]>,
}

impl std::fmt::Debug for UploadOptions {
Expand All @@ -280,7 +293,7 @@ impl UploadOptions {
Self::default()
}

pub fn with_media_key(mut self, key: Vec<u8>) -> Self {
pub fn with_media_key(mut self, key: [u8; 32]) -> Self {
self.media_key = Some(key);
self
}
Expand All @@ -299,14 +312,7 @@ impl Client {
) -> Result<UploadResponse> {
let file_length = data.len() as u64;
let enc = wacore::runtime::blocking(&*self.runtime, move || {
let key_ref = match &options.media_key {
Some(k) => Some(
<&[u8; 32]>::try_from(k.as_slice())
.map_err(|_| anyhow!("media_key must be exactly 32 bytes"))?,
),
None => None,
};
wacore::upload::encrypt_media_with_key(&data, media_type, key_ref)
wacore::upload::encrypt_media_with_key(&data, media_type, options.media_key.as_ref())
})
.await?;

Expand Down
30 changes: 15 additions & 15 deletions tests/e2e/tests/media.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,9 @@ fn build_image_message(upload: &UploadResponse, caption: Option<&str>) -> wa::Me
image_message: Some(Box::new(wa::message::ImageMessage {
url: Some(upload.url.clone()),
direct_path: Some(upload.direct_path.clone()),
media_key: Some(upload.media_key.clone()),
file_sha256: Some(upload.file_sha256.clone()),
file_enc_sha256: Some(upload.file_enc_sha256.clone()),
media_key: Some(upload.media_key.to_vec()),
file_sha256: Some(upload.file_sha256.to_vec()),
file_enc_sha256: Some(upload.file_enc_sha256.to_vec()),
file_length: Some(upload.file_length),
mimetype: Some("image/jpeg".to_string()),
caption: caption.map(|c| c.to_string()),
Expand All @@ -33,9 +33,9 @@ fn build_video_message(
video_message: Some(Box::new(wa::message::VideoMessage {
url: Some(upload.url.clone()),
direct_path: Some(upload.direct_path.clone()),
media_key: Some(upload.media_key.clone()),
file_sha256: Some(upload.file_sha256.clone()),
file_enc_sha256: Some(upload.file_enc_sha256.clone()),
media_key: Some(upload.media_key.to_vec()),
file_sha256: Some(upload.file_sha256.to_vec()),
file_enc_sha256: Some(upload.file_enc_sha256.to_vec()),
file_length: Some(upload.file_length),
mimetype: Some("video/mp4".to_string()),
seconds: Some(seconds),
Expand All @@ -52,9 +52,9 @@ fn build_document_message(upload: &UploadResponse, filename: &str, mimetype: &st
document_message: Some(Box::new(wa::message::DocumentMessage {
url: Some(upload.url.clone()),
direct_path: Some(upload.direct_path.clone()),
media_key: Some(upload.media_key.clone()),
file_sha256: Some(upload.file_sha256.clone()),
file_enc_sha256: Some(upload.file_enc_sha256.clone()),
media_key: Some(upload.media_key.to_vec()),
file_sha256: Some(upload.file_sha256.to_vec()),
file_enc_sha256: Some(upload.file_enc_sha256.to_vec()),
file_length: Some(upload.file_length),
mimetype: Some(mimetype.to_string()),
file_name: Some(filename.to_string()),
Expand All @@ -70,9 +70,9 @@ fn build_audio_message(upload: &UploadResponse, ptt: bool, seconds: u32) -> wa::
audio_message: Some(Box::new(wa::message::AudioMessage {
url: Some(upload.url.clone()),
direct_path: Some(upload.direct_path.clone()),
media_key: Some(upload.media_key.clone()),
file_sha256: Some(upload.file_sha256.clone()),
file_enc_sha256: Some(upload.file_enc_sha256.clone()),
media_key: Some(upload.media_key.to_vec()),
file_sha256: Some(upload.file_sha256.to_vec()),
file_enc_sha256: Some(upload.file_enc_sha256.to_vec()),
file_length: Some(upload.file_length),
mimetype: Some(if ptt {
"audio/ogg; codecs=opus".to_string()
Expand Down Expand Up @@ -301,9 +301,9 @@ async fn test_upload_then_download_via_downloadable_trait() -> anyhow::Result<()
let img_msg = wa::message::ImageMessage {
url: Some(upload.url.clone()),
direct_path: Some(upload.direct_path.clone()),
media_key: Some(upload.media_key.clone()),
file_sha256: Some(upload.file_sha256.clone()),
file_enc_sha256: Some(upload.file_enc_sha256.clone()),
media_key: Some(upload.media_key.to_vec()),
file_sha256: Some(upload.file_sha256.to_vec()),
file_enc_sha256: Some(upload.file_enc_sha256.to_vec()),
file_length: Some(upload.file_length),
mimetype: Some("image/jpeg".to_string()),
..Default::default()
Expand Down
6 changes: 3 additions & 3 deletions tests/e2e/tests/newsletter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,9 +277,9 @@ async fn test_newsletter_send_media_message() -> anyhow::Result<()> {
image_message: Some(Box::new(wa::message::ImageMessage {
url: Some(upload.url.clone()),
direct_path: Some(upload.direct_path.clone()),
media_key: Some(upload.media_key.clone()),
file_sha256: Some(upload.file_sha256.clone()),
file_enc_sha256: Some(upload.file_enc_sha256.clone()),
media_key: Some(upload.media_key.to_vec()),
file_sha256: Some(upload.file_sha256.to_vec()),
file_enc_sha256: Some(upload.file_enc_sha256.to_vec()),
file_length: Some(upload.file_length),
mimetype: Some("image/jpeg".to_string()),
caption: Some("Newsletter image test".to_string()),
Expand Down
4 changes: 2 additions & 2 deletions wacore/binary/src/jid.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -313,9 +313,9 @@ impl JidExt for Jid {
}

impl Jid {
pub fn new(user: &str, server: &str) -> Self {
pub fn new(user: impl Into<String>, server: &str) -> Self {
Self {
user: user.to_string(),
user: user.into(),
server: cow_server_from_str(server),
..Default::default()
}
Expand Down
20 changes: 8 additions & 12 deletions wacore/src/send.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -709,7 +709,7 @@ pub async fn prepare_dm_stanza<
let dsm = wa::Message {
device_sent_message: Some(Box::new(DeviceSentMessage {
destination_jid: Some(to_jid.to_string()),
message: Some(Box::new(message_for_encryption.clone())),
message: Some(Box::new(message_for_encryption)),
phash: Some("".to_string()),
})),
..Default::default()
Expand Down Expand Up @@ -977,7 +977,7 @@ pub async fn prepare_group_stanza<

let mut message_children: Vec<Node> = Vec::new();
let mut includes_prekey_message = false;
let mut resolved_devices_for_phash: Option<Vec<Jid>> = None;
let mut phash_for_stanza: Option<String> = None;
let mut skdm_encrypted_devices: Vec<Jid> = Vec::new();

// Determine if we need to distribute SKDM and to which devices
Expand Down Expand Up @@ -1107,7 +1107,10 @@ pub async fn prepare_group_stanza<
if let Some(ref distribution_list) = distribution_list {
// WA Web computes phash from the full distribution list (target set at
// send time), not the actual encrypted outcome
resolved_devices_for_phash = Some(distribution_list.clone());
match MessageUtils::participant_list_hash(distribution_list) {
Ok(phash) => phash_for_stanza = Some(phash),
Err(e) => log::warn!("Failed to compute phash for group {}: {:?}", to_jid, e),
}
let axolotl_skdm_bytes = create_sender_key_distribution_message_for_group(
stores.sender_key_store,
&to_jid,
Expand Down Expand Up @@ -1227,15 +1230,8 @@ pub async fn prepare_group_stanza<
}

// Add phash if we distributed keys in this message
if let Some(devices) = &resolved_devices_for_phash {
match MessageUtils::participant_list_hash(devices) {
Ok(phash) => {
stanza_builder = stanza_builder.attr("phash", phash);
}
Err(e) => {
log::warn!("Failed to compute phash for group {}: {:?}", to_jid, e);
}
}
if let Some(phash) = phash_for_stanza {
stanza_builder = stanza_builder.attr("phash", phash);
}

// Add any extra stanza nodes provided by the caller
Expand Down
Loading
Loading