perf(noise): pre-key the transport AES-GCM once per connection - #850
Merged
Conversation
Every Noise frame, both directions, re-ran the AES-256 key schedule and re-derived the GHASH subkey through setup_gcm even though the transport key is fixed for the connection lifetime; the flamegraph put that setup at 17% of a typical 1500-byte frame seal, plus a per-frame allocation inside it. Aes256GcmKey holds the expanded cipher and the keyed GHASH; per frame only the nonce-dependent counter init and pad block remain, with the keyed state cloned (a plain copy, no key expansion). NoiseCipher stores it and drives the existing encrypt/decrypt bodies through new_with_key constructors, so the output is byte-identical, pinned by a differential test against the per-call setup across sizes, nonces and aad shapes. The handshake path (keys change per step) keeps the per-call setup; the IK outcome enum boxes its grown variant. Local wall-time A/B on 1500-byte frames: decrypt median 11.51 to 9.19 us (-20%), encrypt 9.76 to 9.08 us; 64 KB frames unchanged (data path dominates).
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds a pre-keyed AES-256-GCM key type and TransportAead trait, exposes them via crypto re-exports, integrates connection-lifetime AEAD into NoiseCipher to avoid per-call key expansion, and boxes the IkServerHelloOutcome::Continue payload. ChangesCryptographic Pre-keying and Noise Protocol Optimization
Sequence DiagramssequenceDiagram
participant Client
participant Aes256GcmKey
participant GcmGhash
participant AES as AES256_Cipher
Client->>Aes256GcmKey: new(raw_key)
Aes256GcmKey->>AES: expand key schedule
Aes256GcmKey->>GcmGhash: precompute keyed GHASH (H = AES(key, 0))
Client->>Aes256GcmKey: setup(nonce, aad)
Aes256GcmKey->>GcmGhash: from_keyed(keyed_h, aad)
GcmGhash->>GcmGhash: compute GHASH pad and init state
Aes256GcmKey-->>Client: (CTR, GcmGhash)
Client->>Client: encrypt/decrypt using CTR + GcmGhash
sequenceDiagram
participant NoiseCipher
participant Provider
participant Aes256GcmKey
participant AEAD as TransportAead
participant Buffer
NoiseCipher->>Provider: provider().transport_aead(key)
Provider-->>NoiseCipher: boxed Aes256GcmKey as TransportAead
rect rgba(100, 150, 200, 0.5)
Note over NoiseCipher,Buffer: Encryption with counter
NoiseCipher->>AEAD: encrypt_in_place(iv, buffer, aad)
AEAD->>Buffer: in-place ciphertext + appended tag
end
rect rgba(150, 150, 100, 0.5)
Note over NoiseCipher,Buffer: Decryption with counter
NoiseCipher->>NoiseCipher: validate buffer ≥ TAG_LEN
NoiseCipher->>AEAD: decrypt_in_place(iv, buffer, aad)
AEAD->>NoiseCipher: success or auth failure
NoiseCipher->>Buffer: truncate to plaintext on success
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
You need this to work end-to-end — verify the provider wiring and the boxed handshake change compile cleanly. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
1 issue found across 4 files
Reply with feedback, questions, or to request a fix.
Fix all with cubic | Re-trigger cubic
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a567ccf452
ℹ️ 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".
Contributor
Merging this PR will improve performance by 10.55%
|
| Mode | Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|---|
| ⚡ | Simulation | bench_frame_encrypt_in_place[1500] |
34 µs | 30.7 µs | +10.55% |
Tip
Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.
Comparing perf/noise-prekeyed-gcm (3cbc332) with main (0bcc3a0)
…vider hook NoiseCipher was driving the RustCrypto GCM types directly, so a custom SignalCryptoProvider would silently stop observing transport crypto. The trait now offers transport_aead(key): the default returns a per-call adapter that keeps routing every frame through the configured provider, and RustCryptoProvider overrides it with the pre-keyed Aes256GcmKey fast path. One virtual call per frame, behavior identical on both paths.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@wacore/libsignal/src/crypto/provider.rs`:
- Around line 106-129: PerCallTransportAead currently captures only the key and
re-dispatches to the global provider() in encrypt_in_place/decrypt_in_place,
which loses the SignalCryptoProvider identity passed into transport_aead();
change PerCallTransportAead to hold a reference/handle to the originating
SignalCryptoProvider (or a boxed trait object/cloneable adapter) and call that
provider's aes_256_gcm_encrypt_in_place/aes_256_gcm_decrypt_in_place methods
instead of provider(); alternatively, restrict construction so transport_aead()
only returns implementations tied to the active-provider helper. Also apply the
same fix to the other adapter at the second occurrence (lines referenced around
221-229) so no transport AEAD re-dispatches to the global provider().
🪄 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: 917d3c37-9e09-456f-96e2-63f8334272bb
📒 Files selected for processing (3)
wacore/libsignal/src/crypto/mod.rswacore/libsignal/src/crypto/provider.rswacore/noise/src/state.rs
…provider The default adapter re-dispatched through the global provider(), losing the identity of the instance whose transport_aead was called. It now holds a reference to that instance; the 'static receiver both enables the capture and restricts construction to the installed provider, which is the only reachable caller.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The CodSpeed flamegraph for
bench_frame_encrypt_in_place[1500](the typical stanza-frame size) showssetup_gcmat 16.9% of the seal: every Noise frame, in both directions, re-runs the AES-256 key schedule, re-derives the GHASH subkeyH = E_K(0), and pays an allocation inside the setup — all key-dependent work, and the transport key is fixed for the whole connection. Same shape as the ltHash fix in #847: a per-call recomputation of a constant.Change
New
crypto::aes_gcm::Aes256GcmKeyholds the expanded AES cipher and the keyed GHASH. Per frame, only the nonce-dependent work remains (CTR init and the one pad block); the keyed state is cloned, which is a plain round-key copy with noaeskeygenassistand no polyval re-keying.Aes256GcmEncryption/Aes256GcmDecryptiongainnew_with_keyconstructors that feed the existing encrypt/decrypt/tag bodies, andNoiseCipherstores the pre-keyed state instead of raw key bytes.Scope notes:
NoiseCiphergrew (round keys live inline), so the IK outcome enum boxes itsContinuevariant per clippy.NoiseCiphernow drives the RustCrypto GCM types directly instead of the pluggableSignalCryptoProviderfree functions. No in-repo or known consumer installs a custom provider, and the signal paths still route through it; flagging in case the hook was meant to cover transport crypto too.Measured
Local wall-time A/B on 1500-byte frames (AES-NI + CLMUL host): decrypt median 11.51 -> 9.19 us (-20%), encrypt 9.76 -> 9.08 us (-7%); 64 KB frames unchanged (data path dominates). The Simulation runner takes the software polyval path where the setup weighs more, so the CodSpeed report should show a larger relative win on the 1500-byte benches.
Tests
Full workspace suites pass (2180 tests, e2e excluded as usual); clippy strict clean.