Skip to content

perf(socket): seal each noise frame where it lands in the batch buffer - #1121

Merged
jlucaso1 merged 3 commits into
mainfrom
perf/noise-par
Jul 26, 2026
Merged

perf(socket): seal each noise frame where it lands in the batch buffer#1121
jlucaso1 merged 3 commits into
mainfrom
perf/noise-par

Conversation

@jlucaso1

@jlucaso1 jlucaso1 commented Jul 26, 2026

Copy link
Copy Markdown
Collaborator

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 with PacketKey::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:

  • it scales with frame size, unlike task-level tricks, so it matters where frames are large (history sync, media)
  • it makes the sender smaller, not larger: one buffer and one code path fewer

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 starts crypto.subtle.encrypt eagerly and uses a PromiseQueue to 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::spawn returns an AbortHandle, not a result-carrying JoinHandle, because it must stay portable to wasm32 and ESP32. Every parallel frame would therefore cost three heap allocations (a oneshot for the result, a Box::pin, and the task). The current path costs zero.
  • The WA Web design does not transfer. crypto.subtle has no synchronous form; its PromiseQueue is 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 on encrypt_frame_into moved from "returns before writing a byte" to "leaves out_buf exactly 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 TransportAead is AES-256-GCM by contract. A set_crypto_provider backend 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:

Mutation Caught by
base = before (seal over the header) a_frame_is_sealed_in_place_behind_the_one_before_it, plus 6 existing ordering tests
FrameBody::resize absolute instead of base + the_frame_body_view_never_reaches_before_its_own_frame
FrameBody::truncate absolute instead of base + same
counter guard moved after the header append a_failed_frame_leaves_the_batch_buffer_byte_identical
size check moved after the length write append_frame_header_rejects_oversize_without_writing
reserve(prefix_len) instead of total_len append_frame_header_reserves_room_for_the_payload
len_bytes[..3] instead of [1..] initially escaped: the test compared the helper against append_frame_into, which delegates to it, so it agreed with itself. Rewritten to assert literal big-endian bytes

The 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-wide OnceLock only poisons the binary it is set in. 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. Removing the guard makes it fail.

The truncate rollback 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_into is extracted rather than duplicated; append_frame_into now delegates to it, so the size check and the length encoding exist once.

Verification

cargo fmt --all, cargo clippy --workspace --all-targets with zero warnings, cargo test -p whatsapp-rust --lib green (1215 tests), wacore-noise green (38 tests), and the new integration test green.

jlucaso1 added 2 commits July 26, 2026 14:00
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.
@coderabbitai

coderabbitai Bot commented Jul 26, 2026

Copy link
Copy Markdown

Review Change Stack

📝 Walkthrough

Walkthrough

Noise frame construction now stages headers and seals payloads directly into a shared BytesMut batch buffer. Encryption failures roll back staged bytes, while framing, sender, and AEAD contract tests validate lengths, frame isolation, and counter behavior.

Changes

Noise frame sealing

Layer / File(s) Summary
Header-only frame staging
wacore/noise/src/framing.rs
Adds append_frame_header_into, centralizes length-prefix staging, reserves capacity, and rejects oversize payloads without modifying the output buffer.
In-place encrypted frame construction
src/socket/noise_socket.rs
Adds the bounded FrameBody view and updates encrypt_frame_into to append headers and plaintext, seal AEAD in place, validate ciphertext length, and roll back on failure.
Batched sender integration and validation
src/socket/noise_socket.rs, tests/noise_frame_size_contract.rs
Uses BytesMut for batched encrypted output and tests frame isolation, bounded buffer operations, rollback, counter handling, and invalid AEAD tag sizes.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

Suggested labels: performance, api-design

Suggested reviewers: greptile-apps

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 9.09% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title clearly summarizes the main change: sealing Noise frames in place inside the batch buffer.
Description check ✅ Passed The description matches the PR well and explains the in-place sealing, rollback behavior, and size guard changes.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch perf/noise-par

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@greptile-apps

greptile-apps Bot commented Jul 26, 2026

Copy link
Copy Markdown

Greptile Summary

The PR removes the scratch-buffer copy from inline Noise frame encryption.

  • Writes the predicted frame prefix and plaintext directly into the batch buffer before sealing the frame in place.
  • Rolls back partial output and rejects ciphertext whose size violates the AES-GCM tag-length contract.
  • Extracts a reusable framing-header helper and adds focused framing, rollback, offset, and custom-provider tests.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

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]
Loading

Reviews (2): Last reviewed commit: "test(socket): cover the ciphertext-size ..." | Re-trigger Greptile

greptile-apps[bot]
greptile-apps Bot previously approved these changes Jul 26, 2026

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Add test coverage for the ciphertext-size guard.

We’ve got the rollback and out_buf.len() - base != body_len checks in place, which is good — but none of the existing noise tests trip a provider that returns Ok with the wrong ciphertext size. Wire a fake set_crypto_provider backend 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

📥 Commits

Reviewing files that changed from the base of the PR and between 4d781fe and 1e5a19a.

📒 Files selected for processing (2)
  • src/socket/noise_socket.rs
  • wacore/noise/src/framing.rs

@github-actions

github-actions Bot commented Jul 26, 2026

Copy link
Copy Markdown

📦 Binary size report

Metric main PR Δ
bin size (stripped) 9.98 MiB 9.98 MiB +384 B (+0.00%) 🔺
bin .text 8.03 MiB 8.03 MiB +192 B (+0.00%) 🔺
bin allocated (text+data+bss) 9.98 MiB 9.98 MiB +56 B (+0.00%) 🔺
llvm-lines wacore 492,078 492,078 0
llvm-lines wacore copies 16,335 16,335 0
llvm-lines whatsapp-rust lib 720,529 720,635 +106 (+0.01%) 🔺
llvm-lines whatsapp-rust lib copies 22,768 22,772 +4 (+0.02%) 🔺
deps crates (Cargo.lock) 471 471 0
.text per crate
Crate main PR Δ
.text whatsapp_rust 1.82 MiB 1.82 MiB +217 B (+0.01%) 🔺
.text wacore 646.55 KiB 646.55 KiB 0
.text wacore_binary 89.34 KiB 89.34 KiB 0
.text wacore_libsignal 161.84 KiB 161.84 KiB 0
.text wacore_appstate 22.36 KiB 22.36 KiB 0
.text wacore_noise 21.60 KiB 21.60 KiB -1 B (-0.00%) 🔽
.text waproto 1.74 MiB 1.74 MiB 0
.text whatsapp_rust_sqlite_storage 514.91 KiB 514.91 KiB 0
.text whatsapp_rust_tokio_transport 39.91 KiB 39.91 KiB 0
.text whatsapp_rust_ureq_http_client 10.40 KiB 10.40 KiB 0
.text std 1.07 MiB 1.07 MiB -81 B (-0.01%) 🔽
.text other deps 1.89 MiB 1.89 MiB -11 B (-0.00%) 🔽

Baseline: 4d781fe87 (latest main run) · Head: 69da273a0 · Graphs

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.
@greptile-apps
greptile-apps Bot dismissed their stale review July 26, 2026 17:20

Dismissed because a newer commit was pushed; Greptile will re-review the current head.

@jlucaso1
jlucaso1 merged commit 7dad909 into main Jul 26, 2026
20 of 21 checks passed
@jlucaso1
jlucaso1 deleted the perf/noise-par branch July 26, 2026 17:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant