feat: rotate signed pre-key on a cadence (WA Web RotateKeyJob) - #968
Conversation
The signed pre-key minted at pairing was otherwise permanent. Mirror WA Web's RotateKeyJob: periodically generate a fresh signed pre-key, upload it via an `encrypt` IQ (`<rotate><skey/></rotate>`), and retain the previous keys so prekey messages already in flight against a rotated-out signed pre-key still decrypt. - Multi-key load fallback: `Device::load_signed_prekey` now falls back to the backend `signed_prekeys` table for a non-current id (previously it only ever returned the current field, so any rotated-out key failed). - New `last_signed_pre_key_rotation_ms` Device column (+ migration) drives a weekly cadence gate; devices upgraded in with it unset get a one-time baseline stamp so the first rotation is a full interval out, not on the next connect. - `SetSignedPreKey` / `SetSignedPreKeyRotationBaseline` DeviceCommands install the rotation atomically; the enum's hand-written Debug redacts key material (KeyPair omits Debug on purpose). - Retain the last 3 signed pre-keys, pruning the rest, to bound the decrypt window for delayed prekey messages. - Rotation is spawned post-login so a slow/failing encrypt IQ never delays connect; upload errors (WA Web 406/409/>=500) log and retry on a later connect rather than failing login. The rotation interval is the one value not grounded in the WA Web bundle (a persisted server-tuned job there); it is a documented, tunable default. Tests: should_rotate truth table + wrap-at-24-bit id; both command apply arms; load fallback for a rotated-out id; the <rotate> IQ node shape.
|
Warning Review limit reached
Next review available in: 3 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Repository UI (base), Organization UI (inherited) Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis PR adds signed pre-key rotation end to end: device state and persistence, backend retention for rotated keys, a rotate IQ, rotation timing and upload logic, and a post-login background task that triggers the flow. ChangesSigned Pre-Key Rotation
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant Persistence
participant Backend
participant Server
Client->>Persistence: read last_signed_pre_key_rotation_ms
alt baseline not set
Client->>Persistence: SetSignedPreKeyRotationBaseline(now)
else rotation due
Client->>Backend: stage outgoing and candidate signed pre-keys
Client->>Client: generate and sign new signed pre-key
Client->>Server: RotateSignedPreKeySpec upload
Server-->>Client: success or error
alt accepted
Client->>Persistence: SetSignedPreKey(new_id, key_pair, signature, rotation_ms)
Client->>Backend: prune older retained keys
else rejected
Client->>Backend: keep or discard staged candidate
end
end
Possibly related PRs
Suggested labels: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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/features/rotate_key.rs | New file implementing the full rotation pipeline; pruning now correctly uses skip(RETENTION - 1) = 2, keeping 3 total addressable keys as documented. Error-path ordering (stage before upload) and retry semantics are sound. |
| src/store/signal.rs | Both load_signed_prekey and contains_signed_prekey now fall back to the backend table for rotated-out ids, fixing the forward-secrecy gap where old id lookups always returned None. |
| wacore/src/store/commands.rs | Added SetSignedPreKey and SetSignedPreKeyRotationBaseline variants with hand-rolled Debug impl redacting private key material; apply arms correctly update all three key fields atomically. |
| wacore/src/store/device.rs | Adds last_signed_pre_key_rotation_ms field with serde(default); Device::new() correctly seeds it to now_millis() so new devices skip the baseline path and get a full weekly interval before first rotation. |
| wacore/src/iq/prekeys.rs | RotateSignedPreKeySpec reuses the existing SignedPreKeyNode encoder for the rotate/skey IQ, ensuring the wire format stays consistent with the upload path; test verifies 3/32/64-byte child layout. |
| storages/sqlite-storage/src/sqlite_store.rs | Adds last_signed_pre_key_rotation_ms to DeviceRow, upsert, and load path; store_signed_prekey already uses INSERT OR REPLACE semantics so retry writes of old_id are idempotent. |
| src/client/node_io.rs | Spawns maybe_rotate_signed_pre_key as a detached task post-login with a generation guard to drop stale tasks; errors are logged but never fail login. |
Sequence Diagram
%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant C as Client (post-login)
participant L as rotation_lock
participant B as Backend (SQLite)
participant S as WhatsApp Server
C->>L: try_lock() → guard
C->>B: load_signed_prekey(new_id)
alt already staged (retry path)
B-->>C: Some(kp1 bytes) → reuse
else fresh rotation
C->>C: generate kp1
C->>B: store_signed_prekey(new_id, kp1)
end
C->>B: store_signed_prekey(old_id, old_kp)
C->>S: IQ set encrypt/rotate skey(new_id, kp1)
alt Server accepts (200)
S-->>C: Ok(())
C->>C: process_command(SetSignedPreKey)
C->>B: flush() → persist new device state
C->>B: remove_signed_prekey(new_id)
C->>B: "load_all + prune to RETENTION-1=2 entries"
else Server rejects 406/409
S-->>C: Err(406 or 409)
C->>B: remove_signed_prekey(new_id)
C-->>C: return Ok() — retry on next connect
else Transport error / 5xx
S-->>C: Err(other)
C-->>C: return Ok() — keep staged key, retry
end
L-->>C: drop guard (lock released)
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant C as Client (post-login)
participant L as rotation_lock
participant B as Backend (SQLite)
participant S as WhatsApp Server
C->>L: try_lock() → guard
C->>B: load_signed_prekey(new_id)
alt already staged (retry path)
B-->>C: Some(kp1 bytes) → reuse
else fresh rotation
C->>C: generate kp1
C->>B: store_signed_prekey(new_id, kp1)
end
C->>B: store_signed_prekey(old_id, old_kp)
C->>S: IQ set encrypt/rotate skey(new_id, kp1)
alt Server accepts (200)
S-->>C: Ok(())
C->>C: process_command(SetSignedPreKey)
C->>B: flush() → persist new device state
C->>B: remove_signed_prekey(new_id)
C->>B: "load_all + prune to RETENTION-1=2 entries"
else Server rejects 406/409
S-->>C: Err(406 or 409)
C->>B: remove_signed_prekey(new_id)
C-->>C: return Ok() — retry on next connect
else Transport error / 5xx
S-->>C: Err(other)
C-->>C: return Ok() — keep staged key, retry
end
L-->>C: drop guard (lock released)
Reviews (5): Last reviewed commit: "fix(rotate-key): narrow discard to 406/4..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![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: 7a02c6219e
ℹ️ 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".
📦 Binary size report
.text per crate
Top movers (cargo-bloat attribution)
Baseline: |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@src/client/node_io.rs`:
- Around line 752-767: The detached signed-pre-key rotation task can run with a
stale connection and still call maybe_rotate_signed_pre_key() after a newer
connection has taken over. Add a generation re-check inside the spawned async
block in the rotate-key job, matching the pattern used by the background-queries
task, and use the same client_clone/rotate_client connection generation guard
before attempting rotation so outdated tasks exit without uploading a duplicate
key.
In `@src/features/rotate_key.rs`:
- Around line 90-164: The signed pre-key rotation path in rotate_key.rs advances
persistence too early by writing the new rotation state in
DeviceCommand::SetSignedPreKey before the upload succeeds. Update
rotate_signed_pre_key so the rotation timestamp and pruning only happen after
execute(RotateSignedPreKeySpec::new(...)) returns Ok, or track the last
successfully uploaded id separately. Use the existing persistence_manager,
process_command, flush, and backend pruning flow to keep the server-acknowledged
key retryable on later connects.
🪄 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 (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 4c6eb98b-4d99-455a-a321-57168c6b64a2
📒 Files selected for processing (11)
src/client/node_io.rssrc/features/mod.rssrc/features/rotate_key.rssrc/store/signal.rsstorages/sqlite-storage/migrations/2026-07-03-000000_add_signed_prekey_rotation/down.sqlstorages/sqlite-storage/migrations/2026-07-03-000000_add_signed_prekey_rotation/up.sqlstorages/sqlite-storage/src/schema.rsstorages/sqlite-storage/src/sqlite_store.rswacore/src/iq/prekeys.rswacore/src/store/commands.rswacore/src/store/device.rs
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
3 issues found across 11 files
Confidence score: 2/5
- In
src/features/rotate_key.rs,last_signed_pre_key_rotation_msis updated before the rotate IQ upload succeeds, so a failed upload can make later connects think rotation is fresh and skip retries for a full interval; this can leave clients on stale signed pre-keys longer than intended — move the timestamp update to only after a confirmed successful upload (or roll it back on failure) before merging. - In
src/client/node_io.rs, the detached post-login rotate task lacks a generation/connection guard, so a stale task can still rotate and persist keys after the active connection has changed, creating cross-session state races and unexpected key state — add the same in-task generation check used elsewhere in this flow before merge. - In
src/features/rotate_key.rs, retention currently keeps one extra historical signed pre-key, expanding the decrypt window beyond the intendedcurrent + 2 oldpolicy; this weakens key-rotation strictness versus design expectations — align the retention filter/count logic with the intended limit and verify with a boundary test.
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
… retention Address review findings on the signed pre-key rotation: - Upload-first: only retain the old key, promote the new one, stamp the cadence, and prune AFTER the server accepts the rotate IQ. Previously the cadence was stamped before upload, so a failed upload both skipped retries for a full interval and, across repeated failures, could prune the key the server was still advertising — breaking new prekey sessions. - Re-check the connection generation inside the detached rotation task, not just before spawning it, so a stale task from a superseded connection can't upload a duplicate key (matches the background-queries guard). - Prune to SIGNED_PRE_KEY_RETENTION total addressable keys: the current key lives in the device field, so keep only RETENTION-1 rotated-out keys in the backend (was keeping RETENTION, giving 4 total instead of 3). Dropped the unreachable new_id guard (new_id is never in the backend table). - Make contains_signed_prekey consistent with load_signed_prekey's backend fallback, so a rotated-out id doesn't read as absent to a gated caller. Extends the load-fallback test to cover contains_signed_prekey.
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/features/rotate_key.rs (1)
138-154: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winDon’t abort promotion after WhatsApp accepts the new key.
At this point the server can advertise
new_id; if retainingold_idfails, the?exits beforeSetSignedPreKey, leaving local state on the old key and breaking new prekey sessions. Keep going with promotion and log the old-key retention failure, or retainold_idbefore the upload.Proposed minimal fix
- backend - .store_signed_prekey(old_id, &old_record.encode_to_vec()) - .await - .map_err(|e| anyhow::anyhow!("failed to retain old signed pre-key: {e}"))?; + if let Err(e) = backend + .store_signed_prekey(old_id, &old_record.encode_to_vec()) + .await + { + log::warn!( + "failed to retain old signed pre-key {old_id}: {e}; \ + continuing with server-accepted signed pre-key {new_id}" + ); + }🤖 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/features/rotate_key.rs` around lines 138 - 154, In rotate_key::rotate_key, the old signed pre-key retention failure is currently propagated before promotion, which can stop local state from moving to the new key after the server has accepted it. Update the flow around backend.store_signed_prekey and the subsequent DeviceCommand::SetSignedPreKey / persistence_manager.flush sequence so promotion still happens even if retaining old_id fails, and only log or otherwise record that retention error instead of returning early.
🤖 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/features/rotate_key.rs`:
- Around line 138-154: In rotate_key::rotate_key, the old signed pre-key
retention failure is currently propagated before promotion, which can stop local
state from moving to the new key after the server has accepted it. Update the
flow around backend.store_signed_prekey and the subsequent
DeviceCommand::SetSignedPreKey / persistence_manager.flush sequence so promotion
still happens even if retaining old_id fails, and only log or otherwise record
that retention error instead of returning early.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: e306a939-4fda-4533-9ef8-4b81581591f4
📒 Files selected for processing (3)
src/client/node_io.rssrc/features/rotate_key.rssrc/store/signal.rs
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 932fb1d50e
ℹ️ 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".
Once the server accepts the new signed pre-key, promotion must proceed even if retaining the outgoing key in the backend fails; propagating that error left local state on the old key while the server advertised the new one, breaking new prekey sessions. Log and continue instead.
There was a problem hiding this comment.
2 issues found across 3 files (changes from recent commits).
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
Close the partial-failure family on the rotate path (P0). The candidate signed pre-key is now written to the backend table BEFORE the upload and reused verbatim on retry, so no partial failure can strand the key the server ends up advertising: - Ambiguous transport error (timeout/disconnect after the server may have accepted new_id): the staged key stays decryptable via the load fallback, and the retry re-uploads THIS exact key instead of minting a fresh one under the same id (which would overwrite the accepted key's private half). - Post-acceptance persistence failure: promotion/flush can fail without stranding decryptability, since new_id is already durable in the backend. - Definitive rejection (406/409/>=500): current key stays in place; cadence is not advanced and no key is pruned, so a later connect retries. On success the redundant staged copy of new_id is dropped (it now lives in the device field) before pruning to RETENTION total addressable keys.
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
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 `@src/features/rotate_key.rs`:
- Around line 86-92: The signed-pre-key rotation in rotate_signed_pre_key can
still race because the generation check in maybe_rotate_signed_pre_key only
gates entry, not the full rotate/upload/prune flow. Add a small single-flight
lock or otherwise extend the guard around the entire rotation path, including
rotate_signed_pre_key and its upload/prune steps, so older post-login tasks
cannot overlap with newer ones.
🪄 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 (base), Organization UI (inherited)
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 984b20b0-f5ed-41c0-a9e8-6b3946a3c5b0
📒 Files selected for processing (1)
src/features/rotate_key.rs
There was a problem hiding this comment.
2 issues found across 1 file (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
Two converging review findings on the rotation flow: - Single-flight: the post-login generation check only gates entry, so two overlapping tasks (reconnect churn) could both pass it and race the stage/upload/promote sequence — two keys minted for the same id, one overwriting the other. Add a dedicated try_lock held across the whole maybe_rotate flow; a second concurrent caller simply skips. - Retain the outgoing key BEFORE upload (required), not best-effort after acceptance. With the candidate already staged pre-upload, the old key can also be persisted pre-upload: on failure we abort before sending anything (current key fully intact), and once the server accepts, both the old (retained) and new (staged) keys are already durable, so promotion cannot strand either. Resolves the earlier best-effort-vs-required tension.
There was a problem hiding this comment.
0 issues found across 3 files (changes from recent commits).
Requires human review: This PR changes cryptographic key management and database migrations, which require human review for security and correctness.
Re-trigger cubic
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ef6610feb4
ℹ️ 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".
The reuse-on-retry staging path would resend a server-rejected key forever: a 4xx (WA Web 406 = bad key, 409 = validation fail) is a deterministic rejection of that specific key, so reusing the staged candidate wedges rotation permanently. Drop the staged row on a 4xx so the next attempt mints a fresh candidate; keep it only for transient 5xx and ambiguous transport failures, where the server may have accepted the same key.
There was a problem hiding this comment.
2 issues found across 1 file (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
Two refinements to the reject-path discard: - Narrow the discard trigger from any 4xx to exactly WA Web's deterministic key-rejection codes (406 bad key, 409 validation fail). Other 4xx (rate limits, transient auth) are retryable; discarding+reminting there only churns key ids without improving success. - Make the discard remove REQUIRED: propagate a remove failure instead of logging it, so a failed cleanup can't silently leave the rejected key staged and re-wedge rotation on the next attempt.
There was a problem hiding this comment.
0 issues found across 1 file (changes from recent commits).
Requires human review: This PR adds a new feature for signed pre-key rotation with database migration, core protocol changes, and concurrency handling. High-impact due to security-sensitive key management and schema migration.
Re-trigger cubic
2 participants
What
The signed pre-key minted at pairing was otherwise permanent — a real forward-secrecy gap. This lands WhatsApp Web's
RotateKeyJob: periodically generate a fresh signed pre-key, upload it, and retain the previous keys so prekey messages already in flight against a rotated-out signed pre-key still decrypt. Follow-up to the merged parity PR #965.The core safety fix
Device::load_signed_prekey(src/store/signal.rs) previously returned a record only whenid == the current field— it ignored the existingsigned_prekeysbackend table entirely. So rotating the single key in place would make any in-flight prekey message naming the old id fail withInvalidSignedPreKeyId. This PR adds the backend fallback for non-current ids, which is what makes any rotation safe. (The table + CRUD already existed; they were just never consulted.)Implemented (with tests)
getSignedPreKeyByIdsrc/store/signal.rsRotateKeyJobcadencesrc/features/rotate_key.rsSetSignedPreKey/SetSignedPreKeyRotationBaselinecommands (atomic install;Debugredacts key material)wacore/src/store/commands.rsrotateSignedPreKey/putSignedPreKeyssrc/features/rotate_key.rs<iq xmlns=encrypt type=set><rotate><skey/></rotate>upload IQ (reuses the upload path's<skey>encoder)WAWebRotateKeyJobwacore/src/iq/prekeys.rslast_signed_pre_key_rotation_ms(column + migration)device.rs, sqlite schema/migrationRotateKeyJoberror laddersrc/client/node_io.rsWire format grounded in the captured bundle (
WAWebRotateKeyJob):Retention & cadence
RotateKeyJobis a persisted, server-tuned background job — so it's a documented, easily-tunable policy default. A device upgraded in with the field unset (0) gets a one-time baseline stamp so its first rotation lands a full interval out, not immediately on the next connect.Validation
cargo build --all✅,cargo clippy --all --tests✅ clean,cargo fmt --all.should_rotatetruth table (incl.last==0and boundary/clock-skew),next_signed_pre_key_idwrap at the 24-bit border, bothDeviceCommandapply arms, the load fallback for a rotated-out id, and the<rotate>IQ node shape (3/32/64-byte children). sqlite-storage suite 47/47 (migration path included).Not in scope
No digest-key re-trigger on a 409 (a warn is complete behavior; the rotation is persisted locally and retried). No change to the pairing-time key generation.
Generated by Claude Code