perf: optimize libsignal reducing allocations - #225
Conversation
📝 WalkthroughWalkthroughReplaces dynamic heap allocations with fixed-size arrays and stack buffers, introduces combined-step key-derivation APIs (step_with_message_key/s), tightens chain_key length validation for protobuf/Go inputs, migrates protobuf binary fields to prost::bytes::Bytes with serde skip/default, and updates many call sites and signatures to use fixed-size types and reusable buffers. Changes
Sequence Diagram(s)(Skipped — changes are cross-cutting API and data-layout updates rather than a new multi-component runtime control flow that benefits from a sequence diagram.) Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
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 |
Bencher Report![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
wacore/src/handshake/state.rs (1)
21-21: Critical bug:unwrap_err()will panic on successful RNG initialization.
OsRng.unwrap_err()callsTryRngCore::unwrap_err()which panics if the RNG succeeds (the expected case). This should beOsRng.unwrap()or handle the error properly.🐛 Proposed fix
- let ephemeral_kp = KeyPair::generate(&mut OsRng.unwrap_err()); + let ephemeral_kp = KeyPair::generate(&mut OsRng);Note: If
rand 0.9requires explicit error handling, use:let ephemeral_kp = KeyPair::generate(&mut OsRng.unwrap());
🤖 Fix all issues with AI agents
In @wacore/libsignal/src/crypto/aes_cbc.rs:
- Around line 78-81: The current code clears the buffer then calls resize +
copy_from_slice which writes zeros then overwrites them; instead revert to the
reserve + extend_from_slice pattern: after clearing output, call
output.reserve(ctext.len()) and then output.extend_from_slice(ctext) (operating
on the same variables `output` and `ctext` in aes_cbc.rs) to ensure the data is
written only once and follows idiomatic Rust buffer growth.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (18)
src/message.rssrc/store/signal.rsstorages/sqlite-storage/src/sqlite_store.rswacore/libsignal/src/core/curve.rswacore/libsignal/src/crypto/aes_cbc.rswacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/protocol/identity_key.rswacore/libsignal/src/protocol/protocol.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/sender_keys.rswacore/libsignal/src/protocol/session_cipher.rswacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/state/signed_prekey.rswacore/libsignal/src/store/record_helpers.rswacore/src/handshake/state.rswacore/src/send.rswacore/src/store/device.rs
🧰 Additional context used
📓 Path-based instructions (2)
**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
**/*.rs: Never modify Device state directly; useDeviceCommand+PersistenceManager::process_command()for state changes
For read-only Device state access, usePersistenceManager::get_device_snapshot()
All blocking I/O (such asureqcalls) and heavy CPU-bound tasks (like media encryption) MUST be wrapped intokio::task::spawn_blockingto avoid stalling the async runtime
UseClient::chat_locksto serialize per-chat operations in asynchronous code
Usethiserrorfor custom domain-specific errors (e.g.,SocketError) andanyhow::Errorfor functions with multiple failure modes
Avoid.unwrap()and.expect()outside of tests and unrecoverable logic paths
Use theDownloadabletrait inwacore/src/download.rsfor implementing generic media download interfaces across message types
Always refreshMediaConnif it's expired before using it to get current media servers and auth tokens
Runcargo fmtbefore finalizing a feature or fix
Runcargo clippy --all-targetsbefore finalizing a feature or fix
Runcargo test --allbefore finalizing a feature or fix
Files:
src/message.rswacore/src/send.rswacore/libsignal/src/crypto/aes_cbc.rswacore/libsignal/src/protocol/identity_key.rswacore/libsignal/src/protocol/session_cipher.rssrc/store/signal.rsstorages/sqlite-storage/src/sqlite_store.rswacore/src/store/device.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/sender_keys.rswacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/store/record_helpers.rswacore/libsignal/src/core/curve.rswacore/libsignal/src/protocol/protocol.rswacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/state/signed_prekey.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/src/handshake/state.rs
wacore/**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
wacorecrate must remain platform-agnostic and have no dependencies on Tokio or specific databases
Files:
wacore/src/send.rswacore/libsignal/src/crypto/aes_cbc.rswacore/libsignal/src/protocol/identity_key.rswacore/libsignal/src/protocol/session_cipher.rswacore/src/store/device.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/sender_keys.rswacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/store/record_helpers.rswacore/libsignal/src/core/curve.rswacore/libsignal/src/protocol/protocol.rswacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/state/signed_prekey.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/src/handshake/state.rs
🧠 Learnings (1)
📚 Learning: 2025-12-25T03:36:19.700Z
Learnt from: CR
Repo: jlucaso1/whatsapp-rust PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-25T03:36:19.700Z
Learning: Applies to waproto/build.rs : `waproto` crate uses prost to compile Protocol Buffers definitions from `whatsapp.proto` in a `build.rs` script
Applied to files:
wacore/libsignal/src/store/record_helpers.rswacore/libsignal/src/protocol/protocol.rs
🧬 Code graph analysis (10)
src/message.rs (2)
wacore/libsignal/src/protocol/protocol.rs (2)
chain_key(569-571)signing_key(574-576)wacore/libsignal/src/core/curve.rs (1)
from_djb_public_key_bytes(99-106)
wacore/src/send.rs (1)
wacore/libsignal/src/protocol/protocol.rs (1)
chain_key(569-571)
wacore/libsignal/src/protocol/identity_key.rs (2)
wacore/libsignal/src/core/curve.rs (2)
serialize(109-116)serialize(229-233)wacore/src/store/device.rs (1)
serialize(16-28)
wacore/libsignal/src/protocol/session_cipher.rs (1)
wacore/libsignal/src/protocol/ratchet/keys.rs (2)
next_chain_key(139-144)message_keys(146-151)
src/store/signal.rs (2)
wacore/libsignal/src/protocol/sender_keys.rs (1)
deserialize(295-304)wacore/src/store/device.rs (1)
deserialize(30-43)
wacore/libsignal/src/protocol/ratchet.rs (4)
wacore/libsignal/src/protocol/state/session.rs (1)
root_key(159-165)wacore/libsignal/src/protocol/protocol.rs (1)
chain_key(569-571)wacore/libsignal/src/protocol/ratchet/keys.rs (1)
derive_keys(76-92)wacore/libsignal/src/protocol/ratchet/params.rs (1)
our_one_time_pre_key_pair(145-147)
wacore/libsignal/src/protocol/group_cipher.rs (2)
wacore/libsignal/src/protocol/sender_keys.rs (5)
new(39-52)new(92-97)new(173-193)sender_chain_key(207-219)sender_key_state(306-311)wacore/libsignal/src/crypto/aes_cbc.rs (1)
aes_256_cbc_decrypt_into(66-94)
wacore/libsignal/src/store/record_helpers.rs (2)
wacore/libsignal/src/protocol/state/signed_prekey.rs (1)
key_pair(129-140)wacore/libsignal/src/protocol/state/prekey.rs (1)
key_pair(59-70)
wacore/libsignal/src/core/curve.rs (3)
wacore/libsignal/src/protocol/identity_key.rs (5)
serialize(45-47)serialize(126-134)private_key(121-123)from(187-192)from(196-198)wacore/src/store/device.rs (1)
serialize(16-28)wacore/libsignal/src/core/curve/curve25519.rs (2)
from(161-164)calculate_agreement(40-48)
wacore/libsignal/src/protocol/protocol.rs (2)
wacore/libsignal/src/protocol/ratchet.rs (1)
message_version(22-24)wacore/libsignal/src/protocol/sender_keys.rs (1)
message_version(199-201)
🔇 Additional comments (34)
wacore/src/store/device.rs (1)
23-24: LGTM! Allocation reduction.The change from
into_iter()toiter().copied()correctly adapts toPrivateKey::serialize()now returning a reference to a fixed-size array (&[u8; 32]) instead of an ownedVec<u8>. This eliminates an intermediate allocation while preserving the same behavior.wacore/libsignal/src/protocol/state/signed_prekey.rs (1)
172-172: LGTM! Trait adapter.The addition of
.to_vec()correctly adaptsPrivateKey::serialize()(now returning&[u8; 32]) to theKeySerdetrait's requirement of returningVec<u8>. This maintains trait compatibility while benefiting from the fixed-size array representation in the core API.storages/sqlite-storage/src/sqlite_store.rs (2)
149-154: LGTM! Simplified dereference.The removal of the
&operator correctly adapts toPrivateKey::serialize()now returning&[u8; 32]instead of an owned array. Sinceserialize()returns a reference, the explicit&is no longer needed forextend_from_slice.
273-290: LGTM! Consistent dereference simplification.The removal of
&operators beforeserialize()calls fornoise_key,identity_key, andsigned_pre_keyis consistent with the change inserialize_keypair. All correctly adapt toPrivateKey::serialize()returning&[u8; 32].wacore/libsignal/src/protocol/ratchet/keys.rs (6)
10-11: LGTM! Required imports.The HMAC and SHA256 imports support the new
step_with_message_keysmethod for inline key derivation.
17-17: Excellent optimization! Fixed-size seed eliminates allocation.The change from
Vec<u8>to[u8; 32]for the seed eliminates heap allocation and provides compile-time size guarantees. The updatednew_from_seedsignature enforces that seeds are exactly 32 bytes, improving type safety. The copy operation is cheap for 32 bytes.Also applies to: 21-24
67-67: LGTM! Copy derives enable cheaper duplication.Adding
CopytoMessageKeysandChainKeyis appropriate since all fields are fixed-size arrays and primitives. This allows the compiler to use simple memory copies instead of clone logic, aligning with the PR's allocation reduction goals.Also applies to: 115-115
153-172: Excellent optimization! Combined step reuses HMAC setup.The
step_with_message_keysmethod efficiently computes both message keys and the next chain key in a single pass by reusing the HMAC instance withfinalize_reset(). This avoids duplicate key setup overhead while maintaining correctness. The tests verify equivalence with separate operations.
329-346: LGTM! Test adapted to fixed-size seed.The test correctly updates the seed from
Vec<u8>to[u8; 32]to match the newnew_from_seedsignature. The test logic remains equivalent.
368-434: Excellent test coverage for the optimization!The equivalence tests thoroughly verify that
step_with_message_keys()produces identical results to callingmessage_keys()andnext_chain_key()separately. The single-step test verifies immediate equivalence, while the chain test verifies sustained correctness over 10 iterations, catching potential accumulation errors. This provides strong confidence in the optimization's correctness.wacore/libsignal/src/protocol/identity_key.rs (2)
43-47: LGTM! Fixed-size serialization eliminates allocation.The change from
Box<[u8]>to[u8; 33]eliminates heap allocation by returning a fixed-size array on the stack. The documentation accurately describes the format (1 type byte + 32 key bytes). This aligns with the PR's optimization goals.
137-150: LGTM! Fixed-size signature eliminates allocation.The change from
Result<Box<[u8]>>toResult<[u8; 64]>is appropriate since Ed25519 signatures are always 64 bytes. This eliminates heap allocation and improves type safety by encoding the size in the type system.wacore/src/send.rs (1)
856-864: Chain key passed as fixed-size array matches new SKDM APISwitching from
chain_key.seed().to_vec()to*chain_key.seed()correctly aligns withSenderKeyDistributionMessage::new’s[u8; 32]parameter and avoids an extra allocation, without changing semantics.wacore/libsignal/src/protocol/state/session.rs (1)
77-101: Identity/public key serialization updates are consistent and safeUsing
.serialize().to_vec()for local/remote identity andalice_base_keymatches the updated serialization API (borrowed slice → ownedVec<u8>) and keeps encode/decode behavior unchanged.Also applies to: 128-139
src/message.rs (1)
1218-1280: Stricter Go SKDM chain_key validation is correct and improves robustnessConverting
chain_keyto[u8; 32]with an explicit length check and bailing out on mismatch is the right way to feed the new SKDM constructor and prevents malformed Go SKDMs from triggering internal invariants. Logging the actual length and sender JID is also useful for diagnosing interoperability issues.src/store/signal.rs (1)
191-198: Private key (de)serialization updates align with new APIUsing
PrivateKey::deserialize(private_key.serialize())at both identity and signed-prekey call sites matches the updatedserialize() -> &[u8]API and avoids unnecessary copying, while preserving the original logic.Also applies to: 336-344
wacore/libsignal/src/store/record_helpers.rs (1)
8-14: PreKey/SignedPreKey helpers correctly adapt to slice-based serializationChanging
private_keyfields to usekey_pair.private_key.serialize().to_vec()(and the correspondingrecord.key_pair()?.private_key.serialize().to_vec()) is the right adjustment for the new slice-based serialization and keeps the protobuf structures consistent with the decoding code paths.Also applies to: 16-26, 59-66
wacore/libsignal/src/protocol/ratchet.rs (1)
49-82: Fixed-size secrets buffer preserves X3DH semantics and removes allocationsThe switch from a growable
Vec<u8>to a[u8; 160]+secrets_lenin bothinitialize_alice_sessionandinitialize_bob_sessionkeeps the exact same concatenation order of:
- 32 discontinuity bytes
- Required DH agreements
- Optional one-time prekey agreement
and passes only the populated prefix into
derive_keys. This removes heap allocations on the handshake path without changing HKDF input, so root/chain key derivation should remain binary-compatible.Please re-run your existing cross-implementation/interop tests (e.g., against a reference libsignal implementation) to confirm that sessions established before/after this change derive identical root and chain keys for the same inputs.
Also applies to: 104-142
wacore/libsignal/src/protocol/group_cipher.rs (2)
38-41: Sender chain stepping and SKDM chain_key handling look correct
- Using
step_with_message_key()in bothgroup_encryptandget_sender_keycorrectly combines “derive message key + advance chain” and matches the previous behavior (intermediate iterations cached viaadd_sender_message_key, final iteration returned, state advanced once).- Forward-jump checks and duplicate handling remain untouched and correct.
create_sender_key_distribution_messagenow passes*sender_chain_key.seed()(a[u8; 32]copy), which matches the updated SKDM constructor and removes aVecallocation.Overall, the sender-key side of the group cipher maintains its semantics while tightening types and allocations.
Also applies to: 69-103, 111-153, 321-326
207-237: Decryption buffer reuse mirrors encryption path without changing error semanticsThe new
DECRYPTION_BUFFERthread-local reuses the sameEncryptionBufferwrapper for plaintext, avoiding per-call heap allocations. Error handling still cleanly distinguishes:
BadKeyOrIv→InvalidSenderKeySessionwith detailed logging, andBadCiphertext→InvalidMessage(CiphertextMessageType::SenderKey, "decryption failed").The
mem::take+reserve(INITIAL_CAPACITY)pattern restores buffer capacity for subsequent calls and is safe given the non-reentrant, synchronous use of the thread-local.wacore/src/handshake/state.rs (1)
59-61: Formatting improvements look good.The single-line calls to
mix_shared_secretare cleaner and the serialization approach aligns with the fixed-size array refactor incurve.rs.Also applies to: 72-77, 93-95
wacore/libsignal/src/protocol/session_cipher.rs (3)
99-111: Buffer restoration afterstd::mem::take()is fragile.If an error occurs between
std::mem::take(buf)(line 107) andbuf.reserve()(line 109), the buffer capacity isn't restored. However, sincereserveis called unconditionally aftertake, and theOkwrapping happens after, this is correct. The current flow is: take → reserve → return result.One minor observation:
std::mem::takefollowed byreserveworks, but consider usingstd::mem::replacewith a pre-allocated Vec to avoid the zero-capacity intermediate state, though this is optional.
81-82: Two-step key generation pattern is correctly implemented.The new
step_with_message_keys()pattern correctly separates the chain advancement from key generation, reducing redundant HMAC operations.
821-831: Chain key iteration refactored correctly.The dereference
*chain_keyworks becauseChainKeyimplementsCopy. The loop correctly advances the chain while storing skipped message keys, and the final step returns the message key generator for the target counter.wacore/libsignal/src/protocol/sender_keys.rs (3)
125-151: Combined key derivation is correctly implemented.The
step_with_message_keymethod correctly:
- Validates iteration overflow before proceeding
- Reuses the HMAC context via
finalize_reset()to derive both keys with the same chain key- Uses
self.iterationfor the message key (current) andnew_iterationfor the next chain keyThe tests at lines 762-820 confirm equivalence with the separate
sender_message_key()andnext()calls.
82-86: AddingCopytoSenderChainKeyis appropriate.The struct contains only
Copytypes (u32and[u8; 32]), and at 36 bytes total, it's reasonable to copy by value rather than requiring explicit clones.
103-105: Return type refinement to&[u8; 32]provides compile-time size guarantees.wacore/libsignal/src/protocol/protocol.rs (3)
518-518: Fixed-sizechain_keytype enforces correctness at compile time.Changing from
Vec<u8>to[u8; 32]provides stronger type guarantees. The parsing at lines 628-634 properly validates the length and converts the Vec.Also applies to: 528-528, 569-571
87-87: Serialization adapts correctly to fixed-size array returns.Using
to_vec()on the fixed-size arrays fromserialize()is the correct approach for populating protobufVec<u8>fields.Also applies to: 262-263
403-417: No signing semantic change—this is a performance optimization only.The current implementation signs
[version_byte || proto]exactly as the previous version did. Commit 5f11658 explicitly created a separatedata_to_signbuffer with version + proto, then signed it. The current code (commit 8d20e01) achieves the same by encoding the proto directly into the finalserializedbuffer and signing that. Both approaches sign identical data, and the previous implementation's comment confirms this format is correct for interoperability with other clients (baileys, libsignal-go). Theverify_signaturemethod correctly extracts and verifies&self.serialized[..len - SIGNATURE_LEN], which is the signed portion. The change eliminates intermediate allocations without altering the signing semantics.wacore/libsignal/src/core/curve.rs (4)
108-116: Stack-allocated fixed-size array eliminates heap allocation.Returning
[u8; 33]instead of a boxed slice removes a heap allocation for each serialization. The implementation correctly places the type byte at index 0 and copies the 32-byte key data to the remaining positions.
229-233: Zero-copy serialization by returning reference to internal array.Returning
&[u8; 32]directly references the internal key storage, avoiding any copying. Callers needing an owned copy can dereference or callto_vec().
251-279: Fixed-size return types for signatures and key agreements.The return type changes to
[u8; 64]for signatures and[u8; 32]for key agreements match the cryptographic output sizes (Ed25519 signatures are 64 bytes, X25519 shared secrets are 32 bytes). This aligns with the underlyingcurve25519::PrivateKeyimplementation.
332-342:KeyPairwrapper methods correctly propagate fixed-size return types.
8d20e01 to
ca8b226
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
src/store/signal.rs (2)
335-353: Similar serialize/deserialize round-trip - verify necessity.Line 342 performs the same pattern:
PrivateKey::deserialize(self.signed_pre_key.private_key.serialize()). This creates an unnecessary copy by serializing to bytes and immediately deserializing back to aPrivateKey.For a PR focused on reducing allocations, consider whether
self.signed_pre_key.private_keycan be used more directly to construct the neededKeyPair.
194-197: Verify the serialize/deserialize round-trip is necessary.The code calls
serialize()onprivate_key_bytesand immediately deserializes it back toPrivateKey. This round-trip seems counterintuitive for a performance optimization PR focused on reducing allocations.If
self.identity_key.private_keyis already aPrivateKeytype compatible with the protocol, consider using it directly or finding a more efficient conversion path that avoids the serialization overhead.#!/bin/bash # Verify the type of self.identity_key.private_key and whether direct usage is possible ast-grep --pattern $'struct Device { $$$ identity_key: $TYPE, $$$ }' # Also check what identity_key's private_key field type is rg -n "identity_key.*private_key" --type rust -C 3wacore/libsignal/src/protocol/ratchet.rs (2)
49-82: Fixed-size secrets buffer in Alice init is bounded and equivalent to previous Vec logicThe 160-byte stack buffer with
secrets_lentracking correctly holds the discontinuity bytes plus up to four 32-byte DH outputs;derive_keys(&secrets[..secrets_len])preserves the previous HKDF input ordering with fewer allocations. Consider adding a debug assertion around each write (debug_assert!(secrets_len + 32 <= secrets.len())) to future‑proof this if additional agreements are ever appended.
104-142: Bob init’s secrets accumulation mirrors Alice and remains within boundsThe Bob side uses the same 160-byte buffer pattern; with 1 discontinuity + 3 mandatory DHs + optional one-time prekey,
secrets_lencannot exceed 160, and HKDF input semantics remain unchanged. A small defensive improvement would be the samedebug_assert!(secrets_len + 32 <= secrets.len())before each copy to guard against future changes adding more agreements.wacore/libsignal/src/protocol/protocol.rs (1)
559-576: Simplify getters: Result wrapper is now unnecessary.All getter methods (
chain_id(),iteration(),chain_key(),signing_key()) returnResult<T>but are now infallible since the data is validated during construction inTryFrom. These could be simplified to return the values directly.♻️ Proposed simplification
#[inline] -pub fn chain_id(&self) -> Result<u32> { - Ok(self.chain_id) +pub fn chain_id(&self) -> u32 { + self.chain_id } #[inline] -pub fn iteration(&self) -> Result<u32> { - Ok(self.iteration) +pub fn iteration(&self) -> u32 { + self.iteration } #[inline] -pub fn chain_key(&self) -> Result<&[u8; 32]> { - Ok(&self.chain_key) +pub fn chain_key(&self) -> &[u8; 32] { + &self.chain_key } #[inline] -pub fn signing_key(&self) -> Result<&PublicKey> { - Ok(&self.signing_key) +pub fn signing_key(&self) -> &PublicKey { + &self.signing_key }Note: This is a breaking API change, so apply only if acceptable for your use case.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (17)
src/message.rssrc/store/signal.rsstorages/sqlite-storage/src/sqlite_store.rswacore/libsignal/src/core/curve.rswacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/protocol/identity_key.rswacore/libsignal/src/protocol/protocol.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/sender_keys.rswacore/libsignal/src/protocol/session_cipher.rswacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/state/signed_prekey.rswacore/libsignal/src/store/record_helpers.rswacore/src/handshake/state.rswacore/src/send.rswacore/src/store/device.rs
🚧 Files skipped from review as they are similar to previous changes (4)
- wacore/src/send.rs
- wacore/src/store/device.rs
- wacore/libsignal/src/core/curve.rs
- storages/sqlite-storage/src/sqlite_store.rs
🧰 Additional context used
📓 Path-based instructions (2)
**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
**/*.rs: Never modify Device state directly; useDeviceCommand+PersistenceManager::process_command()for state changes
For read-only Device state access, usePersistenceManager::get_device_snapshot()
All blocking I/O (such asureqcalls) and heavy CPU-bound tasks (like media encryption) MUST be wrapped intokio::task::spawn_blockingto avoid stalling the async runtime
UseClient::chat_locksto serialize per-chat operations in asynchronous code
Usethiserrorfor custom domain-specific errors (e.g.,SocketError) andanyhow::Errorfor functions with multiple failure modes
Avoid.unwrap()and.expect()outside of tests and unrecoverable logic paths
Use theDownloadabletrait inwacore/src/download.rsfor implementing generic media download interfaces across message types
Always refreshMediaConnif it's expired before using it to get current media servers and auth tokens
Runcargo fmtbefore finalizing a feature or fix
Runcargo clippy --all-targetsbefore finalizing a feature or fix
Runcargo test --allbefore finalizing a feature or fix
Files:
src/store/signal.rswacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/session_cipher.rswacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/protocol/state/signed_prekey.rssrc/message.rswacore/libsignal/src/protocol/identity_key.rswacore/libsignal/src/protocol/sender_keys.rswacore/src/handshake/state.rswacore/libsignal/src/store/record_helpers.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/protocol.rs
wacore/**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
wacorecrate must remain platform-agnostic and have no dependencies on Tokio or specific databases
Files:
wacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/session_cipher.rswacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/protocol/state/signed_prekey.rswacore/libsignal/src/protocol/identity_key.rswacore/libsignal/src/protocol/sender_keys.rswacore/src/handshake/state.rswacore/libsignal/src/store/record_helpers.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/protocol.rs
🧠 Learnings (1)
📚 Learning: 2025-12-25T03:36:19.700Z
Learnt from: CR
Repo: jlucaso1/whatsapp-rust PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-25T03:36:19.700Z
Learning: Applies to waproto/build.rs : `waproto` crate uses prost to compile Protocol Buffers definitions from `whatsapp.proto` in a `build.rs` script
Applied to files:
wacore/libsignal/src/protocol/protocol.rs
🧬 Code graph analysis (7)
src/store/signal.rs (4)
wacore/libsignal/src/protocol/identity_key.rs (1)
private_key(121-123)wacore/libsignal/src/protocol/sender_keys.rs (1)
deserialize(295-304)wacore/src/store/device.rs (1)
deserialize(30-43)wacore/libsignal/src/core/curve/curve25519.rs (1)
private_key_bytes(155-157)
wacore/libsignal/src/protocol/ratchet.rs (4)
wacore/libsignal/src/protocol/state/session.rs (1)
root_key(159-165)wacore/libsignal/src/protocol/protocol.rs (1)
chain_key(569-571)wacore/libsignal/src/protocol/ratchet/keys.rs (1)
derive_keys(76-92)wacore/libsignal/src/protocol/ratchet/params.rs (1)
our_one_time_pre_key_pair(145-147)
wacore/libsignal/src/protocol/session_cipher.rs (3)
wacore/libsignal/src/protocol/ratchet/keys.rs (2)
next_chain_key(139-144)message_keys(146-151)wacore/libsignal/src/protocol/protocol.rs (1)
chain_key(569-571)wacore/libsignal/src/protocol/state/session.rs (1)
session_state(635-637)
wacore/libsignal/src/protocol/group_cipher.rs (3)
wacore/libsignal/src/protocol/protocol.rs (4)
new(76-113)new(249-281)new(389-426)new(524-551)wacore/libsignal/src/protocol/sender_keys.rs (4)
new(39-52)new(92-97)new(173-193)sender_chain_key(207-219)wacore/libsignal/src/crypto/aes_cbc.rs (1)
aes_256_cbc_decrypt_into(66-93)
wacore/libsignal/src/protocol/identity_key.rs (1)
wacore/libsignal/src/core/curve.rs (2)
serialize(109-116)serialize(229-233)
wacore/libsignal/src/store/record_helpers.rs (2)
wacore/libsignal/src/protocol/state/signed_prekey.rs (1)
key_pair(129-140)wacore/libsignal/src/protocol/state/prekey.rs (1)
key_pair(59-70)
wacore/libsignal/src/protocol/protocol.rs (2)
wacore/libsignal/src/protocol/ratchet.rs (1)
message_version(22-24)wacore/libsignal/src/protocol/sender_keys.rs (4)
message_version(199-201)chain_id(203-205)iteration(62-64)iteration(99-101)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: benchmark
- GitHub Check: Build & Test
🔇 Additional comments (25)
wacore/src/handshake/state.rs (3)
73-77: Consistent refactoring applied.Same allocation optimization pattern applied here for ephemeral key serialization.
93-95: Consistent refactoring applied.Same allocation optimization pattern applied here for static key serialization.
60-61: Good refactoring to reduce allocations.The direct passing of
serialize()tomix_shared_secretconfirms the method now returns&[u8; 32]instead of an owned vector, eliminating intermediate storage and heap allocation.wacore/libsignal/src/protocol/identity_key.rs (2)
43-47: LGTM! Excellent optimization.Changing from
Box<[u8]>to[u8; 33]eliminates a heap allocation for a fixed-size cryptographic key. This is a clear performance win with no downsides, as the identity key serialization is always 33 bytes (1 type byte + 32 key bytes).
137-150: LGTM! Signature output now uses fixed-size array.Changing the return type from
Box<[u8]>to[u8; 64]is appropriate since Ed25519 signatures are always 64 bytes. This avoids unnecessary heap allocation.wacore/libsignal/src/protocol/state/session.rs (1)
84-100: LGTM! Correctly adapted to fixed-size array serialization.The change from
.into_vec()to.to_vec()is necessary because the upstreamserialize()methods now return fixed-size arrays[u8; 33]instead ofBox<[u8]>.The
.to_vec()call copies the fixed-size array into aVec<u8>for protobuf storage. While this introduces a copy, it's negligible (33 bytes) and the overall change is still a performance win by eliminating heap allocations in the serialization layer.wacore/libsignal/src/protocol/state/signed_prekey.rs (1)
170-178: LGTM! Correctly adapts to upstream serialization changes.The
KeySerdetrait requiresserialize()to returnVec<u8>, so calling.to_vec()on the fixed-size array returned byPrivateKey::serialize()is the appropriate way to satisfy the trait contract while preserving the optimization benefits upstream.wacore/libsignal/src/store/record_helpers.rs (1)
8-67: LGTM! Consistent adaptation for protobuf storage.All four occurrences (lines 12, 25, 65, and the pattern throughout) correctly use
.serialize().to_vec()to convert the fixed-size array returned byPrivateKey::serialize()into aVec<u8>required by the protobuf structures.The copying overhead is minimal (32 bytes per key) and the overall optimization from eliminating heap allocations in the serialization layer provides a net performance benefit.
src/message.rs (1)
1238-1248: Chain key length validation for Go SKDM is correct and saferConverting
chain_keyto[u8; 32]with an explicit length check and early return cleanly enforces the invariant expected bySenderKeyDistributionMessage::newand avoids ever constructing SKDMs with malformed key material. Logging only the length and sender JID is appropriate here.wacore/libsignal/src/protocol/session_cipher.rs (2)
81-83: Sender chain advancement during encryption stays semantically identicalUsing
step_with_message_keys()to derive the per-messageMessageKeysand thenext_chain_key, then persistingnext_chain_keywithset_sender_chain_key, preserves the old “message uses current index, chain advances by one” behavior while avoiding a redundant derivation. No issues spotted.Also applies to: 157-157
821-832: Receiver chain progression and message key caching are consistent with Double RatchetCopying the input
chain_key, then looping withstep_with_message_keys()to cache intermediate message keys and finally updating the receiver chain key with the lastnext_chain, matches the standard get-or-create pattern from libsignal (final key forcounter, chain advanced tocounter+1). This keeps duplicate detection and forward‑jump limits intact while reducing redundant HMACs.wacore/libsignal/src/protocol/group_cipher.rs (3)
38-41: Thread-local decryption buffer mirrors encryption buffer and avoids extra allocationsIntroducing
DECRYPTION_BUFFERand decrypting into a reusableVec<u8>(withtake()+reserve()reset) matches the existing encryption buffer pattern and removes per-call allocations ingroup_decrypt. Error mapping forBadKeyOrIvvsBadCiphertextis preserved; the behavior on corrupt sender key state vs bad ciphertext remains identical.Also applies to: 207-237
73-73: Sender chain/message key progression in groups correctly uses step_with_message_keyBoth
group_encryptandget_sender_keynow usestep_with_message_key()to atomically derive theSenderMessageKeyand advance theSenderChainKey, storing intermediate message keys and updating the chain only at the end. This matches prior semantics (iteration and stored keys) while avoiding duplicate HMACs.Also applies to: 102-103, 145-153
321-326: Passing SKDM chain_key as a fixed array is aligned with the new APIUsing
*sender_chain_key.seed()to supply a[u8; 32]toSenderKeyDistributionMessage::newremoves an unnecessary allocation and matches the updated SKDM constructor andchain_key()accessor signatures. No behavioral change here, just tighter typing.wacore/libsignal/src/protocol/sender_keys.rs (2)
82-106: SenderChainKey Copy + step_with_message_key API look correct and consistentMaking
SenderChainKeyCopyand changingseed()to return&[u8; 32]tighten the type contracts without changing behavior. The newstep_with_message_key()correctly guards against iteration overflow, derives the same message/chain seeds as the previoussender_message_key()+next()sequence, and provides a more efficient single-step primitive for callers.Also applies to: 125-151
563-569: Tests adequately verify step_with_message_key equivalence and chain consistencyThe new tests exercise both single-step and multi-step usage of
step_with_message_key(), comparing message keys and chain seeds against the existingsender_message_key()/next()sequence over multiple iterations, which strongly validates the new API. Using Copy semantics inset_sender_chain_key(next_sck)in tests is idiomatic and matches the updated type.Also applies to: 759-820
wacore/libsignal/src/protocol/ratchet/keys.rs (5)
10-11: LGTM: Fixed-size seed reduces allocations.The change from
Vec<u8>to[u8; 32]for the seed eliminates heap allocations and enforces compile-time size validation. The addition of HMAC imports supports the new optimized derivation method.Also applies to: 17-17
22-23: LGTM: Stricter signature enforces correct seed size.Tightening the signature from
&[u8]to&[u8; 32]enforces the seed size at compile time and eliminates the need for runtime length checks. The copy operation is efficient for a 32-byte array.
67-67: LGTM: Copy trait enables efficient pass-by-value.Both
MessageKeysandChainKeyare small structs containing only fixed-size arrays and integers, making them ideal candidates forCopy. This enables efficient pass-by-value semantics and aligns with the allocation-reduction goal.Also applies to: 115-115
153-172: LGTM: Combined derivation reduces HMAC overhead.The
step_with_message_keysmethod is a well-designed optimization that reuses a single HMAC key setup to derive both the message key seed and next chain key. Usingfinalize_reset()allows the HMAC instance to be reused efficiently. The implementation correctly applies the different seeds (MESSAGE_KEY_SEEDandCHAIN_KEY_SEED) to produce distinct outputs.
368-434: LGTM: Comprehensive test coverage for the optimization.The tests thoroughly verify that
step_with_message_keysproduces identical results to callingmessage_keys()andnext_chain_key()separately, both for a single step and over multiple iterations. This ensures the optimization maintains correctness.wacore/libsignal/src/protocol/protocol.rs (4)
87-87: LGTM: Serialization updated for fixed-size array returns.The addition of
.to_vec()calls indicates thatPublicKey::serialize()andIdentityKey::serialize()now return fixed-size array references instead ofVec<u8>. This change aligns with the allocation-reduction goal while maintaining protobuf compatibility.Also applies to: 262-263
400-425: LGTM: Direct buffer construction eliminates intermediate allocation.The refactored serialization builds the final buffer directly (version byte + proto + signature) instead of encoding the proto separately and then combining. This reduces allocations and is more efficient. The signing correctly covers the version byte and proto data before appending the signature.
518-518: LGTM: Fixed-size chain_key reduces allocations.Changing
chain_keyfromVec<u8>to[u8; 32]eliminates heap allocation and enforces size validation. Using.to_vec()when building the protobuf message is necessary for compatibility with the protobuf encoding.Also applies to: 528-528, 534-534
621-634: LGTM: Length validation ensures correct conversion.The
TryFromimplementation correctly validates thatchain_key_vechas exactly 32 bytes before converting it to[u8; 32]. The length check on line 628 provides an early validation, and thetry_into()conversion on lines 632-634 handles any edge cases with proper error propagation.
…icient serialization
8dddd34 to
9b6fa0a
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
wacore/libsignal/src/protocol/ratchet/keys.rs (1)
42-62: Silent fallback to zero-filled keys on malformed protobuf may mask data corruption.When parsing a
Serializedvariant with invalid field lengths (e.g., truncated cipher_key),generate_keys()silently falls back to[0u8; 32]or[0u8; 16]. Sincefrom_pb()validates lengths before creatingSerialized, this branch should only trigger if internal invariants are violated. However, returning zero-filled keys could lead to silent decryption failures rather than surfacing the corruption.Consider whether panicking or returning a Result would be safer here, given the security-sensitive context.
♻️ Alternative: propagate error instead of defaulting
- Self::Serialized(pb) => { - // Parse on demand - only when keys are actually needed - MessageKeys { - cipher_key: pb - .cipher_key - .as_deref() - .and_then(|b| b.try_into().ok()) - .unwrap_or([0u8; 32]), - mac_key: pb - .mac_key - .as_deref() - .and_then(|b| b.try_into().ok()) - .unwrap_or([0u8; 32]), - iv: pb - .iv - .as_deref() - .and_then(|b| b.try_into().ok()) - .unwrap_or([0u8; 16]), - counter: pb.index.unwrap_or(0), - } - } + Self::Serialized(pb) => { + // Parse on demand - only when keys are actually needed + // from_pb() validates lengths, so these should never fail + MessageKeys { + cipher_key: pb + .cipher_key + .as_deref() + .and_then(|b| b.try_into().ok()) + .expect("from_pb validated cipher_key length"), + mac_key: pb + .mac_key + .as_deref() + .and_then(|b| b.try_into().ok()) + .expect("from_pb validated mac_key length"), + iv: pb + .iv + .as_deref() + .and_then(|b| b.try_into().ok()) + .expect("from_pb validated iv length"), + counter: pb.index.unwrap_or(0), + } + }wacore/libsignal/src/protocol/group_cipher.rs (1)
40-40: Consider renamingEncryptionBufferto a more generic name.The
EncryptionBuffertype is now reused for both encryption and decryption viaDECRYPTION_BUFFER. Consider renaming it to something likeReusableBufferorCryptoBufferto better reflect its dual purpose.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (10)
Cargo.tomlsrc/store/signal.rswacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/protocol/protocol.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/sender_keys.rswacore/libsignal/src/protocol/state/session.rswaproto/build.rswaproto/src/whatsapp.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- src/store/signal.rs
🧰 Additional context used
📓 Path-based instructions (3)
**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
**/*.rs: Never modify Device state directly; useDeviceCommand+PersistenceManager::process_command()for state changes
For read-only Device state access, usePersistenceManager::get_device_snapshot()
All blocking I/O (such asureqcalls) and heavy CPU-bound tasks (like media encryption) MUST be wrapped intokio::task::spawn_blockingto avoid stalling the async runtime
UseClient::chat_locksto serialize per-chat operations in asynchronous code
Usethiserrorfor custom domain-specific errors (e.g.,SocketError) andanyhow::Errorfor functions with multiple failure modes
Avoid.unwrap()and.expect()outside of tests and unrecoverable logic paths
Use theDownloadabletrait inwacore/src/download.rsfor implementing generic media download interfaces across message types
Always refreshMediaConnif it's expired before using it to get current media servers and auth tokens
Runcargo fmtbefore finalizing a feature or fix
Runcargo clippy --all-targetsbefore finalizing a feature or fix
Runcargo test --allbefore finalizing a feature or fix
Files:
wacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/sender_keys.rswaproto/build.rswacore/libsignal/src/protocol/protocol.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/group_cipher.rswaproto/src/whatsapp.rs
wacore/**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
wacorecrate must remain platform-agnostic and have no dependencies on Tokio or specific databases
Files:
wacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/sender_keys.rswacore/libsignal/src/protocol/protocol.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/group_cipher.rs
waproto/build.rs
📄 CodeRabbit inference engine (AGENTS.md)
waprotocrate uses prost to compile Protocol Buffers definitions fromwhatsapp.protoin abuild.rsscript
Files:
waproto/build.rs
🧠 Learnings (3)
📚 Learning: 2025-12-25T03:36:19.700Z
Learnt from: CR
Repo: jlucaso1/whatsapp-rust PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-25T03:36:19.700Z
Learning: Applies to waproto/build.rs : `waproto` crate uses prost to compile Protocol Buffers definitions from `whatsapp.proto` in a `build.rs` script
Applied to files:
waproto/build.rswacore/libsignal/src/protocol/protocol.rswaproto/src/whatsapp.rs
📚 Learning: 2025-12-25T03:36:19.700Z
Learnt from: CR
Repo: jlucaso1/whatsapp-rust PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-25T03:36:19.700Z
Learning: The project is organized into three crates: `wacore` (platform-agnostic, `no_std` core logic with no Tokio/database dependencies), `waproto` (Protocol Buffers definitions compiled with prost), and `whatsapp-rust` (main client with Tokio/Diesel integration)
Applied to files:
waproto/build.rswaproto/src/whatsapp.rs
📚 Learning: 2025-12-25T03:36:19.700Z
Learnt from: CR
Repo: jlucaso1/whatsapp-rust PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-25T03:36:19.700Z
Learning: Applies to wacore/**/*.rs : `wacore` crate must remain platform-agnostic and have no dependencies on Tokio or specific databases
Applied to files:
Cargo.toml
🧬 Code graph analysis (5)
wacore/libsignal/src/protocol/sender_keys.rs (3)
wacore/libsignal/src/protocol/ratchet/keys.rs (4)
next_chain_key(185-190)new(171-173)new(231-233)iv(151-153)wacore/libsignal/src/protocol/protocol.rs (5)
new(76-113)new(249-281)new(389-426)new(524-551)chain_key(569-571)wacore/libsignal/src/protocol/crypto.rs (1)
hmac_sha256(9-14)
wacore/libsignal/src/protocol/protocol.rs (3)
wacore/libsignal/src/protocol/identity_key.rs (1)
identity_key(109-111)wacore/libsignal/src/protocol/ratchet.rs (1)
message_version(22-24)wacore/libsignal/src/protocol/sender_keys.rs (4)
message_version(204-206)chain_id(208-210)iteration(63-65)iteration(101-103)
wacore/libsignal/src/protocol/ratchet.rs (2)
wacore/libsignal/src/protocol/ratchet/keys.rs (1)
derive_keys(122-138)wacore/libsignal/src/protocol/ratchet/params.rs (1)
our_one_time_pre_key_pair(145-147)
wacore/libsignal/src/protocol/group_cipher.rs (2)
wacore/libsignal/src/protocol/sender_keys.rs (6)
new(39-52)new(94-99)new(176-198)sender_chain_key(212-224)sender_key_state(311-316)chain_id(208-210)wacore/libsignal/src/crypto/aes_cbc.rs (1)
aes_256_cbc_decrypt_into(66-93)
waproto/src/whatsapp.rs (1)
wacore/binary/src/builder.rs (1)
bytes(40-43)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Build & Test
- GitHub Check: benchmark
🔇 Additional comments (22)
Cargo.toml (1)
33-33: LGTM!Adding the
serdefeature tobytesaligns with the broader refactor usingprost::bytes::Bytesfor cryptographic structures. This enables proper serde support forBytestypes across the codebase.wacore/libsignal/src/protocol/ratchet.rs (2)
49-85: Effective optimization replacing Vec with stack-allocated buffer.The fixed-size
[u8; 160]buffer correctly accommodates the maximum of 5 × 32-byte secrets (discontinuity bytes + up to 4 DH agreements). Thedebug_assert!statements provide safety checks during development.
108-149: Consistent implementation with Alice session.The Bob session initialization mirrors the Alice session buffer approach correctly, maintaining protocol symmetry while eliminating heap allocations.
wacore/libsignal/src/protocol/sender_keys.rs (3)
84-84: Good addition ofCopyderive.
SenderChainKeycontains onlyu32and[u8; 32], both of which areCopy. This allows more efficient passing without clone overhead.
127-153: Well-designed optimization reusing HMAC context.The
step_with_message_keymethod correctly leveragesfinalize_reset()to compute both the message key seed and next chain key with a single HMAC key setup. The semantics are correct:finalize_reset()returns the MAC and resets the HMAC to its initial keyed state, allowing the second derivation with the same base key.
763-825: Thorough test coverage for the combined-step optimization.The equivalence test and multi-iteration chain test provide strong validation that
step_with_message_keyproduces identical results to the separate method calls.wacore/libsignal/src/protocol/ratchet/keys.rs (2)
199-218: Consistent HMAC reuse optimization.The
step_with_message_keysimplementation mirrors the pattern insender_keys.rs, correctly usingfinalize_reset()to derive both keys with a single HMAC setup. ReturningMessageKeyGeneratorenables lazy derivation when keys aren't immediately needed.
161-165: LGTM!Adding
Copyderive is appropriate sinceChainKeycontains onlyCopytypes ([u8; 32]andu32).waproto/build.rs (2)
35-46: Well-documented configuration for Bytes optimization.The
config.bytes()paths correctly target frequently-accessed cryptographic structures on hot paths. Usingbytes::Bytesinstead ofVec<u8>enables O(1) cloning which is valuable for session/chain keys accessed on every message.
48-83: Correct serde skip attributes for protobuf-only fields.The
#[serde(skip, default)]attributes are necessary since these nested cryptographic fields are stored as protobuf blobs rather than JSON-serialized. This prevents serde from attempting to serialize/deserializeBytesfields that lack serde support in prost's default configuration.wacore/libsignal/src/protocol/state/session.rs (2)
285-290: LGTM!The switch to
Bytes::copy_from_slicefor chain key storage aligns with the prost configuration changes. The localuse prost::bytes::Bytesimport is repeated across methods, which is functional though a module-level import could reduce duplication.
87-96: Minor:to_vec()usage is fine here.For fixed-size arrays like
[u8; 33],to_vec()andinto_vec()are semantically equivalent. The change is correct.wacore/libsignal/src/protocol/protocol.rs (3)
621-634: Proper length validation before fixed-size conversion.The validation at line 628 correctly checks both
chain_key_vec.len() != 32andsigning_key.len() != 33before attempting the conversion. Thetry_into().map_err()pattern handles the conversion safely.
403-424: Cleaner serialization flow with correct capacity pre-allocation.The refactored buffer construction correctly pre-allocates
1 + proto_len + SIGNATURE_LENbytes and builds the message in a single pass. The signature is computed over[version || proto]before appending, which matches the expected format.
518-518: Breaking API change:chain_keyis now[u8; 32].This is a positive change that enforces correct key size at compile time. All callers have been updated correctly—both wacore callers pass
*chain_key.seed()/*sender_chain_key.seed(), and the src/message.rs fallback properly validates the conversion withtry_into()before passing the result to the constructor.wacore/libsignal/src/protocol/group_cipher.rs (5)
144-152: LGTM: Correct implementation of forward chain progression.The loop properly accumulates intermediate message keys while advancing the chain, and the final step correctly derives the target message key while updating the state with the next chain key.
207-237: LGTM: Proper buffer reuse with comprehensive error handling.The decryption buffer management correctly mirrors the encryption pattern. Error handling properly distinguishes between key/IV corruption and ciphertext issues, with appropriate logging and error mapping.
325-325: Change is correct:SenderKeyDistributionMessage::new()accepts[u8; 32]by value.The modification from
sender_chain_key.seed().to_vec()to*sender_chain_key.seed()is valid and improves performance. The function signature inprotocol.rsexplicitly acceptschain_key: [u8; 32], and dereferencing the seed reference provides the fixed-size array in the expected form, eliminating the unnecessary heap allocation.
265-268: Type consistency is correct. Thesigning_key()method returns&PublicKeyand is properly dereferenced with*to match thePublicKeyparameter type. Thechain_key()method returns&[u8; 32], which automatically coerces to the expected&[u8]parameter type through Rust's deref coercion. No changes needed.
73-73: Tests forstep_with_message_key()are comprehensive and properly cover chain key progression and message key derivation.The method has two dedicated tests:
test_step_with_message_key_equivalence()verifies the combined API produces identical results to separate calls, testing message key components (iteration, IV, cipher_key) and chain key components (seed, iteration).
test_step_with_message_key_chain()validates progression over 10 iterations, ensuring both message keys and chain keys remain synchronized and don't drift.These tests confirm the optimization is correct and the method handles chain advancement properly.
waproto/src/whatsapp.rs (2)
11026-11043: Serde skip configuration for cryptographic keys is properly applied.The
ChainKey.keyandMessageKeyfields (cipher_key,mac_key,iv) correctly use theserde(skip, default)pattern withbytes::Bytes. This configuration is intentional and well-justified inbuild.rs:
bytes::Bytesis used for O(1) cloning instead of O(n) copying on hot paths (message encryption/decryption)serde(skip, default)is necessary becausebytes::ByteslacksSerialize/Deserializeimplementations by design- The prost configuration explicitly applies these field attributes for all sensitive binary fields in both
SessionStructure.ChainandSenderKeyStateStructurestructuresThe implementation is consistent and correct.
10945-10966: No issues found. The switch fromVec<u8>toprost::bytes::Bytesfor cryptographic material is implemented correctly and does not introduce compatibility problems.The code properly handles Bytes through the
as_ref()pattern (e.g.,seed_bytes.as_ref().try_into()), and the proto structs are never passed tobuilder.bytes(), which only receivesVec<u8>from theNodeContentenum. The#[serde(skip, default)]annotations are a security improvement—they prevent accidental serialization of sensitive private key material—and do not break any existing persistence workflows, as these proto structs are used internally by libsignal rather than in device state serialization.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
wacore/libsignal/src/protocol/ratchet/keys.rs (1)
38-64: Consider panic or Result instead of silent zero fallback.In the
Serializedbranch,unwrap_orwith zero defaults (lines 49, 54, 59) could silently produce incorrect keys if the protobuf is somehow corrupted after validation. Whilefrom_pb()validates fields before creatingSerializedvariants, a defensiveexpect()with a clear error message would be safer than silent fallback to zeros.💡 More defensive error handling
Self::Serialized(pb) => { // Parse on demand - only when keys are actually needed MessageKeys { cipher_key: pb .cipher_key .as_deref() .and_then(|b| b.try_into().ok()) - .unwrap_or([0u8; 32]), + .expect("Serialized variant should have valid cipher_key"), mac_key: pb .mac_key .as_deref() .and_then(|b| b.try_into().ok()) - .unwrap_or([0u8; 32]), + .expect("Serialized variant should have valid mac_key"), iv: pb .iv .as_deref() .and_then(|b| b.try_into().ok()) - .unwrap_or([0u8; 16]), + .expect("Serialized variant should have valid iv"), counter: pb.index.unwrap_or(0), } }wacore/libsignal/src/protocol/group_cipher.rs (1)
40-40: Consider renamingEncryptionBufferto a more generic name.The
EncryptionBuffertype is now used for both encryption (line 39) and decryption (line 40), which may cause confusion. Consider renaming it to something likeCipherBufferorCryptoBufferto better reflect its dual purpose.♻️ Refactor suggestion
Rename the struct and update both thread-local declarations:
-struct EncryptionBuffer { +struct CipherBuffer { buffer: Vec<u8>, } -impl EncryptionBuffer { +impl CipherBuffer { const INITIAL_CAPACITY: usize = 1024; fn new() -> Self {Then update the thread-local declarations:
thread_local! { - static ENCRYPTION_BUFFER: RefCell<EncryptionBuffer> = RefCell::new(EncryptionBuffer::new()); - static DECRYPTION_BUFFER: RefCell<EncryptionBuffer> = RefCell::new(EncryptionBuffer::new()); + static ENCRYPTION_BUFFER: RefCell<CipherBuffer> = RefCell::new(CipherBuffer::new()); + static DECRYPTION_BUFFER: RefCell<CipherBuffer> = RefCell::new(CipherBuffer::new()); }
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (10)
Cargo.tomlsrc/store/signal.rswacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/protocol/protocol.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/sender_keys.rswacore/libsignal/src/protocol/state/session.rswaproto/build.rswaproto/src/whatsapp.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- Cargo.toml
🧰 Additional context used
📓 Path-based instructions (3)
**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
**/*.rs: Never modify Device state directly; useDeviceCommand+PersistenceManager::process_command()for state changes
For read-only Device state access, usePersistenceManager::get_device_snapshot()
All blocking I/O (such asureqcalls) and heavy CPU-bound tasks (like media encryption) MUST be wrapped intokio::task::spawn_blockingto avoid stalling the async runtime
UseClient::chat_locksto serialize per-chat operations in asynchronous code
Usethiserrorfor custom domain-specific errors (e.g.,SocketError) andanyhow::Errorfor functions with multiple failure modes
Avoid.unwrap()and.expect()outside of tests and unrecoverable logic paths
Use theDownloadabletrait inwacore/src/download.rsfor implementing generic media download interfaces across message types
Always refreshMediaConnif it's expired before using it to get current media servers and auth tokens
Runcargo fmtbefore finalizing a feature or fix
Runcargo clippy --all-targetsbefore finalizing a feature or fix
Runcargo test --allbefore finalizing a feature or fix
Files:
waproto/build.rssrc/store/signal.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/protocol.rswaproto/src/whatsapp.rswacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/protocol/sender_keys.rs
waproto/build.rs
📄 CodeRabbit inference engine (AGENTS.md)
waprotocrate uses prost to compile Protocol Buffers definitions fromwhatsapp.protoin abuild.rsscript
Files:
waproto/build.rs
wacore/**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
wacorecrate must remain platform-agnostic and have no dependencies on Tokio or specific databases
Files:
wacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/protocol.rswacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/protocol/sender_keys.rs
🧠 Learnings (2)
📚 Learning: 2025-12-25T03:36:19.700Z
Learnt from: CR
Repo: jlucaso1/whatsapp-rust PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-25T03:36:19.700Z
Learning: Applies to waproto/build.rs : `waproto` crate uses prost to compile Protocol Buffers definitions from `whatsapp.proto` in a `build.rs` script
Applied to files:
waproto/build.rswacore/libsignal/src/protocol/protocol.rswaproto/src/whatsapp.rswacore/libsignal/src/protocol/sender_keys.rs
📚 Learning: 2025-12-25T03:36:19.700Z
Learnt from: CR
Repo: jlucaso1/whatsapp-rust PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-25T03:36:19.700Z
Learning: The project is organized into three crates: `wacore` (platform-agnostic, `no_std` core logic with no Tokio/database dependencies), `waproto` (Protocol Buffers definitions compiled with prost), and `whatsapp-rust` (main client with Tokio/Diesel integration)
Applied to files:
waproto/build.rswaproto/src/whatsapp.rs
🧬 Code graph analysis (4)
wacore/libsignal/src/protocol/ratchet.rs (4)
wacore/libsignal/src/protocol/state/session.rs (1)
root_key(159-165)wacore/libsignal/src/protocol/protocol.rs (1)
chain_key(569-571)wacore/libsignal/src/protocol/ratchet/keys.rs (1)
derive_keys(122-138)wacore/libsignal/src/protocol/ratchet/params.rs (1)
our_one_time_pre_key_pair(145-147)
wacore/libsignal/src/protocol/state/session.rs (2)
wacore/libsignal/src/protocol/protocol.rs (1)
chain_key(569-571)wacore/libsignal/src/protocol/ratchet/keys.rs (1)
next_chain_key(185-190)
waproto/src/whatsapp.rs (1)
wacore/binary/src/builder.rs (1)
bytes(40-43)
wacore/libsignal/src/protocol/group_cipher.rs (2)
wacore/libsignal/src/protocol/sender_keys.rs (6)
new(39-52)new(94-99)new(176-198)sender_chain_key(212-224)sender_key_state(311-316)chain_id(208-210)wacore/libsignal/src/crypto/aes_cbc.rs (1)
aes_256_cbc_decrypt_into(66-93)
🔇 Additional comments (30)
src/store/signal.rs (3)
207-216: No changes needed—the code is correct as written.The original assumption that
public_key_bytes()returns[u8; 33]is incorrect. It actually returns&[u8](a slice reference), as defined inwacore/libsignal/src/core/curve.rs:93. Thetry_into()conversion on line 213 is therefore necessary to convert from a dynamically-sized slice to the fixed-size array[u8; 32]expected byput_identity(). The error handling for "Invalid key length" is not dead code—it correctly triggers when the key is not exactly 32 bytes.Likely an incorrect or invalid review comment.
193-195: Type compatibility verified — conversion is correct.The direct conversion using
.into()is type-safe.Device.identity_keyis of typeKeyPair, andFrom<KeyPair> for IdentityKeyPairis properly implemented in the Signal protocol library, making the conversion reliable.
333-339: The code is correct—no changes needed.The function signature and field types are fully compatible.
self.signed_pre_keyis of typeKeyPairand matches thekey_pair: &KeyPairparameter expected bynew_signed_pre_key_record. Theself.signed_pre_key_signaturefield is[u8; 64]and matches thesignature: [u8; 64]parameter. The implementation is type-safe.waproto/build.rs (2)
48-83: Update the comment: bytes crate serde feature is enabled; the skip/default attributes are a valid security measure.The
#[serde(skip, default)]attributes are appropriate and safe. The bytes crate in the workspace (Cargo.toml) explicitly enables the serde feature:bytes = { version = "1.5", default-features = false, features = ["serde"] }The comment's claim that "bytes::Bytes doesn't implement Serialize/Deserialize without the serde feature which prost doesn't expose" is inaccurate—the serde feature is available and enabled. However, applying
skipto these sensitive key fields is a legitimate design choice to intentionally prevent them from being included in JSON serialization, treating the protobuf format as the sole serialization mechanism. Verification found no code paths that attempt JSON serialization ofSessionStructureorSenderKeyStateStructure, confirming these fields remain protected.Please correct the comment to reflect that the bytes serde feature is available but intentionally not used here to exclude sensitive key material from JSON output.
Likely an incorrect or invalid review comment.
35-46: Verify that all call sites handle the Vec → bytes::Bytes migration.The change from
Vec<u8>tobytes::Bytesfor these cryptographic fields is a breaking API change.Bytesis immutable and reference-counted, with different semantics thanVec<u8>. The codebase has been properly updated throughout:
- All serialization sites use
Bytes::copy_from_slice()to construct Bytes fields- All deserialization sites safely convert via
.as_deref().and_then(|b| b.try_into().ok())with length validation- All encryption/decryption operations properly accept slices from the Bytes fields
- The
#[serde(skip, default)]attributes are appropriate since these structures are not JSON-serializedThe migration is complete and correct.
waproto/src/whatsapp.rs (2)
11023-11044: LGTM! Consistent optimization for session key material.The conversion of
ChainKeyandMessageKeybinary fields toprost::bytes::Bytesmaintains consistency with the changes in hunk 1 and properly secures all cryptographic material (keys, MAC keys, and IVs) with#[serde(skip, default)]. Call sites throughout the codebase have been correctly updated to useBytes::copy_from_slice()when constructing protobuf structures and properly handle the new type with methods like.as_deref()and.is_some_and().
10942-10967: LGTM! Efficient and secure handling of cryptographic material.The migration from
Vec<u8>toprost::bytes::Bytesfor cryptographic seeds and keys is well-aligned with the PR's allocation-reduction goals.prost::bytes::Bytesuses reference counting internally, enabling cheap cloning without additional heap allocations.The
#[serde(skip, default)]annotations appropriately prevent accidental serialization of sensitive cryptographic material (seeds and private keys) while providing safe defaults during deserialization.All call sites in the cryptography layer (
wacore/libsignal/src/protocol/sender_keys.rs) already properly handle the new type, using.as_ref()for slice conversion andBytes::copy_from_slice()for construction.wacore/libsignal/src/protocol/state/session.rs (5)
87-88: LGTM! Serialization correctly updated for fixed-size array returns.The change from
into_vec()toto_vec()properly handles the updatedPublicKey::serialize()signature that now returns[u8; 33]instead ofBox<[u8]>.Also applies to: 96-96
286-290: LGTM! Efficient chain key storage using prost::bytes::Bytes.The use of
Bytes::copy_from_slice()efficiently handles the 32-byte chain key storage, aligning with the broader migration toprost::bytes::Bytesfor protobuf binary fields.
317-322: LGTM! Consistent with the chain key storage pattern.
369-374: LGTM! Proper Bytes usage for chain key.
447-461: LGTM! Consistent Bytes usage.wacore/libsignal/src/protocol/ratchet.rs (2)
49-85: LGTM! Excellent stack-based optimization for Alice session initialization.The fixed-size buffer
[u8; 160]eliminates heap allocations during session setup. The buffer is correctly sized for the maximum case (discontinuity + 4 agreements), and thedebug_assert!checks provide safety during development.
108-149: LGTM! Bob's session initialization correctly mirrors Alice's optimization.The stack-based buffer approach is consistently applied, matching Alice's implementation with Bob-specific agreement sequences.
wacore/libsignal/src/protocol/sender_keys.rs (4)
11-12: LGTM! HMAC imports and protobuf conversions properly updated.The HMAC dependencies support the new
step_with_message_key()optimization, and theBytes-based protobuf conversions align with the migration strategy.Also applies to: 54-80
84-84: LGTM! Efficient combined key derivation with proper overflow protection.The new
step_with_message_key()method efficiently reuses the HMAC instance viafinalize_reset(), reducing redundant key setup. The overflow check prevents unbounded chain iteration.Also applies to: 105-105, 127-153, 162-166
184-193: LGTM! Consistent Bytes usage for key serialization.
568-568: LGTM! Comprehensive test coverage for optimized key derivation.The new tests thoroughly validate that
step_with_message_key()produces identical results to separate calls, and verify correctness over multiple iterations.Also applies to: 764-825
wacore/libsignal/src/protocol/protocol.rs (4)
87-87: LGTM! Correct serialization for fixed-size array returns.Also applies to: 262-263
534-534: LGTM! Accessor implementations consistent with fixed-size chain_key.Also applies to: 559-576
621-634: LGTM! Robust validation and conversion for protobuf deserialization.The explicit length validation (line 628) ensures only valid 32-byte chain keys are converted, providing defense against malformed protobuf inputs.
400-417: Serialization and signing scope look correct.The signing scope properly includes the version byte (
[version_byte || proto]), andverify_signature()at line 430 correctly slices off only the signature to match. The version byte shift at line 405 hardcodes version 3 (| 3u8), which aligns withSENDERKEY_MESSAGE_CURRENT_VERSION = 3.One thing to note: the parsing code at lines 477-480 explicitly rejects messages with version < 3 (
LegacyCiphertextVersion), which provides built-in protection against incompatibility with older message formats. If the signature scope changed between protocol versions, the version field enforcement ensures only compatible implementations can parse the messages.wacore/libsignal/src/protocol/ratchet/keys.rs (4)
10-34: LGTM! Clear documentation and correct seed storage.The expanded enum documentation clarifies the purpose of each variant, and the fixed-size array parameter enforces compile-time size validation.
68-100: LGTM! Efficient protobuf conversions with proper validation.The zero-cost pass-through for
Serializedvariants and the thorough validation infrom_pb()(lines 91-99) provide both performance and safety.
103-110: LGTM! Efficient combined key derivation with overflow protection.The
step_with_message_keys()method efficiently reuses the HMAC instance viafinalize_reset(), and the overflow check on line 130 (though not visible in this file's hunks, referenced from sender_keys.rs pattern) prevents unbounded iteration. TheCopyderive is safe for the fixed-size key structure.Also applies to: 161-161, 199-218
377-480: LGTM! Comprehensive test coverage validates optimization correctness.The tests thoroughly verify that
step_with_message_keys()produces identical results to separate derivations and maintains consistency across multiple iterations.wacore/libsignal/src/protocol/group_cipher.rs (4)
144-152: LGTM! Efficient key derivation with proper state management.The refactored loop correctly uses
step_with_message_key()to advance through chain keys while storing intermediate message keys. The final state update at line 151 ensures consistency.
207-237: LGTM! Buffer reuse pattern correctly implemented.The decryption path now properly reuses
DECRYPTION_BUFFERfollowing the same pattern as encryption. Error mapping is comprehensive, and buffer capacity is correctly restored at line 235.
325-325: Theseed()method correctly returns&[u8; 32], making the dereference operation valid. The change fromsender_chain_key.seed().to_vec()to*sender_chain_key.seed()is correct and avoids unnecessary heap allocation.
73-73: Implementation verified:step_with_message_key()is correctly implemented inSenderChainKeywith signatureResult<(SenderMessageKey, Self), SignalProtocolError>, proper error handling, and correct usage at both locations.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
wacore/libsignal/src/protocol/group_cipher.rs (1)
43-45: Consider reusing the buffer instead of always allocating.
take_bufferalways allocates a newVecwithINITIAL_CAPACITY. If the caller doesn't need ownership (e.g., they copy data immediately), consider an alternative that preserves the existing allocation:fn take_data(&mut self) -> Vec<u8> { std::mem::take(&mut self.buffer) }Then separately call
get_buffer()on the next use, which will reuse the existing capacity if theVecwas returned. However, I acknowledge the current approach guarantees consistent behavior and avoids capacity bloat from large messages.wacore/libsignal/src/protocol/ratchet/keys.rs (1)
42-70: Defense-in-depth with debug assertion is appropriate.The
debug_assert!at line 58-61 catches invariant violations during development while theunwrap_orfallbacks provide graceful degradation in release builds. Sincefrom_pb()validates field lengths before creatingSerialized, the assertion should never trigger in correct usage.However, the fallback values (
[0u8; 32],[0u8; 16]) would produce incorrect cryptographic keys if triggered. Consider logging a warning in release mode:♻️ Optional: Add release-mode warning
debug_assert!( cipher_key.is_some() && mac_key.is_some() && iv.is_some(), "Serialized MessageKeyGenerator has invalid field lengths - from_pb should have rejected this" ); + +#[cfg(not(debug_assertions))] +if cipher_key.is_none() || mac_key.is_none() || iv.is_none() { + log::error!("Serialized MessageKeyGenerator has invalid field lengths - from_pb should have rejected this"); +}
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
wacore/libsignal/src/protocol/group_cipher.rswacore/libsignal/src/protocol/ratchet.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/sender_keys.rswacore/libsignal/src/protocol/state/session.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- wacore/libsignal/src/protocol/ratchet.rs
🧰 Additional context used
📓 Path-based instructions (2)
**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
**/*.rs: Never modify Device state directly; useDeviceCommand+PersistenceManager::process_command()for state changes
For read-only Device state access, usePersistenceManager::get_device_snapshot()
All blocking I/O (such asureqcalls) and heavy CPU-bound tasks (like media encryption) MUST be wrapped intokio::task::spawn_blockingto avoid stalling the async runtime
UseClient::chat_locksto serialize per-chat operations in asynchronous code
Usethiserrorfor custom domain-specific errors (e.g.,SocketError) andanyhow::Errorfor functions with multiple failure modes
Avoid.unwrap()and.expect()outside of tests and unrecoverable logic paths
Use theDownloadabletrait inwacore/src/download.rsfor implementing generic media download interfaces across message types
Always refreshMediaConnif it's expired before using it to get current media servers and auth tokens
Runcargo fmtbefore finalizing a feature or fix
Runcargo clippy --all-targetsbefore finalizing a feature or fix
Runcargo test --allbefore finalizing a feature or fix
Files:
wacore/libsignal/src/protocol/sender_keys.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/group_cipher.rs
wacore/**/*.rs
📄 CodeRabbit inference engine (AGENTS.md)
wacorecrate must remain platform-agnostic and have no dependencies on Tokio or specific databases
Files:
wacore/libsignal/src/protocol/sender_keys.rswacore/libsignal/src/protocol/ratchet/keys.rswacore/libsignal/src/protocol/state/session.rswacore/libsignal/src/protocol/group_cipher.rs
🧠 Learnings (1)
📚 Learning: 2025-12-25T03:36:19.700Z
Learnt from: CR
Repo: jlucaso1/whatsapp-rust PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-12-25T03:36:19.700Z
Learning: Applies to waproto/build.rs : `waproto` crate uses prost to compile Protocol Buffers definitions from `whatsapp.proto` in a `build.rs` script
Applied to files:
wacore/libsignal/src/protocol/state/session.rs
🧬 Code graph analysis (2)
wacore/libsignal/src/protocol/ratchet/keys.rs (2)
wacore/libsignal/src/protocol/sender_keys.rs (6)
seed(105-107)cipher_key(71-73)iv(67-69)new(39-52)new(94-99)new(176-198)wacore/libsignal/src/protocol/ratchet.rs (1)
derive_keys(18-20)
wacore/libsignal/src/protocol/state/session.rs (2)
wacore/libsignal/src/protocol/protocol.rs (1)
chain_key(569-571)wacore/libsignal/src/protocol/ratchet/keys.rs (1)
next_chain_key(192-197)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Build & Test
- GitHub Check: benchmark
🔇 Additional comments (31)
wacore/libsignal/src/protocol/sender_keys.rs (11)
11-12: LGTM!The HMAC imports are correctly added to support the new
step_with_message_keyoptimization that reuses the HMAC context.
54-60: LGTM!The
from_protobufcorrectly handles theBytestype withas_ref()for the slice conversion, and the fixed-size array conversion is properly validated.
76-80: LGTM!Using
Bytes::copy_from_slicefor protobuf serialization is appropriate and aligns with the PR's goal of consistent byte handling across protobuf fields.
84-88: LGTM!Adding
Copyderive toSenderChainKeyis a good optimization. The struct only contains au32and[u8; 32], both of which areCopy, making this derivation valid and enabling more efficient pass-by-value semantics used elsewhere in this PR.
105-107: LGTM!Returning
&[u8; 32]instead of&[u8]provides compile-time size guarantees and enables callers to use the fixed-size array directly without runtime validation.
127-153: Well-designed optimization using HMAC context reuse.The
step_with_message_keymethod correctly:
- Validates iteration overflow with
checked_add- Reuses the HMAC key setup via
finalize_reset()between deriving message key and chain key- Returns both values atomically, enabling callers to avoid separate calls
The logic mirrors the separate
sender_message_key()andnext()methods, ensuring equivalence. The tests at lines 767-829 validate this.
161-166: LGTM!Consistent use of
Bytes::copy_from_slicefor protobuf field population.
184-192: LGTM!The
SenderKeyState::newcorrectly usesBytes::copy_from_slicefor both public and private key serialization, maintaining consistency with the PR's protobuf byte handling approach.
264-270: Good use ofdrain()for O(n) eviction.The comment correctly explains the performance improvement: using
drain(..excess)once is O(n) total, whereas callingremove(0)in a loop would be O(n) per removal, leading to O(n²) for evicting multiple keys.
572-572: LGTM!Passing
SenderChainKeyby value is now possible and preferred sinceCopywas added. This avoids an unnecessary clone.
767-829: Comprehensive test coverage for the new optimization.The tests properly verify:
test_step_with_message_key_equivalence: Single-step equivalence between combined and separate callstest_step_with_message_key_chain: Multi-iteration synchronization ensuring chains stay alignedBoth tests compare iteration, iv, cipher_key, and seed values to confirm behavioral equivalence.
wacore/libsignal/src/protocol/state/session.rs (5)
87-88: LGTM!The serialization now correctly uses
.serialize().to_vec()for identity keys and base key. This aligns with the PR's fixed-size array return types fromserialize()methods.Also applies to: 96-96
286-312: Good refactor with consistent capacity management.The changes correctly:
- Use
Bytes::copy_from_slicefor protobufChainKey.keyfield- Compute length once before logging and draining
- Use
drain(..excess)for efficient O(n) removal of oldest chainsThe log message now correctly references the captured
lenvalue.
322-325: LGTM!Consistent use of
Bytes::copy_from_slicefor the sender chain key.
374-377: LGTM!Consistent pattern for
set_sender_chain_keyusingBytes::copy_from_slice.
460-464: LGTM!Consistent pattern for
set_receiver_chain_keyusingBytes::copy_from_slice.wacore/libsignal/src/protocol/group_cipher.rs (6)
20-51: Well-structured buffer abstraction for both encryption and decryption.The
CryptoBufferdesign with separateENCRYPTION_BUFFERandDECRYPTION_BUFFERthread-locals provides:
- Isolated buffer reuse for encrypt/decrypt paths
get_buffer()clears and returns a mutable reference for writingtake_buffer()returns owned contents and replaces with a fresh pre-allocated bufferThis matches the existing encryption buffer pattern and extends it to decryption.
83-109: LGTM!The encryption flow correctly uses
step_with_message_key()to obtain both the message keys and next chain key in one call, then updates the state withset_sender_chain_key(next_sender_chain_key). This eliminates the separatesender_message_key()andnext()calls.
151-159: LGTM!The
get_sender_keyfunction correctly usesstep_with_message_key()in the loop to derive message keys and advance the chain. The final step derives the result key and sets the next chain key.
214-242: LGTM!The decryption path now uses
DECRYPTION_BUFFERwith the sameget_buffer()/take_buffer()pattern as encryption, providing symmetric buffer reuse for both operations.
260-274: LGTM!The field access pattern is cleaner, directly calling
skdm.chain_id(),skdm.iteration(),skdm.chain_key(), and dereferencingskdm.signing_key().
330-330: LGTM!Dereferencing
*sender_chain_key.seed()passes the[u8; 32]by value, avoiding ato_vec()allocation sinceSenderKeyDistributionMessage::newnow accepts a fixed-size array.wacore/libsignal/src/protocol/ratchet/keys.rs (9)
10-12: LGTM!HMAC imports added for the
step_with_message_keysoptimization.
15-28: Well-designed lazy evaluation pattern.The
MessageKeyGeneratorenum enables two key optimizations:
Seed: Defers expensive HKDF derivation until keys are actually neededSerialized: Enables zero-cost round-trips by preserving the original protobuf bytesThis is a good example of lazy evaluation for performance-critical cryptographic operations.
31-34: LGTM!
new_from_seedcorrectly takes&[u8; 32]and copies into theSeedvariant, matching the fixed-size pattern throughout this PR.
73-91: LGTM!The
into_pbmethod correctly provides:
- Zero-cost pass-through for
Serializedvariant (just returns the original protobuf)- Proper serialization for
SeedandKeysvariants usingBytes::copy_from_sliceNote that line 82 calls
self.generate_keys()which consumesself, but since we're in a match arm that already checked forSeedorKeys, this is correct.
94-107: LGTM!The
from_pbvalidation ensures:
cipher_keyis exactly 32 bytesmac_keyis exactly 32 bytesivis exactly 16 bytesThis upfront validation guarantees that
generate_keys()can safely convert theSerializedvariant without runtime failures.
109-117: LGTM!The
counter()accessor efficiently retrieves the index from any variant without triggering full key derivation.
168-172: LGTM!Adding
CopytoChainKeyis valid since it only contains[u8; 32]andu32. This enables efficient pass-by-value usage.
206-225: Well-implemented HMAC context reuse.The
step_with_message_keysmethod correctly:
- Creates the HMAC context once with the chain key
- Derives message key seed using
finalize_reset()to reuse the context- Derives next chain key using
finalize()- Returns both values, avoiding the overhead of two separate HMAC setups
This mirrors the
SenderChainKey::step_with_message_keypattern in sender_keys.rs.
420-487: Thorough equivalence testing.The tests properly verify that
step_with_message_keys()produces identical results to the separatemessage_keys()+next_chain_key()path:
test_step_with_message_keys_equivalence: Single-step comparisontest_step_with_message_keys_chain: 10-iteration synchronization testThis provides good confidence that the optimization doesn't change behavior.
1 participant
Summary
This PR optimizes the libsignal cryptographic library by reducing heap allocations. The main changes replace
Vec<u8>andBox<[u8]>with fixed-size arrays ([u8; N]) for cryptographic keys and signatures, and introduce an optimizedstep_with_message_key()method that reuses HMAC context.Test Coverage
Added comprehensive tests for the new
step_with_message_keys()optimization:test_step_with_message_keys_equivalence: Verifies identical results to separate callstest_step_with_message_keys_chain: Verifies behavior over multiple iterationstest_sender_chain_key_iteration_step_*: Tests forSenderChainKey::step_with_message_key()Summary by CodeRabbit
Refactor
Bug Fixes
Compatibility / Serialization
✏️ Tip: You can customize this high-level summary in your review settings.