feat: pluggable SignalCryptoProvider for AES-CBC/GCM + HMAC-SHA256 - #561
Conversation
Introduces a global provider trait (wacore_libsignal::crypto::SignalCryptoProvider) that centralizes the 5 symmetric primitives used across the codebase: AES-256-CBC encrypt/decrypt, AES-256-GCM encrypt/decrypt, HMAC-SHA256. Default RustCryptoProvider preserves current behavior byte-for-byte (all 1000+ tests green, incl. NIST KATs). The OnceLock-backed global follows the same pattern as wacore::time::TimeProvider so the bridge can swap in a native implementation (node:crypto over wasm-bindgen) without touching any caller. Callers unified: - wacore/noise (hot path) — new NoiseBuffer trait keeps Vec+BytesMut callers working without exposing aes_gcm::aead::Buffer. - wacore/src: poll, pair, pair_code, media_retry. Drops the external aes-gcm crate from wacore-libsignal (was already unused — name collided with the internal crypto::aes_gcm module), from wacore-noise, and from wacore's [dependencies] (kept in [dev-dependencies] for the noise_handshake_test reference vector). Also drops the now-unused pub use aes_gcm re-export in wacore::lib and wacore_noise::lib.
|
Caution Review failedPull request was closed or merged during review 📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR adds a pluggable Signal crypto provider, moves AES‑GCM/AES‑CBC/HMAC logic to that provider, and updates callers across wacore and noise to use libsignal::crypto helpers; it also removes Changes
Sequence Diagram(s)sequenceDiagram
participant Caller as Encryption Caller
participant API as libsignal::crypto helper
participant Registry as Provider Registry
participant Provider as SignalCryptoProvider
participant Backend as RustCrypto Primitives
Caller->>API: aes_256_gcm_encrypt(key, nonce, aad, plaintext, &mut out)
API->>Registry: provider()
Registry-->>API: &dyn SignalCryptoProvider
API->>Provider: aes_256_gcm_encrypt(key, nonce, aad, plaintext, out)
Provider->>Backend: perform AES-GCM encrypt (primitive)
Backend-->>Provider: ciphertext || tag
Provider-->>API: Ok(()) / Err(CryptoProviderError)
API-->>Caller: propagate result
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
@claude review this nitpicky |
|
Claude finished @jlucaso1's task in 2m 13s —— View job Nitpicky Code Review — PR #561
Overall this is a clean, well-structured PR. The trait design is solid, the migration is mechanical and correct, and the dep cleanup is a nice bonus. Here are the nits: 1.
|
Benchmark Results59 unchanged benchmark(s)
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2a7232a2e7
ℹ️ 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".
| let plaintext = std::mem::take(buffer); | ||
| aes_256_gcm_encrypt(&self.key, &iv, b"", &plaintext, buffer) |
There was a problem hiding this comment.
Keep in-place encrypt from reallocating every frame
encrypt_in_place_with_counter now calls std::mem::take(buffer), which replaces the caller’s reusable Vec with a fresh zero-capacity allocation before encryption. In the send hot path (src/socket/noise_socket.rs uses a reused enc_buf), this drops the preallocated capacity on every message and forces repeated allocate/free cycles, causing a significant throughput/latency regression under sustained traffic.
Useful? React with 👍 / 👎.
| let mut out = Vec::with_capacity(buffer.as_slice().len().saturating_sub(TAG_LEN)); | ||
| aes_256_gcm_decrypt(&self.key, &iv, b"", buffer.as_slice(), &mut out) | ||
| .map_err(|e| NoiseError::CryptoError(format!("Decrypt failed: {e}")))?; | ||
| buffer.replace_contents(&out); |
There was a problem hiding this comment.
Restore in-place decrypt instead of temp-buffer copying
decrypt_in_place_with_counter allocates a new Vec for every call and then copies plaintext back into the original buffer via replace_contents. This removes the previous true in-place behavior, so high-volume receive paths (e.g., decrypt_frame on BytesMut) now pay an extra allocation and full-buffer copy per frame, which is a production-facing performance regression.
Useful? React with 👍 / 👎.
Addresses Codex review P1/P2 on #561: P1 - encrypt_in_place_with_counter was calling std::mem::take(buffer), replacing the caller's preallocated Vec (e.g. enc_buf = Vec::with_capacity(4096) in noise_socket) with a zero-capacity one. Every subsequent send paid a reallocate/copy. P2 - decrypt_in_place_with_counter now allocated a scratch Vec and copied plaintext back via replace_contents, so BytesMut frames on the hot receive path paid one allocation + a full-buffer memcpy per frame. Fix: add GcmInPlaceBuffer trait (Vec<u8>, BytesMut) and provider methods aes_256_gcm_encrypt_in_place / aes_256_gcm_decrypt_in_place. Default impl keeps current allocating behavior for backward compat. RustCryptoProvider overrides with truly in-place CTR-XOR + tag append / truncate on the caller buffer, preserving capacity. Empirical before/after on the hot paths: - encrypt: capacity dropped 4096 -> 120 every iter -> stays at 4096. - decrypt: 2 allocs per call -> 0 allocs per call (1000 iters). Tests: libsignal 107, noise 23, appstate 375, wacore 21 pass. Clippy clean. No API change visible to external callers.
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wacore/libsignal/src/crypto/mod.rs`:
- Around line 73-82: Update the documentation for the public function
aes_256_gcm_decrypt_in_place to explicitly state the provider contract: on
authentication failure the underlying provider
(provider::provider().aes_256_gcm_decrypt_in_place) may leave `buffer` in an
indeterminate or clobbered state, so callers must not assume the original
ciphertext or any plaintext is preserved and must discard or reinitialize the
buffer after an Err result; keep the existing signature and behavior unchanged,
only expand the doc comment to warn about this failure mode and recommended
caller behavior.
In `@wacore/libsignal/src/crypto/provider.rs`:
- Around line 181-185: The public function set_crypto_provider currently returns
Result<(), &'static str>; replace this with a typed error type (e.g., a new
CryptoProviderError enum) that derives thiserror::Error and Debug and add a
variant like AlreadySet; change the signature of set_crypto_provider(provider:
impl SignalCryptoProvider) -> Result<(), CryptoProviderError>, map the
CRYPTO_PROVIDER.set(...) failure into CryptoProviderError::AlreadySet (or other
variants if needed), and update any callers/tests to handle the new error type;
ensure you add the thiserror dependency and export the error type alongside
set_crypto_provider.
In `@wacore/noise/src/state.rs`:
- Around line 47-54: encrypt_in_place_with_counter is currently fixed to &mut
Vec<u8>, causing unnecessary copies for other buffer types; change its signature
to accept a generic buffer type (e.g., B: NoiseBuffer or whatever
trait/dependency decrypt uses) like pub fn encrypt_in_place_with_counter<B:
NoiseBuffer>(&self, counter: u32, buffer: &mut B) -> Result<()> and update the
body to call generate_iv(counter) and aes_256_gcm_encrypt_in_place(&self.key,
&iv, b"", buffer) (ensuring aes_256_gcm_encrypt_in_place accepts the same
NoiseBuffer trait) and keep the same
map_err(NoiseError::CryptoError(format!("{e}"))) and semantics (in-place
encryption, preserving allocation/capacity).
🪄 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: 67615d72-61e3-4be6-9b09-6bfb63d851f2
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (4)
wacore/libsignal/Cargo.tomlwacore/libsignal/src/crypto/mod.rswacore/libsignal/src/crypto/provider.rswacore/noise/src/state.rs
| pub fn set_crypto_provider(provider: impl SignalCryptoProvider) -> Result<(), &'static str> { | ||
| CRYPTO_PROVIDER | ||
| .set(Box::new(provider)) | ||
| .map_err(|_| "crypto provider already set") | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Use a typed error for provider installation failures.
This new public API falls back to &'static str, which is the wrong shape for a crypto initialization boundary. Make the failure matchable like the rest of this module.
♻️ Proposed fix
+#[derive(Debug, displaydoc::Display, thiserror::Error)]
+pub enum SetCryptoProviderError {
+ /// crypto provider already set
+ AlreadySet,
+}
+
/// Install a custom provider. Must be called before any crypto call. Returns
/// `Err` if a provider was already set (including by `get_or_init` of the
/// default fallback).
-pub fn set_crypto_provider(provider: impl SignalCryptoProvider) -> Result<(), &'static str> {
+pub fn set_crypto_provider(
+ provider: impl SignalCryptoProvider,
+) -> Result<(), SetCryptoProviderError> {
CRYPTO_PROVIDER
.set(Box::new(provider))
- .map_err(|_| "crypto provider already set")
+ .map_err(|_| SetCryptoProviderError::AlreadySet)
}As per coding guidelines "Use thiserror for typed errors; use anyhow for multi-failure functions; never use .unwrap() outside tests".
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| pub fn set_crypto_provider(provider: impl SignalCryptoProvider) -> Result<(), &'static str> { | |
| CRYPTO_PROVIDER | |
| .set(Box::new(provider)) | |
| .map_err(|_| "crypto provider already set") | |
| } | |
| #[derive(Debug, displaydoc::Display, thiserror::Error)] | |
| pub enum SetCryptoProviderError { | |
| /// crypto provider already set | |
| AlreadySet, | |
| } | |
| /// Install a custom provider. Must be called before any crypto call. Returns | |
| /// `Err` if a provider was already set (including by `get_or_init` of the | |
| /// default fallback). | |
| pub fn set_crypto_provider( | |
| provider: impl SignalCryptoProvider, | |
| ) -> Result<(), SetCryptoProviderError> { | |
| CRYPTO_PROVIDER | |
| .set(Box::new(provider)) | |
| .map_err(|_| SetCryptoProviderError::AlreadySet) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wacore/libsignal/src/crypto/provider.rs` around lines 181 - 185, The public
function set_crypto_provider currently returns Result<(), &'static str>; replace
this with a typed error type (e.g., a new CryptoProviderError enum) that derives
thiserror::Error and Debug and add a variant like AlreadySet; change the
signature of set_crypto_provider(provider: impl SignalCryptoProvider) ->
Result<(), CryptoProviderError>, map the CRYPTO_PROVIDER.set(...) failure into
CryptoProviderError::AlreadySet (or other variants if needed), and update any
callers/tests to handle the new error type; ensure you add the thiserror
dependency and export the error type alongside set_crypto_provider.
There was a problem hiding this comment.
This creates an inconsistency with set_time_provider (same signature Result<(), &'static str>). It would only make sense if both were migrated together.
There was a problem hiding this comment.
@jlucaso1 That's a fair point — set_time_provider and set_crypto_provider are intentionally parallel APIs, so breaking their return-type symmetry here would be more confusing than helpful. Migrating both to a typed error in a single follow-up PR is the right call. I'll drop this suggestion.
✏️ Learnings added
Learnt from: jlucaso1
Repo: oxidezap/whatsapp-rust PR: 561
File: wacore/libsignal/src/crypto/provider.rs:181-185
Timestamp: 2026-04-17T17:30:44.592Z
Learning: In the whatsapp-rust repo, `set_crypto_provider` and `set_time_provider` are intentionally parallel APIs and both return `Result<(), &'static str>`. They should be kept consistent with each other; any migration to typed errors should cover both simultaneously.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
Addresses 2 of 3 inline suggestions from code review; third (typed error on set_crypto_provider) rejected as inconsistent with the identical set_time_provider signature. - aes_256_gcm_decrypt_in_place: document that AuthFailed leaves the buffer with CTR-XOR pseudo-plaintext (not the original ciphertext), so callers must discard / treat session as compromised. Empirically confirmed: buffer[..pt_len] equals the derived plaintext after a corrupted-tag decrypt attempt. - NoiseCipher::encrypt_in_place_with_counter: relax the fixed '&mut Vec<u8>' to '<B: NoiseBuffer>', matching the existing decrypt counterpart. Zero-cost via the existing GcmInPlaceBuffer trait; BytesMut roundtrips verified with capacity preserved.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a4fd63901d
ℹ️ 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".
| provider() | ||
| .aes_256_cbc_decrypt(key, iv, ctext, output) | ||
| .map_err(|_| DecryptionError::BadCiphertext("failed to decrypt")) |
There was a problem hiding this comment.
Preserve backend failures in CBC decrypt error mapping
When a custom SignalCryptoProvider returns CryptoProviderError::BackendFailed from CBC decrypt, this wrapper collapses it into BadCiphertext("failed to decrypt"), making provider/runtime outages indistinguishable from corrupted ciphertext. In downstream paths (for example protocol/session_cipher.rs), BadCiphertext is handled as an invalid incoming message, so a transient backend failure can cause valid messages to be dropped instead of surfacing an operational crypto-provider error that can be retried or failed fast.
Useful? React with 👍 / 👎.
Summary
wacore_libsignal::crypto::SignalCryptoProvider— a pluggable trait covering the 5 symmetric primitives actually used across the codebase: AES-256-CBC enc/dec, AES-256-GCM enc/dec, HMAC-SHA256. SameOnceLock+set_*_providerpattern already in use bywacore::time::TimeProvider.RustCryptoProviderpreserves current behavior. No call-site visible change; no new deps.hmac_sha256, and the GCM users inwacore/noise,wacore/src/{poll,pair,pair_code,media_retry}. The bridge can now install anode:crypto-backed provider in a follow-up PR without touching any consumer code.Why
Profile of messaging steady-state showed ~5–6% CPU in AES/SHA soft primitives (fixslice AES + sha2::compress256). Most of the AES time is GCM (Noise rekey/encrypt/decrypt + media retry + pair) — moving to Node's OpenSSL-backed
createCipheriv/createHmacis meaningfully faster, but only worthwhile if it can land without forking call sites. The provider trait lets the bridge install native primitives once atinitWasmEngineand have every caller benefit transparently.Changes
New
wacore/libsignal/src/crypto/provider.rs— trait,CryptoProviderError,RustCryptoProviderdefault,set_crypto_provider, internalprovider().Delegated to provider (behavior preserved)
aes_256_cbc_encrypt_into/_decrypt_into(free functions)protocol::crypto::hmac_sha256aes_256_gcm_encrypt,aes_256_gcm_decrypt,hmac_sha256inwacore_libsignal::crypto.Migrated callers
wacore/noise/src/state.rs— hot path. To preserve the generic `decrypt_in_place_with_counter`, adds a smallNoiseBuffertrait (impl'd forVec<u8>andbytes::BytesMut) replacing the previousaes_gcm::aead::Bufferbound.wacore/src/{poll,pair,pair_code,media_retry}.rs.Dep cleanup (via
cargo shear)wacore-libsignal/Cargo.toml: dropaes-gcm(was never used — the crate name shadowed the localcrypto::aes_gcmmodule, so cargo-shear had missed it earlier).wacore-noise/Cargo.toml: dropaes-gcm.wacore/Cargo.toml:aes-gcmmoved to[dev-dependencies](onlytests/noise_handshake_test.rskeeps using it, as a cross-check reference vector).pub use aes_gcm;re-exports inwacore::libandwacore_noise::lib.Non-goals (future PR)
JsCryptoProvider(node:crypto overjs_sys::Function). Trait is shaped for it but not wired yet.CryptographicMacstays on RustCrypto. Appstate LT-Hash / ratchet HMAC inputs are too small (<128 B) for the WASM↔JS boundary cost to be worth it.Test plan
cargo test --all --exclude e2e-tests— 1000+ tests pass (libsignal 107, noise 23, appstate 375, wacore 81, whatsapp-rust 515).cargo clippy --all --tests— no new warnings.cargo fmt --all,cargo sort -wclean.