perf(prekeys): avoid full record decode on the pre-key upload path - #900
Conversation
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).
📝 WalkthroughWalkthroughExtract_prekey_public_key is tightened to validate protobuf framing end-to-end: it now rejects truncated length-delimited/fixed fields and unsupported wire types (returning None), matching full ChangesPre-key Extraction Validation Alignment and Upload Refactoring
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Look, this is non-negotiable. We can't have the upload path accepting pre-key records that the consume path would later reject—that's a silent failure at scale, and it breaks the entire security model. The validation tightening is the linchpin here. I need to see that the extractor truly rejects everything malformed the same way full decode does, that the test covers truncation and wire-type edge cases, and that there's no path where we upload a key we wouldn't have accepted on the other end. The refactoring itself is solid—pulling fresh keys during generation instead of reading them back from storage is the right move—but it only works if validation is airtight. If the validation isn't bulletproof, we don't ship this. Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b51b20373e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
📦 Binary size report
.text per crate
Top movers (cargo-bloat attribution)
Baseline: |
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
1 issue found across 1 file
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
Merging this PR will not alter performance
|
| Mode | Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|---|
| ⚡ | Simulation | send_message[20] |
5.6 ms | 5 ms | +10.95% |
| ⚡ | Simulation | send_and_receive[1] |
871.5 µs | 788.4 µs | +10.53% |
| 👁 | Simulation | send_message[1] |
651.1 µs | 847 µs | -23.13% |
Tip
Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.
Comparing perf/prekey-upload-skip-full-decode (5f69f98) with main (931a5d6)
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8e437100a7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@wacore/src/prekeys.rs`:
- Around line 422-446: The test
extract_prekey_public_key_matches_full_decode_validation currently only covers
truncated trailing field scenarios but does not test the edge cases that the
extractor claims to reject: malformed varints and unsupported wire types. Add
additional test cases within this test function to construct and verify
rejection of data containing malformed varints (e.g., 10+ continuation bytes)
and unsupported wire types (wire type 3 or 4). For each edge case, assert that
extract_prekey_public_key returns None and that the full prost decode also
rejects the data, ensuring consistent behavior between the two paths.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 443dc581-3ea2-4729-b605-140fa1b4b5de
📒 Files selected for processing (2)
src/prekeys.rswacore/src/prekeys.rs
…y leftovers 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.
2 participants
What
A profile-guided allocation trim on the connect/registration path.
upload_pre_keys_passwas running a fullPreKeyRecordStructure::decode(prost) on every stored record just to read the public key for the upload IQ — and a full decode also allocates a throwaway copy of the private key the upload never touches. On a fresh registration that means 812 records decoded straight back out of the store, immediately after being written to it.Now the pass carries the public keys straight out of generation: it already mints the keypairs in its blocking offload, so it keeps their public keys in hand and uploads those directly — the freshly generated keys are never written-then-reloaded-and-decoded. Only the rare leftover (already-stored, not-yet-uploaded) window keys are read back, and those go through the full
PreKeyRecordStructure::decode— the exact same call the consume path (Device::load_prekey) uses — with any record that fails to decode skipped (warn), not uploaded.Why
Profiling
connect_to_readyon CodSpeed (the run from #898) shows connect is dominated by generating + uploading the one-time pre-key batch on a fresh registration. Two parts:DEFAULT_WANTED_PRE_KEY_COUNT = 812, mirroringWAWebUploadPreKeysJob). Left untouched — real, required crypto, not waste.Vecallocations per record × 812 ≈ 1,600 fewer allocations, plus 812 fewer store reads).Safety / behavior
The upload and consume paths now parse stored records with the same
PreKeyRecordStructure::decode, so they can no longer disagree about which records are valid — a record that fails to decode is skipped on upload exactly asDevice::load_prekeywould reject it on consume.An earlier revision of this PR read the public key with a hand-rolled byte scan (
extract_prekey_public_key); that helper is now off the upload path entirely. It survives only on the digestKey path (validate_digest_key, a local hash compare that skips on mismatch and never uploads), where this PR also tightens its framing validation (rejects truncated length-delimited/fixed fields and unsupported wire types) and adds malformed-varint + wire-type 3/4 test coverage.Ordering is preserved: leftovers (sorted) precede the fresh keys, and the plan guarantees every fresh id is greater than every leftover id, so the assembled batch stays globally sorted and
last_idstays correct.Verification
cargo fmt --all/cargo clippy --all --tests— cleancargo test -p whatsapp-rust --lib prekeys— 16 tests pass, including thewindow_teststhat driveupload_pre_keys_passagainst a mock backend (upload-window, retry-single-key, collapse/regenerate paths)cargo test -p wacore prekeys— extractor framing tests, including the malformed-record / full-decode-parity casesconnect_to_readyMemory delta on this PR.Scope / follow-ups (intentionally not here)