perf(socket): seal each noise frame where it lands in the batch buffer - #1121
Conversation
The sender staged every frame in a scratch `Vec`, encrypted it there, then copied the framed result into the batch buffer. That is a second full pass over every byte the client sends, on top of the one the encryption itself already makes. A frame's ciphertext is exactly its plaintext plus the 16-byte GCM tag, so the length prefix is known before the bytes it counts exist. Write the prefix, copy the plaintext once into the batch buffer, and hand AES-GCM a view of just that frame's region so it seals in place and appends its tag there. The scratch buffer and one copy per frame both go away. `FrameBody` is that view: it holds the frame's start as an offset, not a slice, so the AEAD can grow the buffer by the tag through it. Every one of its operations is relative to that offset, or it would eat the frames already staged for the same write. Sealing in place is the one path that can leave partial output behind on failure, which the previous code never could, so it rolls the prefix and the plaintext back. Only a `set_crypto_provider` backend can fail there; the counter must stay unburned either way, or the next frame reuses its nonce. quinn (`PacketKey::encrypt` straight into the datagram buffer) and rustls (each fragment encrypted into the record it will send) are both built this way for the same reason.
The length prefix is now written from plaintext.len() + TAG_LEN before the ciphertext exists, which holds only because TransportAead is AES-256-GCM by contract. A set_crypto_provider backend that grew the buffer by anything else would put a frame on the wire whose prefix disagrees with its body and desync the peer's parser for the rest of the connection. One comparison turns that into a refused send.
📝 WalkthroughWalkthroughNoise frame construction now stages headers and seals payloads directly into a shared ChangesNoise frame sealing
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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 |
|
| Filename | Overview |
|---|---|
| src/socket/noise_socket.rs | Replaces scratch-buffer encryption for inline frames with frame-scoped in-place sealing, rollback on failure, and ciphertext-size validation while preserving serial counter allocation. |
| wacore/noise/src/framing.rs | Extracts an error-atomic helper that appends a frame prefix and reserves capacity for the subsequently written payload. |
| tests/noise_frame_size_contract.rs | Adds an isolated custom-provider test proving that a successful seal with the wrong tag growth is rejected before transport output. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart LR
P[Plaintext] --> H[Append predicted length prefix]
H --> B[Append plaintext to batch buffer]
B --> E[Seal frame region in place]
E --> V{Ciphertext length matches prediction?}
V -->|Yes| C[Increment counter and commit frame]
V -->|No| R[Truncate to pre-frame length and return error]
Reviews (2): Last reviewed commit: "test(socket): cover the ciphertext-size ..." | Re-trigger Greptile
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/socket/noise_socket.rs (1)
338-394: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd test coverage for the ciphertext-size guard.
We’ve got the rollback and
out_buf.len() - base != body_lenchecks in place, which is good — but none of the existing noise tests trip a provider that returnsOkwith the wrong ciphertext size. Wire a fakeset_crypto_providerbackend through the crypto-path tests, or add a separate coverage test, so this safety net is exercised in CI.🤖 Prompt for 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. In `@src/socket/noise_socket.rs` around lines 338 - 394, Add a noise test covering the ciphertext-size guard in encrypt_frame_into by configuring a fake set_crypto_provider backend that returns success while producing a ciphertext length different from body_len. Assert the operation returns the crypto error and restores out_buf to its pre-encryption state, using the existing crypto-path test setup and symbols rather than changing production behavior.
🤖 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.
Outside diff comments:
In `@src/socket/noise_socket.rs`:
- Around line 338-394: Add a noise test covering the ciphertext-size guard in
encrypt_frame_into by configuring a fake set_crypto_provider backend that
returns success while producing a ciphertext length different from body_len.
Assert the operation returns the crypto error and restores out_buf to its
pre-encryption state, using the existing crypto-path test setup and symbols
rather than changing production behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: e383c260-ac70-426b-a30b-9b38b088958f
📒 Files selected for processing (2)
src/socket/noise_socket.rswacore/noise/src/framing.rs
📦 Binary size report
.text per crate
Baseline: |
The guard was the one path with no coverage, on the grounds that the crypto provider is a process-wide OnceLock and a deliberately broken backend would poison the unit-test binary. An integration test is its own binary, so it can install one safely. The fake provider delegates everything to the real one except transport_aead, which grows the buffer by 15 bytes instead of 16 and reports success: exactly the shape that would put a frame on the wire whose prefix disagrees with its body. Removing the guard makes the test fail.
Dismissed because a newer commit was pushed; Greptile will re-review the current head.
The noise sender staged every frame in a scratch
Vec, encrypted it there, then copied the framed result into the batch buffer. That second pass over every byte sent is exactly the copy comparable stacks are built to avoid: quinn seals withPacketKey::encrypt(&self, packet, buf, header_len)directly in the datagram buffer, and rustls encrypts each fragment into the record it will send.A frame's ciphertext is exactly its plaintext plus the 16-byte GCM tag, so the length prefix is known before the bytes it counts exist. The sender now writes the prefix, copies the plaintext once into the batch buffer, and hands AES-GCM a view of just that frame's region. The scratch buffer and one copy per frame are gone.
Be clear about the size of this
This is not a measurable win at pingpong frame sizes. At ~200 bytes the removed copy is worth roughly 10 ns per frame, well under 0.1% of CPU on that workload, which is below what the harness can resolve. It was found while testing a different hypothesis (see below) and is offered on two grounds instead:
If the churn on a crypto path is not worth that to you, dropping this is a legitimate call.
What this came out of, and what was rejected
The original hypothesis was concurrent per-frame encryption inside a batch, mirroring WhatsApp Web's
WA/Noise/Socket.js, which startscrypto.subtle.encrypteagerly and uses aPromiseQueueto restore ordering. Measured on this machine, that is a ~4.7x regression, not a marginal loss: 16 x 200 B frames sealed sequentially take 2233 ns; the same work across 16 spawned tasks takes 10583 ns.The arithmetic does not depend on batch size. An AES-256-GCM seal of 200 B is ~140 ns;
tokio::spawn+ join is ~404 ns per task at best (pipelined) and ~3067 ns in the latency-critical case. At this workload's 1.23 frames per write, the parallelizable share is 0.23 frames, so the best case saves ~32 ns and costs ~93 ns. Spawn overhead stays ~3x the work handed off until frames reach ~30 KB, and frames above 16 KB already go off-thread individually.Two further blockers, either of which is disqualifying on its own:
wacore::runtime::Runtime::spawnreturns anAbortHandle, not a result-carryingJoinHandle, because it must stay portable to wasm32 and ESP32. Every parallel frame would therefore cost three heap allocations (a oneshot for the result, aBox::pin, and the task). The current path costs zero.crypto.subtlehas no synchronous form; itsPromiseQueueis ordering repair for an asynchrony the browser imposed, not a parallelism mechanism. In Rust the seal is a synchronous call, so there is no ordering to repair, and copying the shape imports the cost without the cause.Worth noting that the closest analogue in the ecosystem reached the same place: quinn#760, "Parallel symmetric cryptography", was opened by the maintainer in 2020 and has zero comments and zero PRs since. QUIC's packet number is its AEAD nonce, structurally identical to the noise counter. rustls encrypts fragments in a loop and answered a measured 3x throughput deficit against OpenSSL with faster serial code, never with threads.
Correctness
The counter is the AES-GCM nonce, so the sender's existing invariants are the ones this could break. It is still strictly serial, so ordering and counter allocation are untouched; the tests that pin them (
queued_frames_leave_in_one_write_in_counter_order,order_survives_a_full_job_channel,every_encrypted_frame_burns_its_own_counter) all fail under a deliberate base-offset mutation.One contract did change and is documented in place. This is now the only path that can leave partial output on failure, so it rolls the prefix and the plaintext back with
truncate; the doc comment onencrypt_frame_intomoved from "returns before writing a byte" to "leavesout_bufexactly as it found it". Keeping partial output there would let the next frame reuse a nonce, which is why the rollback is load-bearing rather than tidiness.The second commit adds a guard the review asked for. Writing the prefix before the ciphertext exists is sound only because
TransportAeadis AES-256-GCM by contract. Aset_crypto_providerbackend that grew the buffer by anything other than the tag would put a frame on the wire whose prefix disagrees with its body, desyncing the peer's parser for the rest of the connection. One comparison turns silent corruption into a refused send.Tests
Six new tests, each mutation-checked:
base = before(seal over the header)a_frame_is_sealed_in_place_behind_the_one_before_it, plus 6 existing ordering testsFrameBody::resizeabsolute instead ofbase +the_frame_body_view_never_reaches_before_its_own_frameFrameBody::truncateabsolute instead ofbase +a_failed_frame_leaves_the_batch_buffer_byte_identicalappend_frame_header_rejects_oversize_without_writingreserve(prefix_len)instead oftotal_lenappend_frame_header_reserves_room_for_the_payloadlen_bytes[..3]instead of[1..]append_frame_into, which delegates to it, so it agreed with itself. Rewritten to assert literal big-endian bytesThe size guard is covered by
tests/noise_frame_size_contract.rs, an integration test and therefore its own binary, which is what makes installing a deliberately broken provider safe: the process-wideOnceLockonly poisons the binary it is set in. The fake provider delegates everything to the real one excepttransport_aead, which grows the buffer by 15 bytes instead of 16 and reports success. Removing the guard makes it fail.The
truncaterollback remains uncovered: the default AEAD cannot fail on a fixed-size key and nonce, so reaching it needs a provider that errors rather than one that lies about size. It is documented in place so it does not read as dead code.append_frame_header_intois extracted rather than duplicated;append_frame_intonow delegates to it, so the size check and the length encoding exist once.Verification
cargo fmt --all,cargo clippy --workspace --all-targetswith zero warnings,cargo test -p whatsapp-rust --libgreen (1215 tests),wacore-noisegreen (38 tests), and the new integration test green.