Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
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
6 changes: 1 addition & 5 deletions src/history_sync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -164,11 +164,7 @@ impl HistorySecretSeedCollector {
Err(_) => return,
};
for sender in senders.into_iter().flatten() {
let sender_id = if sender.is_same_chat_as(chat) {
Arc::clone(&chat_id)
} else {
Arc::from(sender.to_non_ad_string())
};
let sender_id = MsgSecretEntry::sender_id_for(chat, &chat_id, &sender);
self.entries.push(MsgSecretEntry {
chat: Arc::clone(&chat_id),
sender: sender_id,
Expand Down
47 changes: 28 additions & 19 deletions src/message/msg_secret.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,34 +31,43 @@ impl Client {
let class = wacore::msg_secret::classify(msg, chat_is_bot);
let message_ts = u64::try_from(info.timestamp.timestamp()).ok();

// Build both aliases (primary, plus the bot-DM LID key) and write them
// in one batch so a partial write can't leave only one stored.
let mut entries = Vec::with_capacity(2);
if let Some(entry) = self.build_msg_secret_entry(
let primary = self.build_msg_secret_entry(
&info.source.chat,
&info.source.sender,
&info.id,
secret_bytes,
class,
message_ts,
) {
entries.push(entry);
}
);
// The bot-DM LID key is the only alias a capture ever adds, so a plain
// chat writes exactly one row. Deciding that before touching a Vec keeps
// the common capture off the batch allocation entirely.
let mut bot_alias = None;
if chat_is_bot
&& let Some(sender) = self.dm_sender_identity_for(&info.source.chat).await
&& sender.to_non_ad() != info.source.sender.to_non_ad()
&& let Some(entry) = self.build_msg_secret_entry(
{
bot_alias = self.build_msg_secret_entry(
&info.source.chat,
&sender,
&info.id,
secret_bytes,
class,
message_ts,
)
{
entries.push(entry);
);
}

match (primary, bot_alias) {
// Both aliases go out in one batch so a partial write can't leave
// only one of them stored.
(Some(primary), Some(alias)) => {
self.persist_msg_secret_entries(vec![primary, alias]).await
}
(Some(entry), None) | (None, Some(entry)) => {
self.msg_secret_buffer.queue_one(entry).await
}
(None, None) => {}
}
self.persist_msg_secret_entries(entries).await;
}

/// Build one retention entry, applying the policy gates and computing the
Expand Down Expand Up @@ -92,14 +101,14 @@ impl Client {
message_ts,
wacore::time::now_secs(),
);
Some(wacore::store::traits::MsgSecretEntry {
chat: chat.to_non_ad_string().into(),
sender: sender.to_non_ad_string().into(),
msg_id: msg_id.into(),
secret: *secret,
Some(wacore::store::traits::MsgSecretEntry::new(
chat,
sender,
msg_id,
*secret,
expires_at,
message_ts: message_ts.and_then(|t| i64::try_from(t).ok()).unwrap_or(0),
})
message_ts.and_then(|t| i64::try_from(t).ok()).unwrap_or(0),
))
}

/// Queue a batch of secret aliases on the write-behind buffer: immediately
Expand Down
102 changes: 94 additions & 8 deletions src/message/receive.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,21 @@ impl ParsedSessionMessage {
}
}

/// Append to one of the three per-kind enc buckets, allocating that bucket only
/// once it actually receives a payload.
///
/// A stanza's enc nodes are overwhelmingly all one kind, so reserving all three
/// buckets up front spent two allocations per message on buffers that stayed
/// empty for their whole lifetime. Reserving the stanza's full enc count on the
/// first push keeps a mixed stanza at one allocation per non-empty bucket,
/// exactly as before.
fn push_enc_payload(bucket: &mut Vec<EncPayload>, stanza_enc_count: usize, payload: EncPayload) {
if bucket.capacity() == 0 {
bucket.reserve_exact(stanza_enc_count);
}
bucket.push(payload);
}

async fn decrypt_session_message(
message: &mut ParsedSessionMessage,
signal_address: &wacore::libsignal::protocol::ProtocolAddress,
Expand Down Expand Up @@ -222,9 +237,9 @@ impl Client {
return None;
}

let mut session_payloads = Vec::with_capacity(all_enc_nodes.len());
let mut group_payloads = Vec::with_capacity(all_enc_nodes.len());
let mut bot_payloads = Vec::with_capacity(all_enc_nodes.len());
let mut session_payloads = Vec::new();
let mut group_payloads = Vec::new();
let mut bot_payloads = Vec::new();
let mut max_sender_retry_count = 0;
let mut has_hide_fail = false;
let mut had_unknown_enc = false;
Expand Down Expand Up @@ -301,13 +316,14 @@ impl Client {
}
};

if payload.enc_type.is_bot_secret() {
bot_payloads.push(payload);
let bucket = if payload.enc_type.is_bot_secret() {
&mut bot_payloads
} else if payload.enc_type.is_session() {
session_payloads.push(payload);
&mut session_payloads
} else {
group_payloads.push(payload);
}
&mut group_payloads
};
push_enc_payload(bucket, all_enc_nodes.len(), payload);
}

// WA Web diagnostic: validate skmsg is not first in multi-enc messages.
Expand Down Expand Up @@ -1750,6 +1766,76 @@ impl Client {
}
}

