diff --git a/.github/workflows/bench-integration.yml b/.github/workflows/bench-integration.yml new file mode 100644 index 000000000..c51263ae8 --- /dev/null +++ b/.github/workflows/bench-integration.yml @@ -0,0 +1,186 @@ +name: Integration Benchmark + +on: + pull_request: + branches: [main] + push: + branches: [main] + +concurrency: + group: bench-integration-${{ github.head_ref || github.run_id }} + cancel-in-progress: true + +env: + CARGO_TERM_COLOR: always + PROTOC_VERSION: "3.25.3" + SCCACHE_GHA_ENABLED: "true" + RUSTC_WRAPPER: "sccache" + +jobs: + bench-integration-pr: + if: github.event_name == 'pull_request' + name: Integration Benchmark (PR) + runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: write + services: + mock-server: + image: ghcr.io/whiskeysockets-devtools/bartender:latest + credentials: + username: ${{ github.actor }} + password: ${{ secrets.BARTENDER_GHCR_TOKEN }} + ports: + - 8080:8080 + env: + CHATSTATE_TTL_SECS: "3" + options: --log-driver none + steps: + - uses: actions/checkout@v6 + + - uses: dtolnay/rust-toolchain@master + with: + toolchain: nightly-2026-04-05 + + - name: Install protoc + uses: taiki-e/install-action@v2 + with: + tool: protoc@${{ env.PROTOC_VERSION }} + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Cache Rust registry + uses: Swatinem/rust-cache@v2 + with: + prefix-key: ${{ runner.os }}-cargo-bench-int + cache-targets: "false" + + - name: Wait for mock server + run: | + for i in $(seq 1 30); do + if curl -sk https://localhost:8080/ > /dev/null 2>&1; then + echo "Mock server is ready" + exit 0 + fi + sleep 1 + done + echo "Mock server failed to become ready" + exit 1 + + - name: Run integration benchmarks + env: + MOCK_SERVER_URL: "wss://127.0.0.1:8080/ws/chat" + RUST_LOG: warn + run: | + cargo run -p bench-integration --release \ + > bench_integration_results.json \ + 2> bench_integration.log + cat bench_integration.log >&2 + echo "Results:" + cat bench_integration_results.json + + - name: Upload benchmark artifacts + uses: actions/upload-artifact@v4 + with: + name: bench-integration-results + path: | + bench_integration_results.json + bench_integration.log + retention-days: 30 + + - name: Compare against baseline + uses: benchmark-action/github-action-benchmark@v1 + with: + name: "whatsapp-rust integration benchmarks" + tool: "customSmallerIsBetter" + output-file-path: bench_integration_results.json + github-token: ${{ secrets.GITHUB_TOKEN }} + auto-push: false + save-data-file: false + benchmark-data-dir-path: dev/bench-integration + summary-always: true + + bench-integration-push: + if: github.event_name == 'push' + name: Integration Benchmark (push) + runs-on: ubuntu-latest + permissions: + contents: write + deployments: write + services: + mock-server: + image: ghcr.io/whiskeysockets-devtools/bartender:latest + credentials: + username: ${{ github.actor }} + password: ${{ secrets.BARTENDER_GHCR_TOKEN }} + ports: + - 8080:8080 + env: + CHATSTATE_TTL_SECS: "3" + options: --log-driver none + steps: + - uses: actions/checkout@v6 + + - uses: dtolnay/rust-toolchain@master + with: + toolchain: nightly-2026-04-05 + + - name: Install protoc + uses: taiki-e/install-action@v2 + with: + tool: protoc@${{ env.PROTOC_VERSION }} + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.9 + + - name: Cache Rust registry + uses: Swatinem/rust-cache@v2 + with: + prefix-key: ${{ runner.os }}-cargo-bench-int + cache-targets: "false" + + - name: Wait for mock server + run: | + for i in $(seq 1 30); do + if curl -sk https://localhost:8080/ > /dev/null 2>&1; then + echo "Mock server is ready" + exit 0 + fi + sleep 1 + done + echo "Mock server failed to become ready" + exit 1 + + - name: Run integration benchmarks + env: + MOCK_SERVER_URL: "wss://127.0.0.1:8080/ws/chat" + RUST_LOG: warn + run: | + cargo run -p bench-integration --release \ + > bench_integration_results.json \ + 2> bench_integration.log + cat bench_integration.log >&2 + echo "Results:" + cat bench_integration_results.json + + - name: Upload benchmark artifacts + uses: actions/upload-artifact@v4 + with: + name: bench-integration-results + path: | + bench_integration_results.json + bench_integration.log + retention-days: 30 + + - name: Store baseline + uses: benchmark-action/github-action-benchmark@v1 + with: + name: "whatsapp-rust integration benchmarks" + tool: "customSmallerIsBetter" + output-file-path: bench_integration_results.json + github-token: ${{ secrets.GITHUB_TOKEN }} + auto-push: true + benchmark-data-dir-path: dev/bench-integration + max-items-in-chart: 100 + summary-always: true diff --git a/.gitignore b/.gitignore index 539e2d082..a157f7df1 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ docs .claude __pycache__ .codex +dhat-heap.json diff --git a/Cargo.lock b/Cargo.lock index 498b50e9c..23e7d5b7d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -144,6 +144,22 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "bench-integration" +version = "0.0.0" +dependencies = [ + "anyhow", + "dhat", + "e2e-tests", + "env_logger", + "log", + "serde", + "serde_json", + "tokio", + "wacore", + "whatsapp-rust", +] + [[package]] name = "bincode" version = "1.3.3" @@ -608,7 +624,6 @@ dependencies = [ "wacore", "wacore-binary", "whatsapp-rust", - "whatsapp-rust-sqlite-storage", "whatsapp-rust-tokio-transport", "whatsapp-rust-ureq-http-client", ] @@ -2405,6 +2420,7 @@ version = "0.5.0" dependencies = [ "async-trait", "bincode 2.0.1", + "bytes", "diesel", "diesel_migrations", "libsqlite3-sys", diff --git a/Cargo.toml b/Cargo.toml index edb16d2c7..59507d659 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ members = [ ".", "http_clients/ureq-client", "storages/sqlite-storage", + "tests/bench-integration", "tests/e2e", "transports/tokio-transport", "wacore", @@ -161,3 +162,9 @@ inherits = "release" lto = "thin" debug = 1 strip = false + +# Profiling: optimized with debug symbols for DHAT / heaptrack / perf +[profile.profiling] +inherits = "release" +debug = 1 +strip = false diff --git a/src/appstate_sync.rs b/src/appstate_sync.rs index 06556e479..0426e94e1 100644 --- a/src/appstate_sync.rs +++ b/src/appstate_sync.rs @@ -59,7 +59,7 @@ mod tests { async fn store_prekey(&self, _: u32, _: &[u8], _: bool) -> StoreResult<()> { Ok(()) } - async fn load_prekey(&self, _: u32) -> StoreResult>> { + async fn load_prekey(&self, _: u32) -> StoreResult> { Ok(None) } async fn remove_prekey(&self, _: u32) -> StoreResult<()> { diff --git a/src/client.rs b/src/client.rs index 5f399754a..5c0dec966 100644 --- a/src/client.rs +++ b/src/client.rs @@ -893,9 +893,8 @@ impl Client { /// [`send_node`](Client::send_node) for normal stanza sending. pub async fn send_raw_bytes(&self, plaintext: Vec) -> Result<(), ClientError> { let noise_socket = self.get_noise_socket().await?; - let encrypted_buf = Vec::with_capacity(plaintext.len() + 32); noise_socket - .encrypt_and_send(plaintext, encrypted_buf) + .encrypt_and_send(bytes::Bytes::from(plaintext)) .await?; self.last_data_sent_ms .store(wacore::time::now_millis().max(0) as u64, Ordering::Relaxed); diff --git a/src/features/signal.rs b/src/features/signal.rs index 2ed11995f..9e4b1e706 100644 --- a/src/features/signal.rs +++ b/src/features/signal.rs @@ -151,7 +151,7 @@ impl<'a> Signal<'a> { wacore::send::create_sender_key_distribution_message_for_group( &mut adapter.sender_key_store, group_jid, - &own_jid, + &sender_addr, ) .await?, ) @@ -162,7 +162,7 @@ impl<'a> Signal<'a> { let ciphertext = wacore::send::encrypt_group_message( &mut adapter.sender_key_store, group_jid, - &own_jid, + &sender_addr, plaintext, &mut rng, ) diff --git a/src/handlers/message.rs b/src/handlers/message.rs index cb584c85d..5d3854559 100644 --- a/src/handlers/message.rs +++ b/src/handlers/message.rs @@ -38,51 +38,11 @@ impl StanzaHandler for MessageHandler { } }; - // Single cache lookup: get or create the lane (lock + queue + worker). + // Single-flight: get_with_by_ref guarantees exactly one init runs per key, + // preventing duplicate workers for the same chat (TOCTOU race). let lane = client .chat_lanes - .get_with_by_ref(&chat_jid, async { - let (tx, rx) = async_channel::unbounded::>(); - - let client_for_worker = client.clone(); - let spawn_generation = client - .connection_generation - .load(std::sync::atomic::Ordering::Acquire); - - client - .runtime - .spawn(Box::pin(async move { - while let Ok(msg_node) = rx.recv().await { - if client_for_worker - .connection_generation - .load(std::sync::atomic::Ordering::Acquire) - != spawn_generation - { - log::debug!(target: "MessageQueue", "Stale worker exiting; remaining messages will be redelivered by server"); - break; - } - let start = wacore::time::now_millis() as u64; - let client = client_for_worker.clone(); - Box::pin(client.handle_incoming_message(msg_node)).await; - let elapsed = - (wacore::time::now_millis() as u64).saturating_sub(start); - if elapsed > MAX_MESSAGE_DELAY_MS { - warn!( - target: "MessageQueue", - "Message processing took {:.1}s (MAX_MESSAGE_DELAY is {}s)", - elapsed as f64 / 1000.0, - MAX_MESSAGE_DELAY_MS / 1000 - ); - } - } - })) - .detach(); - - ChatLane { - enqueue_lock: Arc::new(async_lock::Mutex::new(())), - queue_tx: tx, - } - }) + .get_with_by_ref(&chat_jid, async { create_chat_lane(&client) }) .await; // Lock serializes enqueue order for this chat @@ -97,3 +57,47 @@ impl StanzaHandler for MessageHandler { true } } + +/// Construct a ChatLane with a spawned worker task. Extracted to keep the +/// init closure passed to `get_with_by_ref` small. +fn create_chat_lane(client: &Arc) -> ChatLane { + let (tx, rx) = async_channel::unbounded::>(); + + let client_for_worker = client.clone(); + let spawn_generation = client + .connection_generation + .load(std::sync::atomic::Ordering::Acquire); + + client + .runtime + .spawn(Box::pin(async move { + while let Ok(msg_node) = rx.recv().await { + if client_for_worker + .connection_generation + .load(std::sync::atomic::Ordering::Acquire) + != spawn_generation + { + log::debug!(target: "MessageQueue", "Stale worker exiting; remaining messages will be redelivered by server"); + break; + } + let start = wacore::time::now_millis() as u64; + let client = client_for_worker.clone(); + Box::pin(client.handle_incoming_message(msg_node)).await; + let elapsed = (wacore::time::now_millis() as u64).saturating_sub(start); + if elapsed > MAX_MESSAGE_DELAY_MS { + warn!( + target: "MessageQueue", + "Message processing took {:.1}s (MAX_MESSAGE_DELAY is {}s)", + elapsed as f64 / 1000.0, + MAX_MESSAGE_DELAY_MS / 1000 + ); + } + } + })) + .detach(); + + ChatLane { + enqueue_lock: Arc::new(async_lock::Mutex::new(())), + queue_tx: tx, + } +} diff --git a/src/handshake.rs b/src/handshake.rs index 779e2e963..f3f4ee1d2 100644 --- a/src/handshake.rs +++ b/src/handshake.rs @@ -71,7 +71,7 @@ pub async fn do_handshake( // First message includes the WA connection header (with optional edge routing) let framed = wacore::framing::encode_frame(&client_hello_bytes, Some(&header)) .map_err(HandshakeError::Transport)?; - transport.send(framed).await?; + transport.send(bytes::Bytes::from(framed)).await?; // Wait for server response frame let resp_frame = loop { @@ -113,7 +113,7 @@ pub async fn do_handshake( // Subsequent messages don't need the header let framed = wacore::framing::encode_frame(&client_finish_bytes, None) .map_err(HandshakeError::Transport)?; - transport.send(framed).await?; + transport.send(bytes::Bytes::from(framed)).await?; let (write_key, read_key) = handshake_state.finish()?; info!("Handshake complete, switching to encrypted communication"); diff --git a/src/message.rs b/src/message.rs index 6eea83e18..4e4657a6f 100644 --- a/src/message.rs +++ b/src/message.rs @@ -26,6 +26,48 @@ use waproto::whatsapp::{self as wa}; /// After this many retries, we stop sending retry receipts and rely solely on PDO. const MAX_DECRYPT_RETRIES: u8 = 5; +/// Pre-extracted enc node payload. Holds owned copies of the fields needed for +/// decryption so the async decrypt phase doesn't borrow the original NodeRef tree. +pub(crate) struct EncPayload { + pub ciphertext: bytes::Bytes, + pub enc_type: EncType, + pub padding_version: u8, +} + +impl EncPayload { + fn from_parts(ciphertext: bytes::Bytes, enc_node: &NodeRef<'_>) -> 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; + Some(Self { + ciphertext, + enc_type, + padding_version, + }) + } + + /// Zero-copy extraction from an OwnedNodeRef. + pub(crate) fn from_owned_node(owner: &OwnedNodeRef, enc_node: &NodeRef<'_>) -> Option { + Self::from_parts(owner.slice_bytes(enc_node.content_bytes()?), enc_node) + } + + /// Copying extraction from a NodeRef (used in tests where there's no OwnedNodeRef). + #[cfg(test)] + pub(crate) fn from_node_ref(node: &NodeRef<'_>) -> Option { + Self::from_parts(bytes::Bytes::copy_from_slice(node.content_bytes()?), node) + } +} + +/// Parsed and classified message ready for decryption. All data is owned -- +/// the original node tree is no longer borrowed. +pub(crate) struct ClassifiedMessage { + pub info: Arc, + pub sender_encryption_jid: Jid, + pub session_payloads: Vec, + pub group_payloads: Vec, + pub max_sender_retry_count: u8, + pub decrypt_fail_mode: crate::types::events::DecryptFailMode, +} + /// Retry count threshold for logging high retry warnings. /// WhatsApp Web logs metrics when retry count exceeds this value. const HIGH_RETRY_COUNT_THRESHOLD: u8 = 3; @@ -262,6 +304,22 @@ impl Client { } pub(crate) async fn handle_incoming_message(self: Arc, node: Arc) { + // Phase 1: classify borrows the node tree, extracts owned payloads, returns quickly. + // Phase 2: process_classified_message holds no node borrows across heavy .await points, + // keeping the async state machine small. + let classified = match self.classify_incoming_message(&node).await { + Some(c) => c, + None => return, + }; + // node is no longer borrowed here -- drop it before the heavy phase + drop(node); + self.process_classified_message(classified).await; + } + + async fn classify_incoming_message( + self: &Arc, + node: &OwnedNodeRef, + ) -> Option { let nr = node.get(); let info = match self.parse_message_info(nr).await { Ok(info) => Arc::new(info), @@ -269,14 +327,14 @@ impl Client { let id = nr.get_attr("id").map(|v| v.as_str()); let from = nr.get_attr("from").map(|v| v.as_str()); log::warn!("Failed to parse message info (id={id:?}, from={from:?}): {e:?}"); - return; + return None; } }; // Newsletters use instead of <enc> because they are not E2E encrypted. if info.source.chat.is_newsletter() { self.handle_newsletter_message(nr, &info).await; - return; + return None; } // Warm LID-PN cache before resolution so resolve_encryption_jid() finds the mapping @@ -313,7 +371,7 @@ impl Client { info.id, nr.tag ); - return; + return None; } if let Some(unavailable) = unavailable_node { @@ -337,11 +395,11 @@ impl Client { decrypt_fail_mode: crate::types::events::DecryptFailMode::Show, }, )); - return; + return None; } - let mut session_enc_nodes = Vec::with_capacity(all_enc_nodes.len()); - let mut group_content_enc_nodes = Vec::with_capacity(all_enc_nodes.len()); + let mut session_payloads = Vec::with_capacity(all_enc_nodes.len()); + let mut group_payloads = Vec::with_capacity(all_enc_nodes.len()); let mut max_sender_retry_count: u8 = 0; let mut has_hide_fail = false; @@ -402,19 +460,27 @@ impl Client { continue; } - // Fall back to built-in handlers - match EncType::from_wire(enc_type.as_ref()) { - Some(et) if et.is_session() => session_enc_nodes.push(*enc_node), - Some(EncType::SenderKey) => group_content_enc_nodes.push(*enc_node), - _ => log::warn!("Unknown enc type: {enc_type}"), + // Zero-copy: slice_bytes returns a Bytes sub-view into the + // node's backing buffer without memcpy. + // from_owned_node returns None for unknown enc types or missing content. + let payload = match EncPayload::from_owned_node(node, enc_node) { + Some(p) => p, + None => { + log::warn!("Enc node has no content or unknown type: {enc_type}"); + continue; + } + }; + + if payload.enc_type.is_session() { + session_payloads.push(payload); + } else { + group_payloads.push(payload); } } // WA Web diagnostic: validate skmsg is not first in multi-enc messages. - // If skmsg comes first, the SKDM (carried in pkmsg/msg) hasn't been processed yet, - // so the skmsg decryption would fail with NoSenderKey. - if !session_enc_nodes.is_empty() - && !group_content_enc_nodes.is_empty() + if !session_payloads.is_empty() + && !group_payloads.is_empty() && all_enc_nodes.first().is_some_and(|n| { n.get_attr("type") .map(|v| v.as_str()) @@ -429,16 +495,31 @@ impl Client { ); } - // Determine decrypt fail mode from enc nodes (WA Web: hideFail) - let decrypt_fail_mode = if has_hide_fail { - crate::types::events::DecryptFailMode::Hide - } else { - crate::types::events::DecryptFailMode::Show - }; + Some(ClassifiedMessage { + info, + sender_encryption_jid, + session_payloads, + group_payloads, + max_sender_retry_count, + decrypt_fail_mode: if has_hide_fail { + crate::types::events::DecryptFailMode::Hide + } else { + crate::types::events::DecryptFailMode::Show + }, + }) + } + + /// Phase 2: acquire permit, decrypt payloads, flush. No node borrows. + async fn process_classified_message(self: Arc<Self>, msg: ClassifiedMessage) { + let ClassifiedMessage { + info, + sender_encryption_jid, + session_payloads, + group_payloads, + max_sender_retry_count, + decrypt_fail_mode, + } = msg; - // Pre-seed retry cache with sender's retry count to avoid redundant retries. - // Uses max(existing, incoming) so redeliveries with higher counts update the cache, - // but lower counts don't reset our local counter. if max_sender_retry_count > 0 { let cache_key = self .make_retry_cache_key(&info.source.chat, &info.id, &info.source.sender) @@ -486,7 +567,7 @@ impl Client { log::debug!( "Starting PASS 1: Processing {} session establishment messages (pkmsg/msg)", - session_enc_nodes.len() + session_payloads.len() ); // Skip session processing for group/broadcast JIDs — they use sender keys, not 1:1 sessions. @@ -498,20 +579,20 @@ impl Client { session_decrypted_successfully, session_had_duplicates, session_dispatched_undecryptable, - ) = if !is_group_sender && !session_enc_nodes.is_empty() { + ) = if !is_group_sender && !session_payloads.is_empty() { self.clone() .process_session_enc_batch( - &session_enc_nodes, + &session_payloads, &info, &sender_encryption_jid, decrypt_fail_mode, ) .await } else { - if is_group_sender && !session_enc_nodes.is_empty() { + if is_group_sender && !session_payloads.is_empty() { log::debug!( "Skipping {} session messages from group sender {}", - session_enc_nodes.len(), + session_payloads.len(), sender_encryption_jid ); } @@ -520,7 +601,7 @@ impl Client { log::debug!( "Starting PASS 2: Processing {} group content messages (skmsg)", - group_content_enc_nodes.len() + group_payloads.len() ); // Only process group content if: @@ -532,8 +613,8 @@ impl Client { // the SKDM it carried is lost, so skmsg will always fail with NoSenderKey — skip it // to avoid unnecessary retry receipts. The retry for the pkmsg will cause the sender // to resend the entire message including SKDM. - if !group_content_enc_nodes.is_empty() { - let should_process_skmsg = session_enc_nodes.is_empty() + if !group_payloads.is_empty() { + let should_process_skmsg = session_payloads.is_empty() || session_decrypted_successfully || session_had_duplicates; @@ -541,7 +622,7 @@ impl Client { match self .clone() .process_group_enc_batch( - &group_content_enc_nodes, + &group_payloads, &info, &sender_encryption_jid, decrypt_fail_mode, @@ -590,7 +671,7 @@ impl Client { } } else if !session_decrypted_successfully && !session_had_duplicates - && !session_enc_nodes.is_empty() + && !session_payloads.is_empty() { // Edge case: message with only msg/pkmsg that failed to decrypt, no skmsg warn!( @@ -609,16 +690,15 @@ impl Client { .await; } - async fn process_session_enc_batch<'n>( + async fn process_session_enc_batch( self: Arc<Self>, - enc_nodes: &[&'n NodeRef<'n>], + payloads: &[EncPayload], info: &Arc<MessageInfo>, sender_encryption_jid: &Jid, decrypt_fail_mode: crate::types::events::DecryptFailMode, ) -> (bool, bool, bool) { - // Returns (any_success, any_duplicate, dispatched_undecryptable) use wacore::libsignal::protocol::CiphertextMessage; - if enc_nodes.is_empty() { + if payloads.is_empty() { return (false, false, false); } @@ -637,25 +717,13 @@ impl Client { let mut any_duplicate = false; let mut dispatched_undecryptable = false; - for enc_node in enc_nodes { - let ciphertext: &[u8] = match enc_node.content_bytes() { - Some(b) => b, - None => { - log::warn!("Enc node has no byte content (batch session)"); - continue; - } - }; - let enc_type = match enc_node.attrs().optional_string("type") { - Some(t) => t, - None => { - log::warn!("Enc node missing 'type' attribute (batch session)"); - continue; - } - }; - let padding_version = enc_node.attrs().optional_u64("v").unwrap_or(2) as u8; - let enc_type_enum = EncType::from_wire(enc_type.as_ref()); + for payload in payloads { + let ciphertext = &payload.ciphertext[..]; + let enc_type = payload.enc_type; + let enc_type_str = enc_type.as_wire_str(); + let padding_version = payload.padding_version; - let parsed_message = if enc_type_enum == Some(EncType::PreKeyMessage) { + let parsed_message = if enc_type == EncType::PreKeyMessage { match PreKeySignalMessage::try_from(ciphertext) { Ok(m) => CiphertextMessage::PreKeySignalMessage(m), Err(e) => { @@ -673,7 +741,7 @@ impl Client { } }; - if enc_type_enum == Some(EncType::PreKeyMessage) { + if enc_type == EncType::PreKeyMessage { // FLAGGED FOR DEBUGGING: "Bad Mac" Reproducibility #[cfg(feature = "debug-snapshots")] { @@ -682,7 +750,7 @@ impl Client { "id": info.id, "sender_jid": sender_encryption_jid.to_string(), "timestamp": info.timestamp, - "enc_type": enc_type, + "enc_type": enc_type_str, "payload_base64": BASE64_STANDARD.encode(ciphertext), }); @@ -702,6 +770,9 @@ impl Client { } } + // Shadow with wire string for all downstream usage (logging, handlers) + let enc_type = enc_type_str; + let decrypt_res = message_decrypt( &parsed_message, &signal_address, @@ -720,7 +791,7 @@ impl Client { if let Err(e) = self .clone() .handle_decrypted_plaintext( - &enc_type, + enc_type, &padded_plaintext, padding_version, info, @@ -806,7 +877,7 @@ impl Client { if let Err(e) = self .clone() .handle_decrypted_plaintext( - &enc_type, + enc_type, &padded_plaintext, padding_version, info, @@ -844,7 +915,7 @@ impl Client { &parsed_message, &mut adapter, &mut rng, - &enc_type, + enc_type, padding_version, info, ) @@ -906,7 +977,7 @@ impl Client { &parsed_message, &mut adapter, &mut rng, - &enc_type, + enc_type, padding_version, info, ) @@ -972,7 +1043,7 @@ impl Client { &parsed_message, &mut adapter, &mut rng, - &enc_type, + enc_type, padding_version, info, ) @@ -1015,29 +1086,21 @@ impl Client { (any_success, any_duplicate, dispatched_undecryptable) } - async fn process_group_enc_batch<'n>( + async fn process_group_enc_batch( self: Arc<Self>, - enc_nodes: &[&'n NodeRef<'n>], + payloads: &[EncPayload], info: &Arc<MessageInfo>, _sender_encryption_jid: &Jid, decrypt_fail_mode: crate::types::events::DecryptFailMode, ) -> Result<(), DecryptionError> { - if enc_nodes.is_empty() { + if payloads.is_empty() { return Ok(()); } - // Use the signal cache adapter for group decryption so sender keys are read/written - // through the cache, keeping it consistent with SKDM processing. let mut adapter = self.signal_adapter().await; - for enc_node in enc_nodes { - let ciphertext: &[u8] = match enc_node.content_bytes() { - Some(b) => b, - None => { - log::warn!("Enc node has no byte content (batch group)"); - continue; - } - }; - let padding_version = enc_node.attrs().optional_u64("v").unwrap_or(2) as u8; + for payload in payloads { + let ciphertext = &payload.ciphertext[..]; + let padding_version = payload.padding_version; // Always use bare sender for sender key operations. Real WA delivers // skmsg with bare participant but pkmsg (SKDM) with device-qualified @@ -1725,12 +1788,12 @@ mod tests { .bytes(signal_message.serialized().to_vec()) .build(); let enc_node_ref = enc_node.as_node_ref(); - let enc_nodes = vec![&enc_node_ref]; + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; // With SessionNotFound, should return (false, false, true) - no success, no dupe, dispatched event let (success, had_duplicates, dispatched) = client .process_session_enc_batch( - &enc_nodes, + &payloads, &info, &sender_jid, crate::types::events::DecryptFailMode::Show, @@ -1814,12 +1877,12 @@ mod tests { .bytes(signal_message.serialized().to_vec()) .build(); let enc_node_ref = enc_node.as_node_ref(); - let enc_nodes = vec![&enc_node_ref]; + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; let (success, had_duplicates, dispatched) = client .clone() .process_session_enc_batch( - &enc_nodes, + &payloads, &info, &sender_jid, crate::types::events::DecryptFailMode::Show, @@ -2765,13 +2828,13 @@ mod tests { .build(); let enc_node_ref = enc_node.as_node_ref(); - let enc_nodes = vec![&enc_node_ref]; + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; // Call process_session_enc_batch // This should handle any errors gracefully without panicking let (success, _had_duplicates, _dispatched) = client .process_session_enc_batch( - &enc_nodes, + &payloads, &info, &sender_jid, crate::types::events::DecryptFailMode::Show, @@ -2851,14 +2914,16 @@ mod tests { log::info!("Test: Created batch of 2 messages with invalid data"); - let enc_node_refs_owned: Vec<_> = enc_nodes.iter().map(|n| n.as_node_ref()).collect(); - let enc_node_refs: Vec<&NodeRef<'_>> = enc_node_refs_owned.iter().collect(); + let payloads: Vec<EncPayload> = enc_nodes + .iter() + .filter_map(|n| EncPayload::from_node_ref(&n.as_node_ref())) + .collect(); // Process the batch // Should handle all errors gracefully without stopping at first error let (success, _had_duplicates, _dispatched) = client .process_session_enc_batch( - &enc_node_refs, + &payloads, &info, &sender_jid, crate::types::events::DecryptFailMode::Show, @@ -2924,13 +2989,13 @@ mod tests { .build(); let enc_node_ref = enc_node.as_node_ref(); - let enc_nodes = vec![&enc_node_ref]; + let payloads: Vec<EncPayload> = vec![EncPayload::from_node_ref(&enc_node_ref).unwrap()]; // Process the message // Should handle errors gracefully in group context let (success, _had_duplicates, _dispatched) = client .process_session_enc_batch( - &enc_nodes, + &payloads, &info, &sender_phone, crate::types::events::DecryptFailMode::Show, diff --git a/src/prekeys.rs b/src/prekeys.rs index 110bded44..5c5480e01 100644 --- a/src/prekeys.rs +++ b/src/prekeys.rs @@ -143,12 +143,25 @@ impl Client { key_pairs_to_upload.push((pre_key_id, key_pair)); } - // Encode once — reused for both pre-upload store and post-upload mark. - let encoded_batch: Vec<(u32, Vec<u8>)> = { + // Encode all prekey records into a single contiguous buffer, then slice + // into Bytes sub-views. This replaces 812 individual encode_to_vec() allocs + // with one large allocation + zero-copy slicing. + let encoded_batch: Vec<(u32, bytes::Bytes)> = { use prost::Message; - keys_to_upload - .iter() - .map(|(id, record)| (*id, record.encode_to_vec())) + let total_len: usize = keys_to_upload.iter().map(|(_, r)| r.encoded_len()).sum(); + let mut buf = Vec::with_capacity(total_len); + let mut offsets = Vec::with_capacity(keys_to_upload.len()); + for (id, record) in &keys_to_upload { + let start = buf.len(); + record + .encode(&mut buf) + .expect("prost encode into pre-sized Vec"); + offsets.push((*id, start..buf.len())); + } + let shared = bytes::Bytes::from(buf); + offsets + .into_iter() + .map(|(id, range)| (id, shared.slice(range))) .collect() }; @@ -339,7 +352,7 @@ impl Client { // Build a lookup so we preserve the server-requested order. // Dedupe the expected count since the server may send duplicate IDs. - let loaded_map: std::collections::HashMap<u32, Vec<u8>> = loaded.into_iter().collect(); + let loaded_map: std::collections::HashMap<u32, bytes::Bytes> = loaded.into_iter().collect(); let unique_requested: std::collections::HashSet<&u32> = response.prekey_ids.iter().collect(); @@ -351,37 +364,25 @@ impl Client { return Ok(()); } + // Extract public keys directly from stored protobuf bytes without full decode let mut prekey_pubkeys = Vec::with_capacity(response.prekey_ids.len()); for prekey_id in &response.prekey_ids { let Some(record_bytes) = loaded_map.get(prekey_id) else { log::warn!("digestKey: missing local prekey {}, skipping", prekey_id); return Ok(()); }; - use prost::Message; - match waproto::whatsapp::PreKeyRecordStructure::decode(record_bytes.as_slice()) { - Ok(record) => { - if let Some(pk) = record.public_key { - prekey_pubkeys.push(pk); - } else { - log::warn!( - "digestKey: prekey {} has no public key, skipping", - prekey_id - ); - return Ok(()); - } - } - Err(e) => { + match wacore::prekeys::extract_prekey_public_key(record_bytes) { + Some(pk) => prekey_pubkeys.push(pk), + None => { log::warn!( - "digestKey: failed to decode prekey {}: {}, skipping", - prekey_id, - e + "digestKey: prekey {} has no public key, skipping", + prekey_id ); return Ok(()); } } } - // Compute local SHA-1 digest matching WA Web's validateLocalKeyBundle let local_hash = wacore::prekeys::compute_key_bundle_digest( identity_bytes, skey_pub_bytes, diff --git a/src/retry.rs b/src/retry.rs index 35e158561..6f871433b 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -2546,7 +2546,8 @@ mod tests { fn resolve_retry_chat_info_preserves_original_from() { use wacore_binary::builder::NodeBuilder; - // DM with device suffix — original_from should preserve it + // DM with device suffix — original_from preserves the raw receipt from + // (WA Web: variable m = e.from, used as-is for stanza to) let node = NodeBuilder::new("receipt").build(); let receipt = make_test_receipt("5511999999999:33@s.whatsapp.net"); diff --git a/src/send.rs b/src/send.rs index de7800f4c..b886d1124 100644 --- a/src/send.rs +++ b/src/send.rs @@ -883,7 +883,7 @@ impl Client { &mut store_adapter.session_store, &mut store_adapter.identity_store, to, - encryption_jid, + &signal_addr, message, request_id, ) diff --git a/src/socket/noise_socket.rs b/src/socket/noise_socket.rs index 44292b220..48ca7f35d 100644 --- a/src/socket/noise_socket.rs +++ b/src/socket/noise_socket.rs @@ -15,8 +15,7 @@ type SendResult = std::result::Result<(), EncryptSendError>; /// A job sent to the dedicated sender task. struct SendJob { - plaintext_buf: Vec<u8>, - out_buf: Vec<u8>, + plaintext: bytes::Bytes, response_tx: oneshot::Sender<SendResult>, } @@ -79,6 +78,9 @@ impl NoiseSocket { send_job_rx: async_channel::Receiver<SendJob>, ) { let mut write_counter: u32 = 0; + // Reusable buffers -- capacity stays allocated between sends + let mut enc_buf = Vec::with_capacity(4096); + let mut out_buf = Vec::with_capacity(4096); while let Ok(job) = send_job_rx.recv().await { let result = Self::process_send_job( @@ -86,16 +88,14 @@ impl NoiseSocket { &transport, &write_key, &mut write_counter, - job.plaintext_buf, - job.out_buf, + &job.plaintext, + &mut enc_buf, + &mut out_buf, ) .await; - // Send result back to caller. Ignore error if receiver was dropped. let _ = job.response_tx.send(result); } - - // Channel closed - NoiseSocket was dropped, task exits naturally } /// Process a single send job: encrypt and send the message. @@ -104,39 +104,34 @@ impl NoiseSocket { transport: &Arc<dyn Transport>, write_key: &Arc<NoiseCipher>, write_counter: &mut u32, - mut plaintext_buf: Vec<u8>, - mut out_buf: Vec<u8>, + plaintext: &[u8], + enc_buf: &mut Vec<u8>, + out_buf: &mut Vec<u8>, ) -> SendResult { let counter = *write_counter; - // For small messages, encrypt plaintext_buf in-place then frame into out_buf. - // This avoids the previous triple-copy pattern (plaintext→out→plaintext→out). - if plaintext_buf.len() <= INLINE_ENCRYPT_THRESHOLD { - if let Err(e) = write_key.encrypt_in_place_with_counter(counter, &mut plaintext_buf) { + if plaintext.len() <= INLINE_ENCRYPT_THRESHOLD { + // Copy into reusable enc_buf, encrypt in place + enc_buf.clear(); + enc_buf.extend_from_slice(plaintext); + if let Err(e) = write_key.encrypt_in_place_with_counter(counter, enc_buf) { return Err(EncryptSendError::crypto(anyhow::anyhow!(e.to_string()))); } - // Frame the ciphertext from plaintext_buf into out_buf (single copy) out_buf.clear(); - if let Err(e) = wacore::framing::encode_frame_into(&plaintext_buf, None, &mut out_buf) { + if let Err(e) = wacore::framing::encode_frame_into(enc_buf, None, out_buf) { return Err(EncryptSendError::framing(e)); } } else { - // Offload larger messages to a blocking thread + // Large messages: encrypt on blocking thread (reads plaintext, returns new ciphertext) let write_key = write_key.clone(); - - let plaintext_arc = Arc::new(plaintext_buf); - let plaintext_arc_for_task = plaintext_arc.clone(); + let plaintext_owned = plaintext.to_vec(); let encrypt_result = wacore::runtime::blocking(&**runtime, move || { - write_key.encrypt_with_counter(counter, &plaintext_arc_for_task[..]) + write_key.encrypt_with_counter(counter, &plaintext_owned) }) .await; - // Recover ownership so the buffer is dropped at end of scope - plaintext_buf = Arc::try_unwrap(plaintext_arc).unwrap_or_else(|arc| (*arc).clone()); - drop(plaintext_buf); - let ciphertext = match encrypt_result { Ok(c) => c, Err(e) => { @@ -145,28 +140,28 @@ impl NoiseSocket { }; out_buf.clear(); - if let Err(e) = wacore::framing::encode_frame_into(&ciphertext, None, &mut out_buf) { + if let Err(e) = wacore::framing::encode_frame_into(&ciphertext, None, out_buf) { return Err(EncryptSendError::framing(e)); } } - if let Err(e) = transport.send(out_buf).await { + // copy_from_slice so out_buf retains its capacity for the next send + let frame = bytes::Bytes::copy_from_slice(out_buf); + out_buf.clear(); + if let Err(e) = transport.send(frame).await { return Err(EncryptSendError::transport(e)); } - // Only advance the counter after the encrypted frame was successfully sent. - // If transport.send() fails, we can retry with the same counter value. *write_counter = write_counter.wrapping_add(1); Ok(()) } - pub async fn encrypt_and_send(&self, plaintext_buf: Vec<u8>, out_buf: Vec<u8>) -> SendResult { + pub async fn encrypt_and_send(&self, plaintext: bytes::Bytes) -> SendResult { let (response_tx, response_rx) = oneshot::channel(); let job = SendJob { - plaintext_buf, - out_buf, + plaintext, response_tx, }; @@ -216,10 +211,7 @@ mod tests { read_key, ); - let plaintext_buf = Vec::with_capacity(1024); - let encrypted_buf = Vec::with_capacity(1024); - - let result = socket.encrypt_and_send(plaintext_buf, encrypted_buf).await; + let result = socket.encrypt_and_send(bytes::Bytes::new()).await; assert!(result.is_ok(), "encrypt_and_send should succeed"); } @@ -240,8 +232,9 @@ mod tests { #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl crate::transport::Transport for RecordingTransport { - async fn send(&self, mut data: Vec<u8>) -> std::result::Result<(), anyhow::Error> { + async fn send(&self, data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> { if data.len() > 16 { + let mut data = data.to_vec(); // Strip the 3-byte frame header, then decrypt in place data.drain(..3); let counter = self @@ -291,8 +284,7 @@ mod tests { // Use index as the first byte of plaintext to identify this send let mut plaintext = vec![i as u8]; plaintext.extend_from_slice(&[0u8; 99]); - let out_buf = Vec::with_capacity(256); - socket.encrypt_and_send(plaintext, out_buf).await + socket.encrypt_and_send(bytes::Bytes::from(plaintext)).await })); } @@ -324,7 +316,7 @@ mod tests { #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl crate::transport::Transport for SizeRecordingTransport { - async fn send(&self, data: Vec<u8>) -> std::result::Result<(), anyhow::Error> { + async fn send(&self, data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> { self.last_size.store(data.len(), Ordering::SeqCst); Ok(()) } @@ -352,12 +344,8 @@ mod tests { for size in test_sizes { let plaintext = vec![0xABu8; size]; - // This is the formula used in client.rs - let buffer_capacity = plaintext.len() + 32; - let encrypted_buf = Vec::with_capacity(buffer_capacity); - let result = socket - .encrypt_and_send(plaintext.clone(), encrypted_buf) + .encrypt_and_send(bytes::Bytes::from(plaintext.clone())) .await; assert!( @@ -376,15 +364,6 @@ mod tests { "Encrypted size for {} byte payload should be {} (got {})", size, expected_max, actual_encrypted_size ); - - // Verify our buffer sizing formula provides enough capacity - assert!( - buffer_capacity >= actual_encrypted_size, - "Buffer capacity {} should be >= encrypted size {} for payload size {}", - buffer_capacity, - actual_encrypted_size, - size - ); } } @@ -399,7 +378,7 @@ mod tests { #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl crate::transport::Transport for NoOpTransport { - async fn send(&self, _data: Vec<u8>) -> std::result::Result<(), anyhow::Error> { + async fn send(&self, _data: bytes::Bytes) -> std::result::Result<(), anyhow::Error> { Ok(()) } async fn disconnect(&self) {} @@ -418,26 +397,20 @@ mod tests { ); // Test empty payload - let result = socket - .encrypt_and_send(vec![], Vec::with_capacity(32)) - .await; + let result = socket.encrypt_and_send(bytes::Bytes::new()).await; assert!(result.is_ok(), "Empty payload should encrypt successfully"); // Test payload at inline threshold boundary (16KB) - let at_threshold = vec![0u8; 16 * 1024]; - let result = socket - .encrypt_and_send(at_threshold, Vec::with_capacity(16 * 1024 + 32)) - .await; + let at_threshold = bytes::Bytes::from(vec![0u8; 16 * 1024]); + let result = socket.encrypt_and_send(at_threshold).await; assert!( result.is_ok(), "Payload at inline threshold should encrypt successfully" ); // Test payload just above inline threshold - let above_threshold = vec![0u8; 16 * 1024 + 1]; - let result = socket - .encrypt_and_send(above_threshold, Vec::with_capacity(16 * 1024 + 33)) - .await; + let above_threshold = bytes::Bytes::from(vec![0u8; 16 * 1024 + 1]); + let result = socket.encrypt_and_send(above_threshold).await; assert!( result.is_ok(), "Payload above inline threshold should encrypt successfully" diff --git a/src/store/signal.rs b/src/store/signal.rs index 95dce8a2a..6ac558e85 100644 --- a/src/store/signal.rs +++ b/src/store/signal.rs @@ -276,7 +276,7 @@ impl PreKeyStore for Device { match self.backend.load_prekey(prekey_id).await { Ok(Some(bytes)) => { // Try new format first (protobuf-encoded PreKeyRecordStructure) - if let Ok(record) = PreKeyRecordStructure::decode(bytes.as_slice()) { + if let Ok(record) = PreKeyRecordStructure::decode(bytes.as_ref()) { return Ok(Some(record)); } diff --git a/src/transport.rs b/src/transport.rs index 69ffde74a..91959d913 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -21,7 +21,7 @@ pub mod mock { #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl Transport for MockTransport { - async fn send(&self, _data: Vec<u8>) -> Result<(), anyhow::Error> { + async fn send(&self, _data: bytes::Bytes) -> Result<(), anyhow::Error> { Ok(()) } diff --git a/storages/sqlite-storage/Cargo.toml b/storages/sqlite-storage/Cargo.toml index 59ba57ca6..838c44399 100644 --- a/storages/sqlite-storage/Cargo.toml +++ b/storages/sqlite-storage/Cargo.toml @@ -14,6 +14,7 @@ bundled-sqlite = ["libsqlite3-sys/bundled"] [dependencies] async-trait = { workspace = true } bincode = { version = "2.0.1", features = ["serde"] } +bytes = { workspace = true } diesel = { version = "2.3.7", default-features = false, features = [ "sqlite", "r2d2", diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index d79bfae9e..66e31a141 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -1,5 +1,6 @@ use crate::schema::*; use async_trait::async_trait; +use bytes::Bytes; use diesel::prelude::*; use diesel::r2d2::{ConnectionManager, Pool}; use diesel::result::{DatabaseErrorKind, Error as DieselError}; @@ -1279,7 +1280,7 @@ impl SignalStore for SqliteStore { )) } - async fn store_prekeys_batch(&self, keys: &[(u32, Vec<u8>)], uploaded: bool) -> Result<()> { + async fn store_prekeys_batch(&self, keys: &[(u32, Bytes)], uploaded: bool) -> Result<()> { if keys.is_empty() { return Ok(()); } @@ -1287,7 +1288,7 @@ impl SignalStore for SqliteStore { let pool = self.pool.clone(); let db_semaphore = self.db_semaphore.clone(); let device_id = self.device_id; - let keys = keys.to_vec(); + let keys: Vec<(u32, Bytes)> = keys.to_vec(); const MAX_RETRIES: u32 = 5; @@ -1311,13 +1312,16 @@ impl SignalStore for SqliteStore { diesel::insert_into(prekeys::table) .values(( prekeys::id.eq(*id as i32), - prekeys::key.eq(record), + prekeys::key.eq(record.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), prekeys::uploaded.eq(uploaded))) + .set(( + prekeys::key.eq(record.as_ref()), + prekeys::uploaded.eq(uploaded), + )) .execute(conn)?; } Ok::<(), diesel::result::Error>(()) @@ -1346,10 +1350,10 @@ impl SignalStore for SqliteStore { )) } - async fn load_prekey(&self, id: u32) -> Result<Option<Vec<u8>>> { + async fn load_prekey(&self, id: u32) -> Result<Option<Bytes>> { let pool = self.pool.clone(); let device_id = self.device_id; - tokio::task::spawn_blocking(move || -> Result<Option<Vec<u8>>> { + tokio::task::spawn_blocking(move || -> Result<Option<Bytes>> { let mut conn = pool .get() .map_err(|e| StoreError::Connection(e.to_string()))?; @@ -1360,20 +1364,20 @@ impl SignalStore for SqliteStore { .first(&mut conn) .optional() .map_err(|e| StoreError::Database(e.to_string()))?; - Ok(res) + Ok(res.map(Bytes::from)) }) .await .map_err(|e| StoreError::Database(e.to_string()))? } - async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Vec<u8>)>> { + async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Bytes)>> { if ids.is_empty() { return Ok(Vec::new()); } let pool = self.pool.clone(); let device_id = self.device_id; let ids: Vec<i32> = ids.iter().map(|&id| id as i32).collect(); - self.with_semaphore(move || -> Result<Vec<(u32, Vec<u8>)>> { + self.with_semaphore(move || -> Result<Vec<(u32, Bytes)>> { let mut conn = pool .get() .map_err(|e| StoreError::Connection(e.to_string()))?; @@ -1383,7 +1387,10 @@ impl SignalStore for SqliteStore { .filter(prekeys::device_id.eq(device_id)) .load(&mut conn) .map_err(|e| StoreError::Database(e.to_string()))?; - Ok(rows.into_iter().map(|(id, key)| (id as u32, key)).collect()) + Ok(rows + .into_iter() + .map(|(id, key)| (id as u32, Bytes::from(key))) + .collect()) }) .await } diff --git a/tests/bench-integration/Cargo.toml b/tests/bench-integration/Cargo.toml new file mode 100644 index 000000000..30dae0baa --- /dev/null +++ b/tests/bench-integration/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "bench-integration" +version = "0.0.0" +edition = "2024" +publish = false + +[package.metadata.cargo-shear] +ignored = ["dhat"] + +[[bin]] +name = "bench-integration" +path = "src/main.rs" + +[features] +dhat-heap = ["dep:dhat"] + +[dependencies] +anyhow = { workspace = true } +dhat = { version = "0.3", optional = true } +e2e-tests = { path = "../e2e" } +env_logger = { workspace = true } +log = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true, features = ["std"] } +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "sync", "time"] } +wacore = { path = "../../wacore" } +whatsapp-rust = { path = "../..", default-features = false, features = [ + "danger-skip-tls-verify", + "debug-diagnostics", + "moka-cache", + "simd", + "tokio-runtime", + "tokio-native", + "signal", +] } diff --git a/tests/bench-integration/src/counting_alloc.rs b/tests/bench-integration/src/counting_alloc.rs new file mode 100644 index 000000000..74d683e27 --- /dev/null +++ b/tests/bench-integration/src/counting_alloc.rs @@ -0,0 +1,63 @@ +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicU64, Ordering}; + +static ALLOC_COUNT: AtomicU64 = AtomicU64::new(0); +static ALLOC_BYTES: AtomicU64 = AtomicU64::new(0); + +pub struct CountingAlloc; + +unsafe impl GlobalAlloc for CountingAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + ALLOC_COUNT.fetch_add(1, Ordering::Relaxed); + ALLOC_BYTES.fetch_add(layout.size() as u64, Ordering::Relaxed); + unsafe { System.alloc(layout) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + // Track the size difference but not as a new allocation + if new_size > layout.size() { + ALLOC_BYTES.fetch_add((new_size - layout.size()) as u64, Ordering::Relaxed); + } + unsafe { System.realloc(ptr, layout, new_size) } + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + ALLOC_COUNT.fetch_add(1, Ordering::Relaxed); + ALLOC_BYTES.fetch_add(layout.size() as u64, Ordering::Relaxed); + unsafe { System.alloc_zeroed(layout) } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct AllocSnapshot { + alloc_count: u64, + alloc_bytes: u64, +} + +impl AllocSnapshot { + pub fn now() -> Self { + Self { + alloc_count: ALLOC_COUNT.load(Ordering::Relaxed), + alloc_bytes: ALLOC_BYTES.load(Ordering::Relaxed), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct AllocDelta { + pub alloc_count: u64, + pub alloc_bytes: u64, +} + +impl AllocDelta { + pub fn between(before: AllocSnapshot, after: AllocSnapshot) -> Self { + Self { + alloc_count: after.alloc_count.saturating_sub(before.alloc_count), + alloc_bytes: after.alloc_bytes.saturating_sub(before.alloc_bytes), + } + } +} diff --git a/tests/bench-integration/src/main.rs b/tests/bench-integration/src/main.rs new file mode 100644 index 000000000..71481f39f --- /dev/null +++ b/tests/bench-integration/src/main.rs @@ -0,0 +1,352 @@ +#[cfg(not(feature = "dhat-heap"))] +mod counting_alloc; + +#[cfg(feature = "dhat-heap")] +#[global_allocator] +static GLOBAL: dhat::Alloc = dhat::Alloc; + +#[cfg(not(feature = "dhat-heap"))] +#[global_allocator] +static GLOBAL: counting_alloc::CountingAlloc = counting_alloc::CountingAlloc; + +use std::time::Instant; + +use e2e_tests::{TestClient, text_msg}; +use serde::Serialize; + +// --------------------------------------------------------------------------- +// Measurement helpers +// --------------------------------------------------------------------------- + +/// Snapshot-based measurement: runs `f`, returns wall time + alloc delta. +#[cfg(not(feature = "dhat-heap"))] +async fn measure<F, T>( + mut f: F, +) -> anyhow::Result<(T, std::time::Duration, counting_alloc::AllocDelta)> +where + F: AsyncFnMut() -> anyhow::Result<T>, +{ + let before = counting_alloc::AllocSnapshot::now(); + let t0 = Instant::now(); + let result = f().await?; + let elapsed = t0.elapsed(); + let delta = counting_alloc::AllocDelta::between(before, counting_alloc::AllocSnapshot::now()); + Ok((result, elapsed, delta)) +} + +/// DHAT mode: just runs `f` and returns wall time (DHAT captures everything). +#[cfg(feature = "dhat-heap")] +async fn measure<F, T>(mut f: F) -> anyhow::Result<(T, std::time::Duration)> +where + F: AsyncFnMut() -> anyhow::Result<T>, +{ + let t0 = Instant::now(); + let result = f().await?; + let elapsed = t0.elapsed(); + Ok((result, elapsed)) +} + +// --------------------------------------------------------------------------- +// Results collection +// --------------------------------------------------------------------------- + +#[derive(Serialize)] +struct BenchEntry { + name: String, + unit: String, + value: u64, +} + +struct BenchResults(Vec<BenchEntry>); + +impl BenchResults { + fn new() -> Self { + Self(Vec::new()) + } + + fn record(&mut self, scenario: &str, metric: &str, unit: &str, value: u64) { + self.0.push(BenchEntry { + name: format!("integration::{scenario}::{metric}"), + unit: unit.to_string(), + value, + }); + } + + fn record_wall(&mut self, scenario: &str, elapsed: std::time::Duration) { + self.record( + scenario, + "wall_ms", + "milliseconds", + elapsed.as_millis() as u64, + ); + } + + #[cfg(not(feature = "dhat-heap"))] + fn record_measured( + &mut self, + scenario: &str, + elapsed: std::time::Duration, + delta: &counting_alloc::AllocDelta, + ) { + self.record(scenario, "alloc_count", "allocations", delta.alloc_count); + self.record(scenario, "alloc_bytes", "bytes", delta.alloc_bytes); + self.record_wall(scenario, elapsed); + } + + #[cfg(not(feature = "dhat-heap"))] + fn record_measured_amortized( + &mut self, + scenario: &str, + n: u64, + elapsed: std::time::Duration, + delta: &counting_alloc::AllocDelta, + ) { + self.record( + scenario, + "alloc_count", + "allocations", + delta.alloc_count / n, + ); + self.record(scenario, "alloc_bytes", "bytes", delta.alloc_bytes / n); + self.record( + scenario, + "wall_ms", + "milliseconds", + elapsed.as_millis() as u64 / n, + ); + } +} + +// --------------------------------------------------------------------------- +// Scenarios +// --------------------------------------------------------------------------- + +/// Measure allocations from client creation through Connected (ready). +async fn bench_connect_to_ready(results: &mut BenchResults) -> anyhow::Result<()> { + let m = measure(async || TestClient::connect("bench_connect").await).await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (client, elapsed, delta) = m; + results.record_measured("connect_to_ready", elapsed, &delta); + client.disconnect().await; + } + #[cfg(feature = "dhat-heap")] + { + let (client, elapsed) = m; + results.record_wall("connect_to_ready", elapsed); + client.disconnect().await; + } + Ok(()) +} + +/// Measure allocations for sending a single DM (sender side only). +/// +/// Both clients are connected before measurement starts. +/// We track the `send_message` call which covers: protobuf encoding, +/// Signal encrypt, node marshal, and WebSocket write. +async fn bench_send_message(results: &mut BenchResults) -> anyhow::Result<()> { + let client_a = TestClient::connect("bench_send_a").await?; + let mut client_b = TestClient::connect("bench_send_b").await?; + let jid_b = client_b.jid().await; + + // Warm up: establish Signal session with a throwaway message + client_a + .client + .send_message(jid_b.clone(), text_msg("warmup-send")) + .await?; + client_b.wait_for_text("warmup-send", 30).await?; + + // -- Single send -- + let m = measure(async || { + client_a + .client + .send_message(jid_b.clone(), text_msg("bench-send-single")) + .await + }) + .await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (_, elapsed, delta) = m; + results.record_measured("send_message", elapsed, &delta); + } + #[cfg(feature = "dhat-heap")] + { + let (_, elapsed) = m; + results.record_wall("send_message", elapsed); + } + + client_b.wait_for_text("bench-send-single", 30).await?; + + // -- Amortized: send N messages -- + const N: u64 = 20; + let send_texts: Vec<String> = (0..N).map(|i| format!("bench-send-{i}")).collect(); + let m = measure(async || { + for text in &send_texts { + client_a + .client + .send_message(jid_b.clone(), text_msg(text)) + .await?; + } + Ok(()) + }) + .await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (_, elapsed, delta) = m; + results.record_measured_amortized("send_message_x20_amortized", N, elapsed, &delta); + } + #[cfg(feature = "dhat-heap")] + { + let (_, elapsed) = m; + results.record( + "send_message_x20_amortized", + "wall_ms", + "milliseconds", + elapsed.as_millis() as u64 / N, + ); + } + + for text in &send_texts { + client_b.wait_for_text(text, 30).await?; + } + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +/// Measure allocations for a full send+receive round-trip. +/// +/// Both clients are connected and session is warm before measurement. +async fn bench_receive_message(results: &mut BenchResults) -> anyhow::Result<()> { + let client_a = TestClient::connect("bench_recv_a").await?; + let mut client_b = TestClient::connect("bench_recv_b").await?; + let jid_b = client_b.jid().await; + + // Warm up Signal session + client_a + .client + .send_message(jid_b.clone(), text_msg("warmup-recv")) + .await?; + client_b.wait_for_text("warmup-recv", 30).await?; + + // -- Single send+receive -- + let m = measure(async || { + client_a + .client + .send_message(jid_b.clone(), text_msg("bench-recv-single")) + .await?; + client_b.wait_for_text("bench-recv-single", 30).await?; + Ok(()) + }) + .await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (_, elapsed, delta) = m; + results.record_measured("send_and_receive_message", elapsed, &delta); + } + #[cfg(feature = "dhat-heap")] + { + let (_, elapsed) = m; + results.record_wall("send_and_receive_message", elapsed); + } + + // -- Amortized N round-trips -- + const N: u64 = 20; + let m = measure(async || { + for i in 0..N { + let text = format!("bench-recv-{i}"); + client_a + .client + .send_message(jid_b.clone(), text_msg(&text)) + .await?; + client_b.wait_for_text(&text, 30).await?; + } + Ok(()) + }) + .await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (_, elapsed, delta) = m; + results.record_measured_amortized("send_and_receive_x20_amortized", N, elapsed, &delta); + } + #[cfg(feature = "dhat-heap")] + { + let (_, elapsed) = m; + results.record( + "send_and_receive_x20_amortized", + "wall_ms", + "milliseconds", + elapsed.as_millis() as u64 / N, + ); + } + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +/// Measure allocations for a reconnect cycle (disconnect -> reconnect -> ready). +async fn bench_reconnect(results: &mut BenchResults) -> anyhow::Result<()> { + let mut client = TestClient::connect("bench_reconn").await?; + + let m = measure(async || { + client.reconnect_and_wait().await?; + Ok(()) + }) + .await?; + + #[cfg(not(feature = "dhat-heap"))] + { + let (_, elapsed, delta) = m; + results.record_measured("reconnect", elapsed, &delta); + } + #[cfg(feature = "dhat-heap")] + { + let (_, elapsed) = m; + results.record_wall("reconnect", elapsed); + } + + client.disconnect().await; + Ok(()) +} + +// --------------------------------------------------------------------------- +// Main +// --------------------------------------------------------------------------- + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("warn")) + .try_init() + .ok(); + + #[cfg(feature = "dhat-heap")] + let _profiler = dhat::Profiler::new_heap(); + + let mut results = BenchResults::new(); + + eprintln!("--- bench: connect_to_ready ---"); + bench_connect_to_ready(&mut results).await?; + + eprintln!("--- bench: send_message ---"); + bench_send_message(&mut results).await?; + + eprintln!("--- bench: receive_message ---"); + bench_receive_message(&mut results).await?; + + eprintln!("--- bench: reconnect ---"); + bench_reconnect(&mut results).await?; + + // Output customSmallerIsBetter JSON to stdout + let json = serde_json::to_string_pretty(&results.0)?; + println!("{json}"); + + eprintln!("--- done: {} metrics collected ---", results.0.len()); + Ok(()) +} diff --git a/tests/e2e/Cargo.toml b/tests/e2e/Cargo.toml index f97e54366..b5d71c233 100644 --- a/tests/e2e/Cargo.toml +++ b/tests/e2e/Cargo.toml @@ -19,8 +19,15 @@ tokio = { workspace = true, features = ["macros", "rt-multi-thread", "sync", "ti uuid = { workspace = true, features = ["v4"] } wacore = { path = "../../wacore" } wacore-binary = { path = "../../wacore/binary" } -whatsapp-rust = { path = "../..", features = ["danger-skip-tls-verify", "debug-diagnostics"] } -whatsapp-rust-sqlite-storage = { path = "../../storages/sqlite-storage" } +whatsapp-rust = { path = "../..", default-features = false, features = [ + "danger-skip-tls-verify", + "debug-diagnostics", + "moka-cache", + "simd", + "tokio-runtime", + "tokio-native", + "signal", +] } whatsapp-rust-tokio-transport = { path = "../../transports/tokio-transport", features = [ "danger-skip-tls-verify", ] } diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index c246c47f8..baf323450 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use wacore::store::InMemoryBackend; use wacore::store::traits::TcTokenEntry; use wacore::types::events::{ChannelEventHandler, Event}; use wacore_binary::node::Node; @@ -7,20 +8,9 @@ use whatsapp_rust::Jid; use whatsapp_rust::bot::Bot; use whatsapp_rust::store::traits::Backend; use whatsapp_rust::waproto::whatsapp as wa; -use whatsapp_rust_sqlite_storage::SqliteStore; use whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory; use whatsapp_rust_ureq_http_client::UreqHttpClient; -/// Creates a SqliteStore with a unique in-memory database for test isolation. -pub async fn create_test_store(prefix: &str) -> anyhow::Result<SqliteStore> { - let db = format!( - "file:{}_{}?mode=memory&cache=shared", - prefix, - uuid::Uuid::new_v4() - ); - Ok(SqliteStore::new(&db).await?) -} - /// Returns the mock server WebSocket URL from env, or the default. pub fn mock_server_url() -> String { std::env::var("MOCK_SERVER_URL").unwrap_or_else(|_| "wss://127.0.0.1:8080/ws/chat".to_string()) @@ -64,9 +54,8 @@ impl TestClient { Self::connect_inner(prefix, Some(push_name.to_string())).await } - async fn connect_inner(prefix: &str, push_name: Option<String>) -> anyhow::Result<Self> { - let store = create_test_store(prefix).await?; - let backend = Arc::new(store) as Arc<dyn Backend>; + async fn connect_inner(_prefix: &str, push_name: Option<String>) -> anyhow::Result<Self> { + let backend = Arc::new(InMemoryBackend::new()) as Arc<dyn Backend>; let transport_factory = TokioWebSocketTransportFactory::new().with_url(mock_server_url()); let (event_handler, event_rx) = ChannelEventHandler::new(); diff --git a/tests/e2e/tests/digest_key.rs b/tests/e2e/tests/digest_key.rs index 595e49106..22ca59a7d 100644 --- a/tests/e2e/tests/digest_key.rs +++ b/tests/e2e/tests/digest_key.rs @@ -68,23 +68,14 @@ async fn test_digest_key_hash_matches_server() -> anyhow::Result<()> { ) }); - use prost::Message; - let record = whatsapp_rust::waproto::whatsapp::PreKeyRecordStructure::decode( - record_bytes.as_slice(), - )?; - let pk = record - .public_key + let pk = wacore::prekeys::extract_prekey_public_key(&record_bytes) .unwrap_or_else(|| panic!("Prekey {} has no public_key field", prekey_id)); - prekey_pubkeys.push(pk); + prekey_pubkeys.push(pk.to_vec()); } - // Compute SHA-1 digest the same way the server does - let local_hash = wacore::prekeys::compute_key_bundle_digest( - identity_pub, - skey_pub, - skey_sig, - &prekey_pubkeys, - ); + let pubkey_refs: Vec<&[u8]> = prekey_pubkeys.iter().map(|v| v.as_slice()).collect(); + let local_hash = + wacore::prekeys::compute_key_bundle_digest(identity_pub, skey_pub, skey_sig, &pubkey_refs); assert_eq!( local_hash, diff --git a/tests/e2e/tests/retry_dm_multidevice.rs b/tests/e2e/tests/retry_dm_multidevice.rs index bcf613d0a..bfdd36571 100644 --- a/tests/e2e/tests/retry_dm_multidevice.rs +++ b/tests/e2e/tests/retry_dm_multidevice.rs @@ -3,6 +3,7 @@ use e2e_tests::{TestClient, send_and_expect_text, text_msg}; use log::info; use wacore::types::events::Event; +use wacore_binary::JidExt as _; use wacore_binary::node::Node; use whatsapp_rust::{NodeFilter, SendOptions}; @@ -107,11 +108,17 @@ async fn test_dm_retry_recovers_after_session_deletion() -> anyhow::Result<()> { Some("1"), "Retry resend should mark the payload with count=1" ); - let jid_b_str = jid_b.to_string(); - assert_eq!( - retry_node.attrs().optional_string("to").as_deref(), - Some(jid_b_str.as_str()), - "Retry resend should keep the user-level chat target" + // WA Web uses the raw receipt `from` (with device suffix) as the stanza `to`. + // The mock server may or may not include a device suffix depending on timing, + // so check user part only. + let retry_to = retry_node + .attrs() + .optional_jid("to") + .expect("Retry resend should have a 'to' attribute"); + assert!( + retry_to.is_same_user_as(&jid_b), + "Retry resend 'to' should target the same user (got {retry_to}, expected user {})", + jid_b ); info!("Retry recovered after B deleted its session with A"); diff --git a/transports/tokio-transport/src/lib.rs b/transports/tokio-transport/src/lib.rs index cd8ebd462..d2569bd61 100644 --- a/transports/tokio-transport/src/lib.rs +++ b/transports/tokio-transport/src/lib.rs @@ -15,7 +15,7 @@ use wacore::net::{Transport, TransportEvent, TransportFactory, WHATSAPP_WEB_WS_U pub use tokio_websockets::Connector; -const EVENT_CHANNEL_CAPACITY: usize = 10_000; +const EVENT_CHANNEL_CAPACITY: usize = 1_024; static CRYPTO_PROVIDER_INIT: Once = Once::new(); @@ -131,7 +131,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send + 'static> WsTransport<S> { #[async_trait] impl<S: AsyncRead + AsyncWrite + Unpin + Send + 'static> Transport for WsTransport<S> { - async fn send(&self, data: Vec<u8>) -> Result<(), anyhow::Error> { + async fn send(&self, data: bytes::Bytes) -> Result<(), anyhow::Error> { let mut guard = self.sink.lock().await; let sink = guard .as_mut() diff --git a/wacore/benches/send_receive_benchmark.rs b/wacore/benches/send_receive_benchmark.rs index 2df58f1ff..3d44002b2 100644 --- a/wacore/benches/send_receive_benchmark.rs +++ b/wacore/benches/send_receive_benchmark.rs @@ -590,11 +590,12 @@ fn setup_group_recv() -> GrpRecvData { #[library_benchmark] #[bench::text(setup = setup_dm_send)] fn bench_dm_send(mut d: DmSendData) { + let signal_addr = d.bob_jid.to_protocol_address(); let node = futures::executor::block_on(prepare_peer_stanza( &mut d.alice.sessions, &mut d.alice.identity, - d.bob_jid.clone(), d.bob_jid, + &signal_addr, &d.msg, "b-001".into(), )) diff --git a/wacore/binary/src/lib.rs b/wacore/binary/src/lib.rs index bde27e613..ee41589d1 100644 --- a/wacore/binary/src/lib.rs +++ b/wacore/binary/src/lib.rs @@ -11,6 +11,7 @@ pub mod marshal; pub mod node; pub mod token; pub mod util; +pub mod zlib_pool; pub use attrs::{AttrParser, AttrParserRef}; pub use compact_str::CompactString; diff --git a/wacore/binary/src/marshal.rs b/wacore/binary/src/marshal.rs index 26f9ba03d..1c21b9221 100644 --- a/wacore/binary/src/marshal.rs +++ b/wacore/binary/src/marshal.rs @@ -144,7 +144,15 @@ fn should_auto_reserve_node(node: &Node) -> bool { match &node.content { Some(NodeContent::Bytes(bytes)) => bytes.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, Some(NodeContent::String(text)) => text.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, - Some(NodeContent::Nodes(children)) => children.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD, + Some(NodeContent::Nodes(children)) => { + if children.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD { + return true; + } + // Check one level deeper for large nested lists (e.g., <iq> -> <list> -> 812 keys) + children.iter().any(|child| { + matches!(&child.content, Some(NodeContent::Nodes(gc)) if gc.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD) + }) + } None => false, } } @@ -158,7 +166,15 @@ fn should_auto_reserve_node_ref(node: &NodeRef<'_>) -> bool { match node.content.as_deref() { Some(NodeContentRef::Bytes(bytes)) => bytes.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, Some(NodeContentRef::String(text)) => text.len() >= AUTO_RESERVE_SCALAR_THRESHOLD, - Some(NodeContentRef::Nodes(children)) => children.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD, + Some(NodeContentRef::Nodes(children)) => { + if children.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD { + return true; + } + // Check one level deeper for large nested lists (e.g., <iq> -> <list> -> 812 keys) + children.iter().any(|child| { + matches!(child.content.as_deref(), Some(NodeContentRef::Nodes(gc)) if gc.len() >= AUTO_RESERVE_CHILDREN_THRESHOLD) + }) + } None => false, } } diff --git a/wacore/binary/src/node.rs b/wacore/binary/src/node.rs index 52fbccb38..e95e5b5b8 100644 --- a/wacore/binary/src/node.rs +++ b/wacore/binary/src/node.rs @@ -802,6 +802,21 @@ impl OwnedNodeRef { self.inner.get().to_owned() } + /// Return a zero-copy `Bytes` sub-view for a slice that borrows from this + /// node's backing buffer. Panics if `slice` does not point within the buffer. + pub fn slice_bytes(&self, slice: &[u8]) -> Bytes { + let cart = &self.inner.backing_cart().0; + let base = cart.as_ptr() as usize; + let end = base + cart.len(); + let ptr = slice.as_ptr() as usize; + assert!( + ptr >= base && ptr + slice.len() <= end, + "slice is not within the backing buffer" + ); + let offset = ptr - base; + cart.slice(offset..offset + slice.len()) + } + /// The tag name of this node. #[inline] pub fn tag(&self) -> &str { diff --git a/wacore/binary/src/util.rs b/wacore/binary/src/util.rs index 0be9f7fc2..8ac9f320c 100644 --- a/wacore/binary/src/util.rs +++ b/wacore/binary/src/util.rs @@ -1,27 +1,14 @@ use crate::error::{BinaryError, Result}; +use crate::zlib_pool::decompress_zlib_pooled; use bytes::{Buf, Bytes, BytesMut}; -use flate2::read::ZlibDecoder; use std::borrow::Cow; -use std::io::Read; /// Protocol frames larger than 16 MiB after decompression are rejected. -/// WhatsApp messages are typically small; this guards against malicious -/// or corrupt compressed payloads that would expand into huge allocations. const MAX_DECOMPRESSED_SIZE: u64 = 16 * 1024 * 1024; fn decompress_zlib(compressed: &[u8]) -> Result<Vec<u8>> { - let estimated = (compressed.len() * 4).clamp(256, 64 * 1024); - let mut out = Vec::with_capacity(estimated); - ZlibDecoder::new(compressed) - .take(MAX_DECOMPRESSED_SIZE + 1) - .read_to_end(&mut out) - .map_err(|e| BinaryError::Zlib(e.to_string()))?; - if out.len() as u64 > MAX_DECOMPRESSED_SIZE { - return Err(BinaryError::Zlib(format!( - "decompressed payload exceeds {MAX_DECOMPRESSED_SIZE} bytes" - ))); - } - Ok(out) + decompress_zlib_pooled(compressed, MAX_DECOMPRESSED_SIZE) + .map_err(|e| BinaryError::Zlib(e.to_string())) } pub fn unpack(data: &[u8]) -> Result<Cow<'_, [u8]>> { diff --git a/wacore/binary/src/zlib_pool.rs b/wacore/binary/src/zlib_pool.rs new file mode 100644 index 000000000..ce7b9711e --- /dev/null +++ b/wacore/binary/src/zlib_pool.rs @@ -0,0 +1,89 @@ +use flate2::{Decompress, FlushDecompress, Status}; +use std::cell::RefCell; +use std::io; + +thread_local! { + static DECOMPRESSOR: RefCell<(Decompress, Vec<u8>)> = RefCell::new(( + Decompress::new(true), + Vec::with_capacity(4096), + )); +} + +/// Decompress zlib data using a pooled decompressor. +/// +/// Reuses the `flate2::Decompress` internal state (~48 KB) and the output +/// buffer across calls on the same thread, avoiding repeated heap allocations. +pub fn decompress_zlib_pooled(compressed: &[u8], max_size: u64) -> io::Result<Vec<u8>> { + DECOMPRESSOR.with(|cell| { + let (decompressor, scratch) = &mut *cell.borrow_mut(); + decompressor.reset(true); + scratch.clear(); + + let estimated = (compressed.len() * 4).clamp(256, 64 * 1024); + if scratch.capacity() < estimated { + scratch.reserve(estimated - scratch.capacity()); + } + + // Cap output growth to max_size + 1 so we detect oversized payloads + // without allocating unbounded memory from a compressed bomb. + let cap = (max_size as usize).saturating_add(1); + + let mut input_offset = 0; + loop { + // Enforce cap before decompress_vec can grow the buffer + if scratch.len() >= cap { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("decompressed payload exceeds {max_size} bytes"), + )); + } + + let prev_in = decompressor.total_in(); + let prev_out = decompressor.total_out(); + + let status = decompressor + .decompress_vec( + &compressed[input_offset..], + scratch, + FlushDecompress::Finish, + ) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + + input_offset = decompressor.total_in() as usize; + + if scratch.len() as u64 > max_size { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("decompressed payload exceeds {max_size} bytes"), + )); + } + + match status { + Status::StreamEnd => break, + Status::Ok => { + // Grow but never past the cap + let want = scratch.capacity().max(4096).min(cap - scratch.len()); + scratch.reserve(want); + } + Status::BufError => { + if decompressor.total_in() == prev_in && decompressor.total_out() == prev_out { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "zlib stream truncated (no progress)", + )); + } + let want = scratch.capacity().max(4096).min(cap - scratch.len()); + scratch.reserve(want); + } + } + } + + // Move the Vec out (zero-copy), then restore scratch with fresh capacity. + // Callers (unpack_bytes, history_sync) wrap in Bytes::from() which takes + // ownership of the Vec's allocation, so no extra copy occurs. + let result = std::mem::take(scratch); + // Pre-allocate for next call so the first decompress_vec doesn't start at 0 + scratch.reserve(4096); + Ok(result) + }) +} diff --git a/wacore/src/history_sync.rs b/wacore/src/history_sync.rs index c511352b2..46fec6046 100644 --- a/wacore/src/history_sync.rs +++ b/wacore/src/history_sync.rs @@ -1,7 +1,6 @@ use bytes::Bytes; -use flate2::read::ZlibDecoder; -use std::io::Read; use thiserror::Error; +use wacore_binary::zlib_pool::decompress_zlib_pooled; #[derive(Debug, Error)] pub enum HistorySyncError { @@ -49,23 +48,14 @@ pub fn process_history_sync( compressed_data: Vec<u8>, own_user: Option<&str>, retain_blob: bool, - compressed_size_hint: Option<u64>, + _compressed_size_hint: Option<u64>, ) -> Result<HistorySyncResult, HistorySyncError> { // Hard limit to prevent OOM on malformed blobs. // Typical InitialBootstrap: 5-20 MB decompressed. const MAX_DECOMPRESSED: u64 = 64 * 1024 * 1024; - let estimated = compressed_size_hint - .and_then(|s| usize::try_from(s).ok()) - .map(|s| s * 4) - .unwrap_or_else(|| compressed_data.len() * 4) - .clamp(256, MAX_DECOMPRESSED as usize); - let mut decompressed = Vec::with_capacity(estimated); - { - let decoder = ZlibDecoder::new(compressed_data.as_slice()); - let mut limited = decoder.take(MAX_DECOMPRESSED); - limited.read_to_end(&mut decompressed)?; - } + let decompressed = decompress_zlib_pooled(&compressed_data, MAX_DECOMPRESSED) + .map_err(HistorySyncError::DecompressionError)?; drop(compressed_data); let buf = Bytes::from(decompressed); diff --git a/wacore/src/iq/prekeys.rs b/wacore/src/iq/prekeys.rs index 801d2b070..c0094bff6 100644 --- a/wacore/src/iq/prekeys.rs +++ b/wacore/src/iq/prekeys.rs @@ -367,13 +367,13 @@ impl IqSpec for PreKeyUploadSpec { fn build_iq(&self) -> InfoQuery<'static> { let content = PreKeyUtils::build_upload_prekeys_request( self.registration_id, - self.identity_key.public_key_bytes().to_vec(), + self.identity_key.public_key_bytes(), self.signed_pre_key_id, - self.signed_pre_key_public.public_key_bytes().to_vec(), - self.signed_pre_key_signature.clone(), + self.signed_pre_key_public.public_key_bytes(), + &self.signed_pre_key_signature, self.pre_keys .iter() - .map(|(id, pk)| (*id, pk.public_key_bytes().to_vec())), + .map(|(id, pk)| (*id, pk.public_key_bytes())), ); InfoQuery::set( diff --git a/wacore/src/iq/props.rs b/wacore/src/iq/props.rs index 28181449a..7b99f8800 100644 --- a/wacore/src/iq/props.rs +++ b/wacore/src/iq/props.rs @@ -26,6 +26,7 @@ use crate::iq::spec::IqSpec; use crate::protocol::ProtocolNode; use crate::request::InfoQuery; +use wacore_binary::CompactString; use wacore_binary::builder::NodeBuilder; use wacore_binary::{Jid, Server}; use wacore_binary::{Node, NodeContent, NodeRef}; @@ -68,8 +69,8 @@ pub const PROPS_PROTOCOL_VERSION: &str = "1"; pub struct AbProp { /// The config code (property identifier). pub config_code: u32, - /// The config value. - pub config_value: String, + /// The config value. CompactString inlines values <=24 bytes (covers most props). + pub config_value: CompactString, /// Optional experiment exposure key. pub config_expo_key: Option<u32>, } @@ -82,7 +83,7 @@ impl crate::protocol::ProtocolNode for AbProp { fn into_node(self) -> Node { let mut builder = NodeBuilder::new("prop") .attr("config_code", self.config_code.to_string()) - .attr("config_value", &self.config_value); + .attr("config_value", &*self.config_value); if let Some(expo_key) = self.config_expo_key { builder = builder.attr("config_expo_key", expo_key.to_string()); @@ -105,8 +106,8 @@ impl crate::protocol::ProtocolNode for AbProp { return Err(anyhow::anyhow!("config_code must be >= 1")); } let config_value = optional_attr(node, "config_value") - .ok_or_else(|| anyhow::anyhow!("missing config_value in prop"))? - .into_owned(); + .ok_or_else(|| anyhow::anyhow!("missing config_value in prop"))?; + let config_value = CompactString::from(config_value.as_ref()); let config_expo_key = optional_attr(node, "config_expo_key").and_then(|s| s.parse().ok()); Ok(Self { @@ -483,7 +484,7 @@ mod tests { fn test_ab_prop_protocol_node_round_trip() { let prop = AbProp { config_code: 123, - config_value: "test_value".to_string(), + config_value: "test_value".into(), config_expo_key: Some(456), }; @@ -499,7 +500,7 @@ mod tests { fn test_ab_prop_protocol_node_no_expo_key() { let prop = AbProp { config_code: 789, - config_value: "another_value".to_string(), + config_value: "another_value".into(), config_expo_key: None, }; @@ -522,7 +523,7 @@ mod tests { props: vec![ AbPropConfig::Experiment(AbProp { config_code: 100, - config_value: "value1".to_string(), + config_value: "value1".into(), config_expo_key: None, }), AbPropConfig::Sampling(SamplingProp { @@ -531,7 +532,7 @@ mod tests { }), AbPropConfig::Experiment(AbProp { config_code: 200, - config_value: "value2".to_string(), + config_value: "value2".into(), config_expo_key: Some(99), }), ], diff --git a/wacore/src/net.rs b/wacore/src/net.rs index 2859cf524..83fd28b1a 100644 --- a/wacore/src/net.rs +++ b/wacore/src/net.rs @@ -24,7 +24,7 @@ pub enum TransportEvent { #[cfg_attr(not(target_arch = "wasm32"), async_trait)] pub trait Transport: Send + Sync { /// Sends raw data to the server. - async fn send(&self, data: Vec<u8>) -> Result<(), anyhow::Error>; + async fn send(&self, data: Bytes) -> Result<(), anyhow::Error>; /// Closes the connection. async fn disconnect(&self); diff --git a/wacore/src/prekeys.rs b/wacore/src/prekeys.rs index a909ed00a..d03132717 100644 --- a/wacore/src/prekeys.rs +++ b/wacore/src/prekeys.rs @@ -10,12 +10,12 @@ pub struct PreKeyUtils; /// Compute SHA-1 digest of a key bundle for validation against server. /// /// Matches WA Web's `validateLocalKeyBundle` hash computation: -/// SHA-1(identity_pub_key || signed_prekey_pub || signed_prekey_signature || prekey_pub_1 || prekey_pub_2 || ...) +/// SHA-1(identity_pub_key || signed_prekey_pub || signed_prekey_signature || prekey_pub_1 || ...) pub fn compute_key_bundle_digest( identity_pub_key: &[u8], signed_prekey_pub: &[u8], signed_prekey_signature: &[u8], - prekey_pubkeys: &[Vec<u8>], + prekey_pubkeys: &[&[u8]], ) -> Vec<u8> { use sha1::Digest; let mut hasher = sha1::Sha1::new(); @@ -28,6 +28,73 @@ pub fn compute_key_bundle_digest( hasher.finalize().to_vec() } +/// Extract the `publicKey` field (tag 2) from a protobuf-encoded PreKeyRecordStructure +/// without full prost decode. Uses last-one-wins semantics per protobuf spec. +/// Skips unknown fields gracefully. +pub fn extract_prekey_public_key(record: &[u8]) -> Option<&[u8]> { + let mut pos = 0; + let mut result: Option<&[u8]> = None; + while pos < record.len() { + let (tag_byte, consumed) = decode_varint(&record[pos..])?; + pos += consumed; + let field_number = (tag_byte >> 3) as u32; + let wire_type = (tag_byte & 0x7) as u32; + match wire_type { + // varint + 0 => { + let (_, c) = decode_varint(&record[pos..])?; + pos += c; + } + // length-delimited + 2 => { + let (len, c) = decode_varint(&record[pos..])?; + pos += c; + let len = len as usize; + if pos + len > record.len() { + return result; + } + if field_number == 2 { + result = Some(&record[pos..pos + len]); + } + pos += len; + } + // fixed64 + 1 => { + if pos + 8 > record.len() { + return result; + } + pos += 8; + } + // fixed32 + 5 => { + if pos + 4 > record.len() { + return result; + } + pos += 4; + } + // Unknown wire type -- skip gracefully + _ => return result, + } + } + result +} + +fn decode_varint(buf: &[u8]) -> Option<(u64, usize)> { + let mut result: u64 = 0; + for (i, &byte) in buf.iter().enumerate().take(10) { + // The 10th byte (i==9) carries the highest bits; only the low bit + // is valid payload (64 - 9*7 = 1). Reject if more bits are set. + if i == 9 && (byte & 0x7F) > 1 { + return None; + } + result |= ((byte & 0x7F) as u64) << (i * 7); + if byte & 0x80 == 0 { + return Some((result, i + 1)); + } + } + None +} + impl PreKeyUtils { pub fn build_fetch_prekeys_request(jids: &[Jid], reason: Option<&str>) -> Node { let user_nodes = jids.iter().map(|jid| { @@ -41,30 +108,29 @@ impl PreKeyUtils { NodeBuilder::new("key").children(user_nodes).build() } - pub fn build_upload_prekeys_request( + pub fn build_upload_prekeys_request<'a>( registration_id: u32, - identity_key_bytes: Vec<u8>, + identity_key_bytes: &[u8], signed_pre_key_id: u32, - signed_pre_key_public_bytes: Vec<u8>, - signed_pre_key_signature: Vec<u8>, - pre_keys: impl IntoIterator<Item = (u32, Vec<u8>)>, + signed_pre_key_public_bytes: &[u8], + signed_pre_key_signature: &[u8], + pre_keys: impl IntoIterator<Item = (u32, &'a [u8])>, ) -> Vec<Node> { let pre_keys = pre_keys.into_iter(); let (lower, upper) = pre_keys.size_hint(); let mut pre_key_nodes = Vec::with_capacity(upper.unwrap_or(lower)); for (pre_key_id, public_bytes) in pre_keys { - let id_bytes = pre_key_id.to_be_bytes()[1..].to_vec(); let node = NodeBuilder::new("key") .children([ - NodeBuilder::new("id").bytes(id_bytes).build(), + NodeBuilder::new("id") + .bytes(pre_key_id.to_be_bytes()[1..].to_vec()) + .build(), NodeBuilder::new("value").bytes(public_bytes).build(), ]) .build(); pre_key_nodes.push(node); } - let registration_id_bytes = registration_id.to_be_bytes().to_vec(); - let signed_pre_key_node = NodeBuilder::new("skey") .children([ NodeBuilder::new("id") @@ -79,13 +145,11 @@ impl PreKeyUtils { ]) .build(); - let type_bytes = vec![5u8]; - vec![ NodeBuilder::new("registration") - .bytes(registration_id_bytes) + .bytes(registration_id.to_be_bytes().to_vec()) .build(), - NodeBuilder::new("type").bytes(type_bytes).build(), + NodeBuilder::new("type").bytes(vec![5u8]).build(), NodeBuilder::new("identity") .bytes(identity_key_bytes) .build(), diff --git a/wacore/src/send.rs b/wacore/src/send.rs index 44ad70363..d9b8cef7f 100644 --- a/wacore/src/send.rs +++ b/wacore/src/send.rs @@ -1,7 +1,7 @@ use crate::client::context::{GroupInfo, SendContextResolver}; use crate::libsignal::protocol::{ - CiphertextMessage, SENDERKEY_MESSAGE_CURRENT_VERSION, SenderKeyMessage, SenderKeyStore, - SignalProtocolError, UsePQRatchet, message_encrypt, process_prekey_bundle, + CiphertextMessage, ProtocolAddress, SENDERKEY_MESSAGE_CURRENT_VERSION, SenderKeyMessage, + SenderKeyStore, SignalProtocolError, UsePQRatchet, message_encrypt, process_prekey_bundle, }; use crate::messages::MessageUtils; use crate::reporting_token::{ @@ -247,7 +247,7 @@ pub fn should_hide_decrypt_fail(msg: &wa::Message) -> bool { pub async fn encrypt_group_message<S, R>( sender_key_store: &mut S, group_jid: &Jid, - sender_jid: &Jid, + sender_address: &ProtocolAddress, plaintext: &[u8], csprng: &mut R, ) -> Result<SenderKeyMessage> @@ -255,8 +255,7 @@ where S: SenderKeyStore + ?Sized, R: Rng + CryptoRng, { - let sender_address = sender_jid.to_protocol_address(); - let sender_key_name = make_sender_key_name(group_jid, &sender_address); + let sender_key_name = make_sender_key_name(group_jid, sender_address); log::debug!( "Attempting to load sender key for group {} sender {}", sender_key_name.group_id(), @@ -830,7 +829,7 @@ pub async fn prepare_peer_stanza<S, I>( session_store: &mut S, identity_store: &mut I, transport_jid: Jid, - encryption_jid: Jid, + signal_address: &ProtocolAddress, message: &wa::Message, request_id: String, ) -> Result<Node> @@ -839,10 +838,9 @@ where I: crate::libsignal::protocol::IdentityKeyStore, { 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?; + message_encrypt(&plaintext, signal_address, session_store, identity_store).await?; let (enc_type, _, serialized_bytes) = extract_ciphertext(encrypted_message) .ok_or_else(|| anyhow!("Unexpected peer encryption message type"))?; @@ -1053,6 +1051,8 @@ pub async fn prepare_group_stanza< let mut phash_for_stanza: Option<String> = None; let mut skdm_encrypted_devices: Vec<Jid> = Vec::new(); + let sender_address = own_sending_jid.to_protocol_address(); + // Determine if we need to distribute SKDM and to which devices let distribution_list: Option<Vec<Jid>> = if let Some(target_devices) = skdm_target_devices { // Use the specific list of devices that need SKDM @@ -1183,7 +1183,7 @@ pub async fn prepare_group_stanza< let axolotl_skdm_bytes = create_sender_key_distribution_message_for_group( stores.sender_key_store, &to_jid, - &own_sending_jid, + &sender_address, ) .await?; @@ -1247,7 +1247,7 @@ pub async fn prepare_group_stanza< let skmsg = encrypt_group_message( stores.sender_key_store, &to_jid, - &own_sending_jid, + &sender_address, &plaintext, &mut rand::make_rng::<rand::rngs::StdRng>(), ) @@ -1334,10 +1334,9 @@ pub async fn prepare_group_stanza< pub async fn create_sender_key_distribution_message_for_group( store: &mut (dyn SenderKeyStore + Send + Sync), group_jid: &Jid, - own_sending_jid: &Jid, + sender_address: &ProtocolAddress, ) -> Result<Vec<u8>> { - let sender_address = own_sending_jid.to_protocol_address(); - let sender_key_name = make_sender_key_name(group_jid, &sender_address); + let sender_key_name = make_sender_key_name(group_jid, sender_address); let mut rng = rand::make_rng::<rand::rngs::StdRng>(); let skdm = crate::libsignal::protocol::create_sender_key_distribution_message( diff --git a/wacore/src/store/ab_props.rs b/wacore/src/store/ab_props.rs index 59eafdb60..b6afd0bcb 100644 --- a/wacore/src/store/ab_props.rs +++ b/wacore/src/store/ab_props.rs @@ -10,12 +10,13 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicBool, Ordering}; use async_lock::RwLock; +use wacore_binary::CompactString; use crate::iq::props::{AbPropConfig, PropsResponse}; /// In-memory cache of AB experiment properties, populated on connect. pub struct AbPropsCache { - props: RwLock<HashMap<u32, String>>, + props: RwLock<HashMap<u32, CompactString>>, /// Guards against applying a delta into an empty cache on cold start. seeded: AtomicBool, } @@ -50,7 +51,7 @@ impl AbPropsCache { } } - pub async fn get(&self, config_code: u32) -> Option<String> { + pub async fn get(&self, config_code: u32) -> Option<CompactString> { self.props.read().await.get(&config_code).cloned() } @@ -103,7 +104,7 @@ mod tests { fn experiment(code: u32, value: &str) -> AbPropConfig { AbPropConfig::Experiment(AbProp { config_code: code, - config_value: value.to_string(), + config_value: value.into(), config_expo_key: None, }) } diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index a5f4a2b68..ea942fc19 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -13,6 +13,7 @@ use crate::store::error::Result; use crate::store::traits::*; use async_lock::Mutex; use async_trait::async_trait; +use bytes::Bytes; use wacore_appstate::processor::AppStateMutationMAC; /// Key for the sent-message store: `(chat_jid, message_id)`. @@ -26,7 +27,7 @@ struct SentMessageEntry { /// Key for pre-keys: `id`. struct PreKeyEntry { - record: Vec<u8>, + record: Bytes, } /// Key for base-key collision detection: `(address, message_id)`. @@ -168,13 +169,13 @@ impl SignalStore for InMemoryBackend { self.state.lock().await.prekeys.insert( id, PreKeyEntry { - record: record.to_vec(), + record: Bytes::copy_from_slice(record), }, ); Ok(()) } - async fn store_prekeys_batch(&self, keys: &[(u32, Vec<u8>)], _uploaded: bool) -> Result<()> { + async fn store_prekeys_batch(&self, keys: &[(u32, Bytes)], _uploaded: bool) -> Result<()> { let mut state = self.state.lock().await; for (id, record) in keys { state.prekeys.insert( @@ -187,7 +188,7 @@ impl SignalStore for InMemoryBackend { Ok(()) } - async fn load_prekey(&self, id: u32) -> Result<Option<Vec<u8>>> { + async fn load_prekey(&self, id: u32) -> Result<Option<Bytes>> { Ok(self .state .lock() @@ -197,7 +198,7 @@ impl SignalStore for InMemoryBackend { .map(|e| e.record.clone())) } - async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Vec<u8>)>> { + async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Bytes)>> { let state = self.state.lock().await; let mut result = Vec::with_capacity(ids.len()); for &id in ids { diff --git a/wacore/src/store/traits.rs b/wacore/src/store/traits.rs index 875e1b6fe..f3717c389 100644 --- a/wacore/src/store/traits.rs +++ b/wacore/src/store/traits.rs @@ -10,6 +10,7 @@ use crate::appstate::hash::HashState; use crate::store::error::Result; use async_trait::async_trait; +use bytes::Bytes; use serde::{Deserialize, Serialize}; use wacore_appstate::processor::AppStateMutationMAC; @@ -116,7 +117,7 @@ pub trait SignalStore: Send + Sync { /// Store multiple pre-keys in a single batch operation. /// Default implementation falls back to individual `store_prekey` calls. - async fn store_prekeys_batch(&self, keys: &[(u32, Vec<u8>)], uploaded: bool) -> Result<()> { + async fn store_prekeys_batch(&self, keys: &[(u32, Bytes)], uploaded: bool) -> Result<()> { for (id, record) in keys { self.store_prekey(*id, record, uploaded).await?; } @@ -124,11 +125,11 @@ pub trait SignalStore: Send + Sync { } /// Load a pre-key by ID. - async fn load_prekey(&self, id: u32) -> Result<Option<Vec<u8>>>; + async fn load_prekey(&self, id: u32) -> Result<Option<Bytes>>; /// Load multiple pre-keys by ID in a single batch operation. /// Returns only the keys that exist. - async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Vec<u8>)>> { + async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Bytes)>> { let mut result = Vec::with_capacity(ids.len()); for &id in ids { if let Some(record) = self.load_prekey(id).await? {