fix(prekey): remove the consumed one-time prekey atomically with the session flush - #722
Merged
Merged
Conversation
…session flush On an inbound pkmsg the promoted SessionRecord was put into the volatile SignalStoreCache while the consumed one-time prekey was deleted from the backend synchronously. The cache flush has no prekey lane and runs after the payload loop, so the prekey was durably gone while the new session was still only in memory. A crash or a failed flush in that window lost both, leaving a redelivered pkmsg permanently undecryptable (InvalidPreKeyId). Buffer the consumed prekey id in the cache and delete it from the backend only in the same flush that persists the sessions, after the session batch is durable. This mirrors WAWebSignalProtocolStoreUnifiedApi, which buffers removePreKey and commits it alongside the session put under one storage lock, and removes the per-pkmsg synchronous backend round-trip.
📝 WalkthroughSummary by CodeRabbit
WalkthroughDecryption reports consumed_prekey_id; decrypt callsites buffer consumed one-time prekeys into SignalStoreCache per-session. SignalStoreCache defers backend deletion until a flush persists the owning session and then deletes buffered prekeys under the sessions lock; clear() drops buffered entries. ChangesPre-key removal deferral
Sequence DiagramsequenceDiagram
participant MessageHandler as Message handler
participant PreKeyAdapter as PreKeyAdapter
participant Cache as SignalStoreCache
participant Backend as Backend store
MessageHandler->>PreKeyAdapter: buffer_consumed_prekey(prekey_id, session_addr)
PreKeyAdapter->>Cache: remove_prekey(prekey_id, session_addr)
Cache->>Cache: store in removed_prekeys
MessageHandler->>Cache: flush()
Cache->>Backend: persist session batch
Backend-->>Cache: success
Cache->>Backend: remove_prekey(prekey_id) for each buffered deletable
Backend-->>Cache: success
Cache->>Cache: remove entries from removed_prekeys
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
🚥 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 |
![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: 6be5678528
ℹ️ 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".
Benchmark Results1 regression(s) detected (>2% threshold):
66 unchanged benchmark(s)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
The consumed-prekey buffer was drained in a separate lock scope, acquired after the sessions lock was already released. Flushes are not serialized against decrypts, so in the gap between releasing the sessions lock and taking removed_prekeys a concurrent sender could promote its session into the volatile cache and buffer its consumed prekey, which the in-flight flush would then durably delete while that session was still only in memory. A crash before the second sender's own flush lost its session while its prekey was gone, making a redelivered pkmsg undecryptable (InvalidPreKeyId) again. Move the prekey drain inside the sessions lock scope so the session commit and the prekey delete are atomic against concurrent buffering. A sender that buffers after the session batch is snapshotted is blocked from inserting into removed_prekeys until the lock frees, so a flush can never delete a prekey whose session it did not also persist. This matches WAWebSignalProtocolStoreUnifiedApi holding its cache mutexes across the whole flush. Identity and sender-key lanes keep independent scopes.
The flush drained all buffered consumed-prekey IDs even when a dirty (promoted-but-not-yet-durable) session was skipped because a concurrent reader had it checked out. process_session_enc_batch drops the per-address session lock before the post-loop flush, so a concurrent send/retry can check out the just-promoted session; the flush then deleted its one-time prekey while the session stayed volatile, recreating the crash-orphan window. Track whether any dirty session was checked out during the session batch and defer the whole prekey drain to a later flush when so (IDs stay buffered, bounded by consumed prekeys between flushes). Adds a regression test.
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/src/store/signal_cache.rs`:
- Around line 571-585: The current flush logic uses dirty_keys (u32 IDs) and a
single dirty_session_checked_out flag so a CheckedOut session
(SessionEntry::CheckedOut) prevents deleting all buffered prekey removals, which
can leak other addresses' consumed prekeys; change the approach to track
removals per-address/session: when iterating dirty_keys, consult
state.cache.get(address.as_ref()) and collect deletable addresses into a
dedicated list (e.g., deletable_prekeys) only for SessionEntry::Present records
you successfully serialized (the ones you push into batch), skip or defer
entries that are CheckedOut without setting a global suppression flag, and then
delete/drain only the deletable_prekeys from the removal buffer (instead of
calling clear()), ensuring clear() is not used to drop the whole buffer
prematurely and that ownership ties each buffered removal to its address/session
so only safe removals are applied.
🪄 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: e743eef6-a550-4657-8e57-5fc7d7cfd215
📒 Files selected for processing (1)
wacore/src/store/signal_cache.rs
The flush deferred the entire consumed-prekey buffer whenever any dirty session happened to be checked out, so a prekey whose session WAS persisted this flush stayed buffered. A later clear() (disconnect) drops the buffer, leaving that prekey durable forever while its session is live: a one-time prekey that should have been consumed leaks. Key the buffer by the session address that consumed each prekey, and have message_decrypt report the consumed id instead of removing it internally (DecryptionResult carries consumed_prekey_id; the receive path buffers it with the peer address via PreKeyAdapter::buffer_consumed_prekey). The flush now deletes a prekey as soon as ITS own session is durable and defers only the prekeys of sessions still checked out, never holding back the prekeys of sessions it did persist. clear() still drops the whole buffer: every still-buffered prekey has a volatile session at that point, so the durable prekey is kept so a redelivered pkmsg can rebuild. The drain stays under the sessions lock, so the session commit and prekey delete remain atomic against a concurrent decrypt that promotes a session and buffers its prekey, matching WAWebSignalProtocolStoreUnifiedApi bundling bulkPutSession with bulkRemovePreKey under one storage lock. remove_pre_key becomes the plain immediate-removal primitive (no longer on the consume path).
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 48330559e9
ℹ️ 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".
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/src/store/signal_cache.rs`:
- Around line 55-65: The long explanatory comment describing "Consumed one-time
pre-key IDs awaiting durable deletion" in signal_cache.rs should be replaced
with a tight "why" statement: compress to one or two sentences that state the
invariant and rationale (e.g., that pre-keys are buffered keyed by session
address so deletes occur only when the session is durably persisted to avoid
losing both session and prekey on crash). Apply the same concise rewrite pattern
to the other verbose comment blocks noted (around lines 611-633 and 912-918),
keeping only the essential invariant and rationale, removing flow/mechanics
details.
🪄 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: 3e4fc3b9-2b7c-46e6-b6ca-6cb9d231bfd2
📒 Files selected for processing (5)
src/features/signal.rssrc/message.rssrc/store/signal_adapter.rswacore/libsignal/src/protocol/session_cipher.rswacore/src/store/signal_cache.rs
Resolve the signal_cache.rs conflict with #721's amortized eviction: keep main's high_watermark eviction and the read-path evict calls, keep this branch's per-address consumed-prekey buffer and selective flush drain, and keep both test modules (consumed_prekey_atomicity_tests + eviction_tests). Also fold in the Codex P1 fix: gate each buffered prekey on membership in the batch THIS flush persisted, not on the live cache entry. A session that was checked out, deleted, or dropped by a concurrent clear() before its buffer insert landed is absent from the batch, so its prekey is deferred instead of durably deleted with no session behind it (which would make a redelivered pkmsg permanently undecryptable). This also fixes the symmetric eviction leak the entry-based predicate had (a persisted-then-evicted session read as None). Adds prekey_without_a_persisted_session_survives_flush as the regression lock and tightens the now-verbose buffer/flush comments.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c92f59058a
ℹ️ 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 batch-membership predicate only deleted a buffered prekey when ITS session was persisted by the same flush. If a concurrent flush ran between the decrypt's store_session and the receive path's buffer_consumed_prekey, it persisted the session (empty buffer) and cleared its dirty bit; the prekey then buffered for an address that never re-enters a batch, so it stayed durable forever (one-time prekey leak). Decide durability per session instead: Present (clean at drain) is persisted, so delete; CheckedOut is the still-volatile promoted copy, so defer; an absent/deleted/evicted/cleared entry is ambiguous, so fall back to backend.has_session. This deletes prekeys of already-durable sessions (the leak), still defers prekeys whose session never reached the backend (no InvalidPreKeyId on redelivery), and handles a durable-but-evicted session. The backend probe only runs for the rare non-Present cases; the hot path stays in memory. Adds prekey_buffered_after_session_already_durable_is_deleted as the regression lock.
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)
wacore/src/store/signal_cache.rs (1)
748-755:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMake
clear()drop sessions and buffered prekeys atomically.Line 749 clears
sessionsand releases that guard before Line 755 clearsremoved_prekeys. A concurrent decrypt canput_session()in that gap and then have its buffered prekey removal erased byclear(), so the next flush can persist the session while the consumed one-time prekey stays durable forever. Clear those two structures under the same lock order thatflush()already uses.Suggested fix
pub async fn clear(&self) { - self.sessions.lock().await.clear(); + { + let mut sessions = self.sessions.lock().await; + let mut removed_prekeys = self.removed_prekeys.lock().await; + sessions.clear(); + removed_prekeys.clear(); + } self.identities.lock().await.clear(); self.sender_keys.lock().await.clear(); - // Drop buffered prekey removals together with the volatile sessions they - // belong to: the promoted session is gone, so the still-durable prekey - // must stay so a redelivered pkmsg can rebuild the session. - self.removed_prekeys.lock().await.clear(); }🤖 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 `@wacore/src/store/signal_cache.rs` around lines 748 - 755, The clear() method currently unlocks sessions before cleared removed_prekeys which allows a race; fix it by acquiring the same lock ordering flush() uses and clearing sessions and removed_prekeys while holding both guards atomically — specifically, in clear() obtain the sessions lock and then the removed_prekeys lock (matching flush()), clear both collections while both guards are held, then release; keep clearing identities and sender_keys as before but ensure sessions + removed_prekeys are cleared together to prevent a concurrent put_session() from winning between the two clears.
🤖 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 `@wacore/src/store/signal_cache.rs`:
- Around line 748-755: The clear() method currently unlocks sessions before
cleared removed_prekeys which allows a race; fix it by acquiring the same lock
ordering flush() uses and clearing sessions and removed_prekeys while holding
both guards atomically — specifically, in clear() obtain the sessions lock and
then the removed_prekeys lock (matching flush()), clear both collections while
both guards are held, then release; keep clearing identities and sender_keys as
before but ensure sessions + removed_prekeys are cleared together to prevent a
concurrent put_session() from winning between the two clears.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: df8de63b-1766-492c-9a59-03d7f48de698
📒 Files selected for processing (1)
wacore/src/store/signal_cache.rs
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
On an inbound pkmsg decrypt,
message_decrypt_prekeypromotes the newSessionRecordand stores it via the session adapter, which is cache-only: the record lands in the volatileSignalStoreCacheand is not durable until the post-loopflush_signal_cacheruns. Right after that, the consumed one-time prekey is removed viaPreKeyStore::remove_pre_key, which on our adapter went straight to the backend DB synchronously.The cache flush only had session, identity and sender-key lanes, so there was a window where:
A crash or a failed flush in that window loses both. The peer can redeliver the same pkmsg (it does on reconnection), and decryption then fails with
InvalidPreKeyIdbecause the prekey is gone and there is no session to fall back on, leaving that message permanently undecryptable.Change
Make the consumed-prekey removal durable atomically with the session, matching
WAWebSignalProtocolStoreUnifiedApi, which buffersremovePreKeyin memory and commits it alongside the session put under one storage lock.SignalStoreCachegains a small buffer of consumed prekey ids.remove_pre_keynow buffers the id instead of deleting from the backend.flushdeletes the buffered prekeys only after the session batch is committed, draining the buffer under its lock and clearing ids only once every backend delete succeeds (a failed flush keeps them for the next attempt).clear(disconnect) drops the buffer together with the volatile sessions it belongs to, so a still-durable prekey survives and a redelivered pkmsg can rebuild the session.This also removes the per-pkmsg synchronous backend round-trip for the prekey delete; it now rides the existing batched flush.
The
PreKeyStoretrait is unchanged, so other callers are unaffected. Only the decrypt-path session adapter changes behavior.Tests
consumed_prekey_stays_durable_until_session_flush: after decrypt but before flush, the prekey is still in the backend and the session is not yet durable; after flush, the session is durable and the prekey is deleted.clear_before_flush_keeps_prekey_so_pkmsg_can_rebuild: a disconnect before flush drops the volatile session and keeps the still-durable prekey, and a later flush of the now-empty buffer does not delete it.failed_session_flush_does_not_delete_prekey: when the session write fails, the flush aborts before the prekey lane, the prekey stays durable, and the buffered id remains for the next attempt.remove_pre_key_defers_backend_delete_to_flush: the session adapter'sremove_pre_keydoes not touch the backend; the prekey is only deleted by the subsequent flush.Performance
Removes one synchronous backend delete (and its SQLite transaction) per inbound pkmsg; the deletion now batches into the existing flush.