From b51b20373e341fdfe2cce1ad7eb8a011d05ae841 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 16:52:16 +0000 Subject: [PATCH 1/3] perf(prekeys): avoid full record decode on the pre-key upload path upload_pre_keys_pass re-decoded every stored PreKeyRecordStructure with prost just to read the public key for the upload IQ. A full decode also allocates a copy of the private key the upload never uses; at the default batch of 812 one-time pre-keys (WA Web fidelity) that is ~2 throwaway Vec allocations per record on the connect/registration path. Read the public-key field straight from the encoded record via the existing wacore::prekeys::extract_prekey_public_key helper instead. Behavior-preserving: same public-key bytes, same skip-with-warning handling for a record missing the field. Trims the allocation count/volume CodSpeed attributes to connect_to_ready (the dominant cost there remains the intentional 812 X25519 keygens, untouched). --- src/prekeys.rs | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/src/prekeys.rs b/src/prekeys.rs index ec131b411..61854e829 100644 --- a/src/prekeys.rs +++ b/src/prekeys.rs @@ -551,17 +551,19 @@ impl Client { } let pre_key_pairs = { - use prost::Message; let mut pairs: Vec<(u32, PublicKey)> = Vec::with_capacity(rows.len()); for (id, record) in &rows { - let public_key = waproto::whatsapp::PreKeyRecordStructure::decode(&record[..]) - .map_err(anyhow::Error::from) - .and_then(|structure| { - let raw = structure - .public_key - .ok_or_else(|| anyhow::anyhow!("record missing public key"))?; - Ok(PublicKey::from_djb_public_key_bytes(&raw)?) - }); + // Pull the public key straight out of the encoded record (field 2) + // rather than a full prost decode of the PreKeyRecordStructure: the + // upload only needs the public key, while a full decode also copies + // the private key into its own Vec. At the default batch of 812 keys + // that is ~2 throwaway allocations per record on the connect path. + let public_key = match wacore::prekeys::extract_prekey_public_key(record) { + Some(raw) => { + PublicKey::from_djb_public_key_bytes(raw).map_err(anyhow::Error::from) + } + None => Err(anyhow::anyhow!("record missing public key")), + }; match public_key { Ok(public_key) => pairs.push((*id, public_key)), Err(e) => log::warn!("skipping undecodable prekey record {id}: {e:?}"), From 8e437100a78265a879687e3172c9ac1b6fc8c3ef Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 17:11:00 +0000 Subject: [PATCH 2/3] fix(prekeys): reject malformed records in extract_prekey_public_key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review on the upload-path optimization (#900) flagged that extract_prekey_public_key kept the last-seen public key even when the record's protobuf tail was truncated or used an invalid wire type. The upload path would then ship a key the consume path (get_pre_key's full PreKeyRecordStructure::decode) later rejects — handing a peer a key this device can't use to decrypt their first message. Validate the record framing end-to-end: a malformed varint, a truncated length-delimited/fixed field, or an unsupported wire type now yields None, matching what a full prost decode rejects. The upload and digestKey callers already skip on None. Add a test asserting parity with prost decode on a record with a valid publicKey but a malformed tail. --- src/prekeys.rs | 6 +++++- wacore/src/prekeys.rs | 49 ++++++++++++++++++++++++++++++++++++------- 2 files changed, 46 insertions(+), 9 deletions(-) diff --git a/src/prekeys.rs b/src/prekeys.rs index 61854e829..5cb9ee2f2 100644 --- a/src/prekeys.rs +++ b/src/prekeys.rs @@ -558,11 +558,15 @@ impl Client { // upload only needs the public key, while a full decode also copies // the private key into its own Vec. At the default batch of 812 keys // that is ~2 throwaway allocations per record on the connect path. + // The extractor validates the record framing, so a malformed record + // returns None and is skipped here — same rejection the consume path + // (`get_pre_key`'s full decode) makes, so we never upload a key this + // device couldn't later decode with. let public_key = match wacore::prekeys::extract_prekey_public_key(record) { Some(raw) => { PublicKey::from_djb_public_key_bytes(raw).map_err(anyhow::Error::from) } - None => Err(anyhow::anyhow!("record missing public key")), + None => Err(anyhow::anyhow!("record missing or malformed public key")), }; match public_key { Ok(public_key) => pairs.push((*id, public_key)), diff --git a/wacore/src/prekeys.rs b/wacore/src/prekeys.rs index 97aa93320..3eba6e46b 100644 --- a/wacore/src/prekeys.rs +++ b/wacore/src/prekeys.rs @@ -28,9 +28,16 @@ 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. +/// Extract the `publicKey` field (tag 2) from a protobuf-encoded +/// PreKeyRecordStructure without a full prost decode. +/// +/// Validates the record framing end-to-end: a malformed varint, a truncated +/// length-delimited/fixed field, or an unsupported wire type (e.g. the +/// deprecated group types) yields `None` — the same records a full +/// `PreKeyRecordStructure::decode` rejects. This keeps callers that skip on +/// `None` (the pre-key upload and the digestKey check) from admitting a record +/// this device could not later decode. Uses last-one-wins semantics for a +/// repeated `publicKey` field, per the protobuf spec. pub fn extract_prekey_public_key(record: &[u8]) -> Option<&[u8]> { let mut pos = 0; let mut result: Option<&[u8]> = None; @@ -51,7 +58,7 @@ pub fn extract_prekey_public_key(record: &[u8]) -> Option<&[u8]> { pos += c; let len = len as usize; if pos + len > record.len() { - return result; + return None; } if field_number == 2 { result = Some(&record[pos..pos + len]); @@ -61,19 +68,19 @@ pub fn extract_prekey_public_key(record: &[u8]) -> Option<&[u8]> { // fixed64 1 => { if pos + 8 > record.len() { - return result; + return None; } pos += 8; } // fixed32 5 => { if pos + 4 > record.len() { - return result; + return None; } pos += 4; } - // Unknown wire type -- skip gracefully - _ => return result, + // Unsupported / invalid wire type (groups, reserved): reject the record. + _ => return None, } } result @@ -412,6 +419,32 @@ mod tests { use wacore_binary::NodeValue; + #[test] + fn extract_prekey_public_key_matches_full_decode_validation() { + use prost::Message; + let public_key = vec![0x05u8; 33]; + let record = waproto::whatsapp::PreKeyRecordStructure { + id: Some(1), + public_key: Some(public_key.clone()), + private_key: Some(vec![0x09u8; 32]), + } + .encode_to_vec(); + + // Well-formed record: the extractor returns the public key. + assert_eq!( + extract_prekey_public_key(&record), + Some(public_key.as_slice()) + ); + + // Truncating into the trailing private_key field leaves a valid publicKey + // earlier in the buffer but a malformed tail. A full prost decode rejects + // it; the extractor must agree (return None) so the upload path never ships + // a record the consume path's full decode would later reject. + let truncated = &record[..record.len() - 1]; + assert!(waproto::whatsapp::PreKeyRecordStructure::decode(truncated).is_err()); + assert_eq!(extract_prekey_public_key(truncated), None); + } + fn create_mock_bundle(device_id: u32) -> PreKeyBundle { let mut rng = rand::make_rng::(); let identity_pair = IdentityKeyPair::generate(&mut rng); From 5f69f980b3105ced4ec662bb02c20f959ee84d9e Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 18:01:33 +0000 Subject: [PATCH 3/3] perf(prekeys): carry generated public keys to upload, full-decode only leftovers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review on #900 kept finding parity gaps in the hand-rolled extractor used on the upload path: it skips fields by wire type without validating protobuf keys the way prost does (e.g. field number 0, or a known field with the wrong wire type), so a record the consume path's full PreKeyRecordStructure::decode would reject could still be uploaded — handing a peer a key this device can't decrypt with. Sidestep the parser entirely on the hot path. The batch generates 812 fresh keypairs and already holds their public keys, so keep them from generation and upload them directly instead of re-reading and re-parsing what was just written. Only the rare leftover (already-stored) window keys still need a read-back, and those use the full, correct PreKeyRecordStructure::decode — identical to the consume path, so a record accepted for upload is always one this device can later decrypt with. This also drops the redundant store reload of the freshly generated window: on the common connect/registration path (an all-fresh window) the upload now reads and decodes nothing, trimming the allocations CodSpeed attributes to connect_to_ready beyond what the extractor approach achieved. extract_prekey_public_key stays for the digestKey check (a local hash compare that only skips on mismatch, never uploads); its strict-framing tests gain malformed-varint and unsupported-wire-type cases. --- src/prekeys.rs | 101 ++++++++++++++++++++++++------------------ wacore/src/prekeys.rs | 16 +++++++ 2 files changed, 74 insertions(+), 43 deletions(-) diff --git a/src/prekeys.rs b/src/prekeys.rs index 5cb9ee2f2..e0a41bdda 100644 --- a/src/prekeys.rs +++ b/src/prekeys.rs @@ -455,6 +455,11 @@ impl Client { wanted, ); + // Public keys of the freshly generated batch, kept from generation so the + // upload never reads them back out of the store and never decodes protobuf + // to recover a key it just held in hand. Empty when the plan only re-offers + // leftover window keys (gen_count == 0). + let mut fresh_pre_keys: Vec<(u32, PublicKey)> = Vec::new(); if plan.gen_count > 0 { let gen_start = plan.gen_start; let gen_count = plan.gen_count as usize; @@ -462,16 +467,18 @@ impl Client { // batch size is caller-configurable, so offload the whole batch to keep // the async executor responsive. Records are encoded into one contiguous // buffer with zero-copy Bytes slices instead of an alloc per record. - let encoded_batch = wacore::runtime::blocking(&*self.runtime, move || { + let (encoded_batch, generated) = wacore::runtime::blocking(&*self.runtime, move || { use prost::Message; // Seed one CSPRNG and advance it per key, rather than reseeding from // entropy on every iteration. let mut rng = rand::make_rng::(); let mut records = Vec::with_capacity(gen_count); + let mut pubkeys = Vec::with_capacity(gen_count); for i in 0..gen_count { let pre_key_id = gen_start + i as u32; let key_pair = KeyPair::generate(&mut rng); + pubkeys.push((pre_key_id, key_pair.public_key)); records.push((pre_key_id, new_pre_key_record(pre_key_id, &key_pair))); } @@ -490,7 +497,7 @@ impl Client { .into_iter() .map(|(id, range)| (id, shared.slice(range))) .collect(); - encoded_batch + (encoded_batch, pubkeys) }) .await; @@ -500,6 +507,7 @@ impl Client { // Propagate errors — uploading a key we can't store locally would cause // decryption failures when the server hands it out. backend.store_prekeys_batch(&encoded_batch, false).await?; + fresh_pre_keys = generated; } // Advance NEXT at GENERATION time (WA Web savePreKeys) and initialise @@ -522,57 +530,64 @@ impl Client { .await .map_err(|e| anyhow::anyhow!("failed to flush prekey watermarks: {e:?}"))?; - // Load the upload window: leftover keys plus the fresh ones. Gaps are - // tolerated (a window key consumed via a retry receipt leaves a hole). - let window_ids: Vec = (0..wanted as u32).map(|i| plan.window_start + i).collect(); - let mut rows = backend.load_prekeys_batch(&window_ids).await?; - rows.sort_unstable_by_key(|(id, _)| *id); - if rows.is_empty() { - // A fully consumed/missing window with no generation would bail - // forever (available > 0 keeps gen_count at 0). Collapse the - // window and rerun the pass so a one-shot caller still uploads. - if plan.gen_count == 0 { - self.persistence_manager - .process_command(DeviceCommand::SetPreKeyWatermarks { - next_pre_key_id: plan.new_next, - first_unupload_pre_key_id: plan.new_next, - }) - .await; - if allow_collapse_retry { - log::warn!( - "prekey window [{}, {}) fully missing; collapsed, regenerating", - plan.window_start, - plan.new_next - ); - return Box::pin(self.upload_pre_keys_pass(false)).await; - } + // Only the leftover (already-stored) window keys are read back and decoded; + // the fresh ones are already in `fresh_pre_keys`. On the common connect path + // the window is all-fresh, so this reads and decodes nothing. Leftover gaps + // are tolerated (a window key consumed via a retry receipt leaves a hole). + let leftover_ids: Vec = (0..plan.available).map(|i| plan.window_start + i).collect(); + let mut leftover_rows = if leftover_ids.is_empty() { + Vec::new() + } else { + backend.load_prekeys_batch(&leftover_ids).await? + }; + leftover_rows.sort_unstable_by_key(|(id, _)| *id); + + if plan.gen_count == 0 && leftover_rows.is_empty() { + // A fully consumed/missing leftover window with no generation would bail + // forever (available > 0 keeps gen_count at 0). Collapse the window and + // rerun the pass so a one-shot caller still uploads. + self.persistence_manager + .process_command(DeviceCommand::SetPreKeyWatermarks { + next_pre_key_id: plan.new_next, + first_unupload_pre_key_id: plan.new_next, + }) + .await; + if allow_collapse_retry { + log::warn!( + "prekey window [{}, {}) fully missing; collapsed, regenerating", + plan.window_start, + plan.new_next + ); + return Box::pin(self.upload_pre_keys_pass(false)).await; } anyhow::bail!("no prekey available to upload"); } let pre_key_pairs = { - let mut pairs: Vec<(u32, PublicKey)> = Vec::with_capacity(rows.len()); - for (id, record) in &rows { - // Pull the public key straight out of the encoded record (field 2) - // rather than a full prost decode of the PreKeyRecordStructure: the - // upload only needs the public key, while a full decode also copies - // the private key into its own Vec. At the default batch of 812 keys - // that is ~2 throwaway allocations per record on the connect path. - // The extractor validates the record framing, so a malformed record - // returns None and is skipped here — same rejection the consume path - // (`get_pre_key`'s full decode) makes, so we never upload a key this - // device couldn't later decode with. - let public_key = match wacore::prekeys::extract_prekey_public_key(record) { - Some(raw) => { - PublicKey::from_djb_public_key_bytes(raw).map_err(anyhow::Error::from) - } - None => Err(anyhow::anyhow!("record missing or malformed public key")), - }; + let mut pairs: Vec<(u32, PublicKey)> = + Vec::with_capacity(leftover_rows.len() + fresh_pre_keys.len()); + // Leftover keys live only in the store, so decode them in full — the same + // `PreKeyRecordStructure::decode` the consume path runs, so a record + // accepted here is one this device can later decrypt with. Fresh keys skip + // decode entirely; their public keys never left memory. + use prost::Message; + for (id, record) in &leftover_rows { + let public_key = waproto::whatsapp::PreKeyRecordStructure::decode(&record[..]) + .map_err(anyhow::Error::from) + .and_then(|s| { + let raw = s + .public_key + .ok_or_else(|| anyhow::anyhow!("record missing public key"))?; + PublicKey::from_djb_public_key_bytes(&raw).map_err(anyhow::Error::from) + }); match public_key { Ok(public_key) => pairs.push((*id, public_key)), Err(e) => log::warn!("skipping undecodable prekey record {id}: {e:?}"), } } + // Fresh ids exceed every leftover id and were generated in ascending + // order, so appending keeps `pairs` sorted (last_id reads the tail). + pairs.extend(fresh_pre_keys); if pairs.is_empty() { anyhow::bail!("no decodable prekey available to upload"); } diff --git a/wacore/src/prekeys.rs b/wacore/src/prekeys.rs index 3eba6e46b..b29bf59ed 100644 --- a/wacore/src/prekeys.rs +++ b/wacore/src/prekeys.rs @@ -443,6 +443,22 @@ mod tests { let truncated = &record[..record.len() - 1]; assert!(waproto::whatsapp::PreKeyRecordStructure::decode(truncated).is_err()); assert_eq!(extract_prekey_public_key(truncated), None); + + // A malformed varint in a trailing field (10 continuation bytes that never + // terminate): a full decode errors on the bad varint, and the extractor + // agrees even though a valid publicKey precedes it. + let mut bad_varint = record.clone(); + bad_varint.push(0x08); // field 1, wire type 0 (varint) + bad_varint.extend_from_slice(&[0xFF; 10]); + assert!(waproto::whatsapp::PreKeyRecordStructure::decode(&bad_varint[..]).is_err()); + assert_eq!(extract_prekey_public_key(&bad_varint), None); + + // An unsupported wire type (3 = start group): prost rejects the dangling + // group, and the extractor rejects every wire type it cannot frame. + let mut bad_wire = record.clone(); + bad_wire.push(0x0B); // field 1, wire type 3 (start group) + assert!(waproto::whatsapp::PreKeyRecordStructure::decode(&bad_wire[..]).is_err()); + assert_eq!(extract_prekey_public_key(&bad_wire), None); } fn create_mock_bundle(device_id: u32) -> PreKeyBundle {