#[cfg(test)]
mod enc_bucket_tests {
use super::push_enc_payload;
use crate::message::EncPayload;
use wacore::message_processing::EncType;

fn payload(enc_type: EncType) -> EncPayload {
EncPayload {
ciphertext: bytes::Bytes::from_static(b"ct"),
enc_type,
padding_version: 2,
}
}

/// The DM shape: one enc node, one bucket used, two buckets untouched. The
/// used bucket must be sized to the stanza exactly (not to `Vec`'s default
/// growth step) and the unused ones must own no buffer at all.
#[test]
fn a_one_enc_stanza_sizes_one_bucket_exactly_and_leaves_the_rest_empty() {
let mut used: Vec<EncPayload> = Vec::new();
let unused: Vec<EncPayload> = Vec::new();

push_enc_payload(&mut used, 1, payload(EncType::Message));

assert_eq!(used.len(), 1);
assert_eq!(
used.capacity(),
1,
"the bucket must be reserved to the stanza's enc count, not grown"
);
assert_eq!(unused.capacity(), 0, "an empty bucket must own no buffer");
assert!(unused.is_empty());
}

/// The first push must reserve the stanza's whole enc count, so a multi-enc
/// stanza still pays exactly one allocation for the bucket it fills. Eight
/// is past `Vec`'s own first growth step, so plain pushes would reallocate.
#[test]
fn the_first_push_reserves_the_whole_stanza() {
const ENC_COUNT: usize = 8;
let mut bucket: Vec<EncPayload> = Vec::new();
push_enc_payload(&mut bucket, ENC_COUNT, payload(EncType::Message));
let reserved = bucket.as_ptr();
assert_eq!(bucket.capacity(), ENC_COUNT);

for _ in 1..ENC_COUNT {
push_enc_payload(&mut bucket, ENC_COUNT, payload(EncType::PreKeyMessage));
}
assert_eq!(bucket.len(), ENC_COUNT);
assert_eq!(
bucket.as_ptr(),
reserved,
"filling up to the stanza's enc count must not reallocate"
);
assert_eq!(bucket[0].enc_type, EncType::Message);
assert_eq!(bucket[ENC_COUNT - 1].enc_type, EncType::PreKeyMessage);
}

/// A degenerate count must not make the helper skip its reservation and
/// leave the bucket re-reserving on every later push.
#[test]
fn a_zero_count_still_stores_the_payload() {
let mut bucket: Vec<EncPayload> = Vec::new();
push_enc_payload(&mut bucket, 0, payload(EncType::SenderKey));
assert_eq!(bucket.len(), 1);
push_enc_payload(&mut bucket, 0, payload(EncType::SenderKey));
assert_eq!(bucket.len(), 2);
}
}

#[cfg(test)]
mod tests {
use crate::test_utils::create_test_client_with_failing_http;
Expand Down
25 changes: 25 additions & 0 deletions src/msg_secret_buffer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -616,6 +616,31 @@ mod tests {
buf.wait_flushed().await;
}

/// Building a row costs one allocation per *distinct* identifier and
/// nothing else. The naive spelling (`to_non_ad_string().into()` per JID)
/// cost five for a direct message; this is the per-message saving on both
/// the inbound capture and the outbound persist.
#[test]
fn entry_construction_allocates_once_per_distinct_identifier() {
let chat: wacore_binary::Jid = "5511987650001@s.whatsapp.net".parse().unwrap();
let peer_device: wacore_binary::Jid = "5511987650001:33@s.whatsapp.net".parse().unwrap();
let me: wacore_binary::Jid = "5511987650002@s.whatsapp.net".parse().unwrap();
let secret = [0u8; wacore::reporting_token::MESSAGE_SECRET_SIZE];

// A DM: chat and sender are the same user, so the row holds two
// allocations (the shared identifier and the message id).
let dm = crate::test_alloc::min_allocs(2, || {
MsgSecretEntry::new(&chat, &peer_device, "3EB0AABBCCDDEEFF0011", secret, 0, 0)
});
assert_eq!(dm, 2, "a direct-message row: shared identifier + msg id");

// Outbound: the sender is us, a different user, so it needs its own.
let outbound = crate::test_alloc::min_allocs(3, || {
MsgSecretEntry::new(&chat, &me, "3EB0AABBCCDDEEFF0011", secret, 0, 0)
});
assert_eq!(outbound, 3, "distinct users cost one identifier each");
}

/// Cloning a buffered entry must stay allocation-free: identifiers share
/// their Arc allocations and the protocol-sized secret lives inline.
#[test]
Expand Down
11 changes: 3 additions & 8 deletions src/send/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2432,14 +2432,9 @@ impl Client {
u64::try_from(now).ok(),
now,
);
let entry = wacore::store::traits::MsgSecretEntry {
chat: chat.to_non_ad_string().into(),
sender: sender.to_non_ad_string().into(),
msg_id: msg_id.into(),
secret: *secret,
expires_at,
message_ts: now,
};
let entry = wacore::store::traits::MsgSecretEntry::new(
chat, sender, msg_id, *secret, expires_at, now,
);
// Same write-behind buffer as inbound captures: visible immediately,
// flushed off the send path (msmsg replies read buffer-first).
self.msg_secret_buffer.queue_one(entry).await;
Expand Down
56 changes: 56 additions & 0 deletions wacore/binary/src/jid.rs
Original file line number Diff line number Diff line change
Expand Up @@ -659,6 +659,20 @@ impl Jid {
buf
}

/// [`Self::to_non_ad_string`] as a shareable `Arc<str>`, in exactly one
/// allocation. Going through the `String` first costs two — the buffer, then
/// the `Arc<str>` its bytes are copied into — and the message-secret rows
/// build two of these per message.
pub fn to_non_ad_arc_str(&self) -> std::sync::Arc<str> {
let mut writer = JidStackWriter::new();
if write_jid_fallible(&mut writer, &self.user, self.server, 0, 0).is_ok() {
return std::sync::Arc::from(writer.as_str());
}
// A user part too long for the stack buffer (never seen on the wire)
// still renders, just back through the heap.
std::sync::Arc::from(self.to_non_ad_string())
}

/// Check if this JID matches the user or their LID.
/// Useful for checking if a participant is "us" in group messages.
#[inline]
Expand Down Expand Up @@ -1976,6 +1990,48 @@ mod tests {
}
}

/// The stack-buffered `Arc<str>` form must render exactly what the `String`
/// form does, including for inputs that overflow the stack buffer and fall
/// back to the heap, and for multibyte user parts (the buffer is bounded in
/// bytes, and a split fragment would be invalid UTF-8).
#[test]
fn to_non_ad_arc_str_matches_to_non_ad_string() {
let long_user = "9".repeat(80);
let multibyte_user = "ẞünïcodé-ñ".repeat(3);
let owned = [
format!("{long_user}:12@s.whatsapp.net"),
format!("{multibyte_user}@g.us"),
format!("{multibyte_user}@s.whatsapp.net"),
];
let cases = [
"1234567890:33@s.whatsapp.net",
"1234567890@s.whatsapp.net",
"100000012345678:25@lid",
"867051314767696:0@bot",
"120363021033254949@g.us",
"status@broadcast",
]
.into_iter()
.chain(owned.iter().map(String::as_str));

for s in cases {
let jid: Jid = s.parse().unwrap_or_else(|e| panic!("parse {s}: {e}"));
assert_eq!(
&*jid.to_non_ad_arc_str(),
jid.to_non_ad_string().as_str(),
"mismatch for {s}"
);
}

// A default (empty-user) JID has no wire form to render but must still
// agree with the String path rather than panic in the stack writer.
let empty = Jid::default();
assert_eq!(
&*empty.to_non_ad_arc_str(),
empty.to_non_ad_string().as_str()
);
}

Comment thread
coderabbitai[bot] marked this conversation as resolved.
#[test]
fn test_into_non_ad_matches_to_non_ad() {
// into_non_ad (consuming) must produce a JID identical to to_non_ad (cloning).
Expand Down
Loading
Loading