From 6d001d52321372cb9fff0ada9ae2329dbcf9d8d3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 14 Jul 2026 15:42:32 -0300 Subject: [PATCH 1/3] perf(signal): lease outbound counters in batches instead of flushing every send MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every send used to serialize the whole SessionRecord and commit a storage transaction before the stanza could hit the wire, because outbound message keys/IVs derive deterministically from the sender-chain counter and a crash between wire-write and persistence would re-derive a spent (key, IV) pair. At 100k messages that was ~229 MB of protobuf encoding and one SQLite transaction per message on the hot path. The durable requirement is not the ratchet state itself — it is the promise that no counter ever repeats. That promise can be bought in batches: the record now leases SENDER_CHAIN_RESERVATION_BATCH counters ahead (reservedSenderChainIndex, a record-level field hand-encoded next to the RecordStructure fields, since the whatspec proto cannot carry local fields), and deserialize fast-forwards the loaded chain past the whole lease, making every possibly-spent counter underivable after a reload. Send-path effect: only the send that raises the lease (or advances a group sender-key chain, which has no lease yet) still flushes synchronously before the wire; every lease-covered send just schedules the same coalesced write-behind the receive path uses. Steady-state ping-pong (counter 0 of a freshly ratcheted chain every time) never re-raises the lease, so the per-message serialize + transaction disappears entirely from that profile. Safety interleavings covered by tests: - crash/reload mid-lease resumes past the ceiling; the peer decrypts across the burned gap (bounded, well under MAX_FORWARD_JUMPS) - a DH ratchet whose new chain never reached storage is unrecoverable by construction (fresh random ephemeral), and the reloaded old chain resumes past its own lease - archived-state promotion burns the lease into the promoted chain (promote_state); freshly ratcheted states reset it (promote_fresh_state) - a raised lease gates the wire until a flush SUCCEEDS (failed flushes keep the gate closed; checked-out sessions stay gated across a flush) - group sender-key encrypts still gate the wire, but group decrypt dirtiness no longer forces a sync flush onto unrelated DM sends - legacy records (no lease field) load with a zero reservation; the encoding round-trips and rejects implausible (corrupt) reservations Old readers skip the unknown record field, but a downgrade after a crash would ignore the lease — release notes should flag that. --- src/client/adapters.rs | 23 +- src/features/signal.rs | 4 +- src/send/mod.rs | 23 +- src/signal_flush.rs | 43 +- tests/e2e/tests/session_reuse.rs | 121 +++-- wacore/libsignal/src/protocol/consts.rs | 16 + wacore/libsignal/src/protocol/group_cipher.rs | 5 + wacore/libsignal/src/protocol/sender_keys.rs | 26 +- wacore/libsignal/src/protocol/session.rs | 7 +- .../libsignal/src/protocol/session_cipher.rs | 7 + .../libsignal/src/protocol/state/session.rs | 229 +++++++- wacore/libsignal/tests/counter_lease.rs | 496 ++++++++++++++++++ wacore/src/store/signal_cache.rs | 188 ++++++- 13 files changed, 1124 insertions(+), 64 deletions(-) create mode 100644 wacore/libsignal/tests/counter_lease.rs diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 380e9ae54..10ae56ca2 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -72,7 +72,8 @@ impl Client { /// returning once the flush completes (or fails). /// /// The live receive path schedules a coalesced flush (see `signal_flush.rs`) - /// instead of writing through (sends flush synchronously). On success the + /// instead of writing through, and lease-covered sends do the same (only a + /// send that raises its counter lease flushes synchronously). On success the /// backend normally trails the cache by about the coalescing window, but /// that is not a hard wall-clock bound — the timer can slip under runtime /// starvation and the flush can wait on locks or slow/failing storage (a @@ -160,6 +161,26 @@ impl Client { self.flush_signal_cache().await } + /// Pre-wire durability gate for the send path. Flushes synchronously only + /// when an outbound crypto advance actually demands it — a raised session + /// counter lease or a sender-key chain advance not yet persisted (see + /// `SignalStoreCache::needs_pre_wire_flush`). Otherwise the dirty state is + /// covered by an existing durable lease, so it only needs to land + /// eventually: it rides the same coalesced write-behind as the receive + /// path instead of paying a serialize + storage transaction per message. + /// A failure must abort the send — transmitting a ciphertext whose lease + /// could not be persisted reintroduces the counter-reuse window. + pub(crate) async fn persist_signal_state_pre_wire(&self) -> Result<(), anyhow::Error> { + if self.signal_cache.needs_pre_wire_flush().await { + return self.flush_signal_cache_batch_safe().await; + } + self.schedule_signal_flush( + self.connection_generation + .load(std::sync::atomic::Ordering::Acquire), + ); + Ok(()) + } + /// [`flush_signal_cache_batch_safe`](Self::flush_signal_cache_batch_safe) /// with error logging instead of propagation. pub(crate) async fn flush_signal_cache_batch_safe_logged( diff --git a/src/features/signal.rs b/src/features/signal.rs index 9172b9d98..cf6cd4354 100644 --- a/src/features/signal.rs +++ b/src/features/signal.rs @@ -70,7 +70,9 @@ impl<'a> Signal<'a> { .await?; drop(_guard); - self.client.flush_signal_cache_batch_safe().await?; + // Same pre-wire gate as the send path: the caller transmits these + // bytes, so a raised lease must be durable before they leave here. + self.client.persist_signal_state_pre_wire().await?; let (_, is_prekey, bytes) = wacore::send::extract_ciphertext(encrypted) .ok_or_else(|| SignalError::Unsupported("unexpected ciphertext variant".into()))?; diff --git a/src/send/mod.rs b/src/send/mod.rs index 845d34328..87b97f4f5 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -962,9 +962,9 @@ impl Client { .ensure_status_participants(prepared.node, &group_info) .await?; - // Persist the sender-key ratchet advance before the stanza hits the - // wire (same rule as the DM/group send path); a failure aborts the send. - self.flush_signal_cache_batch_safe().await?; + // Gate the stanza on the sender-key ratchet advance being durable + // (same rule as the DM/group send path); a failure aborts the send. + self.persist_signal_state_pre_wire().await?; let ack = if let Some(phash) = stanza .attrs() @@ -1482,14 +1482,15 @@ impl Client { .await? }; - // Persist the outbound ratchet advance BEFORE the stanza hits the wire - // (WA Web flushes the Signal store ahead of send). Reusing an outbound - // counter reuses its message key + IV, so the advance must be durable - // before anyone can act on the ciphertext — and a persistence failure - // must abort the send rather than transmit an advance we couldn't save. - // Only the receive path, where a lost advance re-derives forward, - // coalesces. - self.flush_signal_cache_batch_safe().await?; + // The outbound advance must be durable BEFORE the stanza hits the wire: + // reusing an outbound counter reuses its message key + IV. Counters are + // leased in batches (see `SessionRecord::reserve_sender_chain_counters`), + // so most sends are already covered by a durable lease and only + // schedule the coalesced write-behind; a send that raised the lease (or + // advanced a sender-key chain) flushes synchronously, and a persistence + // failure must abort the send rather than transmit an advance we + // couldn't save. + self.persist_signal_state_pre_wire().await?; let ack = if let Some(phash) = dm_phash && let Some(msg_id) = stanza_to_send diff --git a/src/signal_flush.rs b/src/signal_flush.rs index f50c79465..58bf0906f 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -7,13 +7,16 @@ //! window. //! //! Scope and durability model (deliberate, bounded): -//! - Only the receive path coalesces. A lost receive-side advance re-derives +//! - The receive path coalesces: a lost receive-side advance re-derives //! forward (the Double Ratchet receiving chain derives `CK_n → CK_n+1`), and //! a consumed one-time prekey stays buffered until its session is durable, so -//! a crash inside the window is recoverable. The SEND path flushes -//! synchronously before returning: reusing an outbound counter would reuse -//! its message key + IV, so that advance must be durable before `send_message` -//! reports success. +//! a crash inside the window is recoverable. The SEND path coalesces too +//! whenever its advance is covered by a durable counter lease (see +//! `SessionRecord::reserve_sender_chain_counters` and +//! `Client::persist_signal_state_pre_wire`); only a send that raises the +//! lease — or advances a group sender-key chain — still flushes +//! synchronously before the wire, because reusing an outbound counter would +//! reuse its message key + IV. //! - The offline drain, retry recovery, identity-change recovery and teardown //! keep their own synchronous flushes: those gate acks, receipts or //! follow-up reads on durability and are not routed here. @@ -284,6 +287,36 @@ mod tests { } } + /// The send-path gate: a raised counter lease flushes synchronously + /// before returning; lease-covered dirty state only schedules the + /// coalesced worker and still lands within its window. + #[tokio::test] + async fn pre_wire_gate_flushes_leases_synchronously_and_coalesces_the_rest() { + let client = crate::test_utils::create_test_client().await; + + // Lease-covered advance (no raised reservation): no synchronous write — + // the session is still absent from the backend when the gate returns. + let covered = dirty_session(&client, "15550003001"); + client.persist_signal_state_pre_wire().await.unwrap(); + assert!( + backend_session(&client, &covered).await.is_none(), + "a covered advance must not flush synchronously" + ); + // ...but it rides the coalescer and still becomes durable. + wait_for_backend_session(&client, &covered).await; + + // A raised lease must be durable when the call returns. + let addr = ProtocolAddress::new("15550003002".to_string(), 1.into()); + let mut record = SessionRecord::new_fresh(); + record.reserve_sender_chain_counters(0); + assert!(client.signal_cache.try_put_session(&addr, record).is_ok()); + client.persist_signal_state_pre_wire().await.unwrap(); + assert!( + backend_session(&client, &addr).await.is_some(), + "the leased session must be durable when the gate returns" + ); + } + /// A burst of requests rides one armed worker and persists every dirty /// entry written before it. #[tokio::test] diff --git a/tests/e2e/tests/session_reuse.rs b/tests/e2e/tests/session_reuse.rs index 9dc29c635..5e07dcf45 100644 --- a/tests/e2e/tests/session_reuse.rs +++ b/tests/e2e/tests/session_reuse.rs @@ -61,25 +61,31 @@ async fn durable_sender_chain_index( Ok(None) } -/// The outbound ratchet advance must be durable by the time `send_message` -/// returns: reusing an outbound counter reuses its message key + IV, so a crash -/// after a successful send must never leave the advance only in memory. This -/// reads the backend IMMEDIATELY after `send_message` (no wait for delivery, no -/// explicit settle): a coalesced send would still be inside its window and the -/// counter would be stale, so only the synchronous outbound flush passes. +/// The durable snapshot must always be able to resume PAST every counter that +/// may have hit the wire: reusing an outbound counter reuses its message key + +/// IV. Counters are leased in batches (`SENDER_CHAIN_RESERVATION_BATCH`): the +/// send that raises the lease flushes synchronously before the wire, and every +/// lease-covered send may defer its advance to the coalesced write-behind +/// because `SessionRecord::deserialize` fast-forwards the reloaded chain past +/// the whole lease. This reads the backend IMMEDIATELY after `send_message` +/// (no delivery wait, no settle): the resume position — what a crash restore +/// would actually use — must already cover every counter spent so far. #[tokio::test] -async fn test_outbound_ratchet_is_durable_when_send_returns() -> anyhow::Result<()> { +async fn test_durable_resume_position_always_covers_spent_counters() -> anyhow::Result<()> { let _ = env_logger::builder().is_test(true).try_init(); - let mut client_a = TestClient::connect("e2e_sig_durable_a").await?; - let mut client_b = TestClient::connect("e2e_sig_durable_b").await?; - let jid_a = client_a.jid().await; + let client_a = TestClient::connect("e2e_sig_durable_a").await?; + let client_b = TestClient::connect("e2e_sig_durable_b").await?; let jid_b = client_b.jid().await; let lid_b = client_b.client.get_lid(); - // Establish the outbound session A→B. - send_and_expect_text(&client_a.client, &mut client_b, &jid_b, "establish", 30).await?; - send_and_expect_text(&client_b.client, &mut client_a, &jid_a, "reply", 30).await?; + // The first send raises the lease, so its flush is synchronous: the + // durable resume position must already be past counter 0 the moment + // send_message returns, with no settle. + client_a + .client + .send_message(jid_b.clone(), e2e_tests::text_msg("establish")) + .await?; let backend_a = client_a.client.persistence_manager().backend(); let read_index = async |user: &str, server: &str| { @@ -92,26 +98,35 @@ async fn test_outbound_ratchet_is_durable_when_send_returns() -> anyhow::Result< _ => (jid_b.user.clone(), "c.us"), }; + // `durable_sender_chain_index` deserializes the stored record, which + // applies the crash-restore fast-forward: this IS the resume position. let mut last = read_index(&user, server) .await? - .expect("an outbound session must exist after the roundtrip"); + .expect("the lease raise must persist the session before the wire"); + let mut spent = 1u32; // counter 0 went out with "establish" + assert!( + last >= spent, + "resume position {last} must cover the {spent} spent counter(s)" + ); - // send_message returns only after the synchronous pre-wire flush, so the - // advanced counter is already durable — read it with no delivery wait and - // no settle. A coalesced (window-deferred) flush would leave it unchanged. + // Lease-covered sends may leave the durable snapshot trailing (that is + // the optimization), but the resume position must never fall behind the + // wire and never regress. for i in 0..3 { client_a .client .send_message(jid_b.clone(), e2e_tests::text_msg(&format!("m{i}"))) .await?; + spent += 1; let now = read_index(&user, server) .await? .expect("session persists across sends"); assert!( - now > last, - "send #{i} must persist the advanced sender-chain counter before returning \ - (durable {last} -> {now}); a coalesced send would leave it stale" + now >= spent, + "send #{i}: resume position {now} fell behind the {spent} spent counter(s); \ + a crash here would re-derive a (key, IV) pair" ); + assert!(now >= last, "resume position must never regress"); last = now; } @@ -120,26 +135,21 @@ async fn test_outbound_ratchet_is_durable_when_send_returns() -> anyhow::Result< Ok(()) } -/// A send whose outbound-ratchet persistence fails must abort BEFORE the stanza -/// reaches the wire: the flush precedes the send on the send path, so if the -/// advance cannot be stored, `send_message` returns `Err` and the peer receives +/// A send that RAISES the counter lease gates the wire on persisting it: if +/// the flush fails, `send_message` returns `Err` and the peer receives /// nothing. Otherwise a crash after a wire-committed send would leave the -/// advance only in memory and the next send would reuse that counter's key + IV. +/// lease only in memory and a reload would re-derive that counter's key + IV. +/// The first send on a fresh session always raises the lease, so the failure +/// is injected before it. #[tokio::test] -async fn test_send_aborts_before_wire_when_persist_fails() -> anyhow::Result<()> { +async fn test_send_aborts_before_wire_when_lease_persist_fails() -> anyhow::Result<()> { let _ = env_logger::builder().is_test(true).try_init(); - let mut client_a = TestClient::connect("e2e_sig_abort_a").await?; + let client_a = TestClient::connect("e2e_sig_abort_a").await?; let mut client_b = TestClient::connect("e2e_sig_abort_b").await?; - let jid_a = client_a.jid().await; let jid_b = client_b.jid().await; - // Establish the session both ways so the next A→B send is a steady-state - // encrypt (its only new durable write is the ratchet advance we fail). - send_and_expect_text(&client_a.client, &mut client_b, &jid_b, "establish", 30).await?; - send_and_expect_text(&client_b.client, &mut client_a, &jid_a, "reply", 30).await?; - - // Persisting the outbound advance now fails. + // Persisting the outbound lease now fails. client_a.backend.set_fail_session_writes(true); let writes_before = client_a.backend.session_batch_write_count(); // `send_node` resolves this before marshaling the node, so a still-pending @@ -154,13 +164,13 @@ async fn test_send_aborts_before_wire_when_persist_fails() -> anyhow::Result<()> .await; assert!( result.is_err(), - "send must fail when the ratchet advance cannot be persisted, got {result:?}" + "send must fail when the raised lease cannot be persisted, got {result:?}" ); - // The send reached the (failing) persistence step, proving the flush runs on - // the send path before the wire rather than being skipped or deferred. + // The send reached the (failing) persistence step, proving the lease flush + // runs on the send path before the wire rather than being skipped or deferred. assert!( client_a.backend.session_batch_write_count() > writes_before, - "the send path must attempt to persist the ratchet advance before the wire" + "the send path must attempt to persist the raised lease before the wire" ); // Deterministic: no `message` node was ever marshaled, so `send_node` (and // thus the wire) was never reached. `Ok(None)` == pending, sender still alive. @@ -198,6 +208,43 @@ async fn test_send_aborts_before_wire_when_persist_fails() -> anyhow::Result<()> Ok(()) } +/// The counterpart of the abort test: a send COVERED by an already-durable +/// lease does not depend on this flush for safety — a crash would reload the +/// durable snapshot and fast-forward past the whole lease, so its counter can +/// never be re-derived. Such a send must therefore succeed even while the +/// backend is refusing session writes (the advance lands later via the +/// coalescer's retry), instead of turning a storage hiccup into message loss. +#[tokio::test] +async fn test_lease_covered_send_survives_persist_failure() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut client_a = TestClient::connect("e2e_sig_covered_a").await?; + let mut client_b = TestClient::connect("e2e_sig_covered_b").await?; + let jid_a = client_a.jid().await; + let jid_b = client_b.jid().await; + + // Establish both ways: the lease raise happens (and is persisted) here. + send_and_expect_text(&client_a.client, &mut client_b, &jid_b, "establish", 30).await?; + send_and_expect_text(&client_b.client, &mut client_a, &jid_a, "reply", 30).await?; + + // Storage starts refusing session writes; the next send is lease-covered, + // so it must still deliver. + client_a.backend.set_fail_session_writes(true); + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "covered by the lease", + 30, + ) + .await?; + + client_a.backend.set_fail_session_writes(false); + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + /// Multiple sequential sends without a reply should all be delivered. #[tokio::test] async fn test_one_way_multiple_sends() -> anyhow::Result<()> { diff --git a/wacore/libsignal/src/protocol/consts.rs b/wacore/libsignal/src/protocol/consts.rs index 7ece062c7..92d118930 100644 --- a/wacore/libsignal/src/protocol/consts.rs +++ b/wacore/libsignal/src/protocol/consts.rs @@ -24,3 +24,19 @@ pub const MAX_SENDER_KEY_STATES: usize = 5; /// Eviction only triggers when buffer exceeds MAX_MESSAGE_KEYS + PRUNE_THRESHOLD, /// reducing O(n) drain() calls from every insert to once every PRUNE_THRESHOLD inserts. pub const MESSAGE_KEY_PRUNE_THRESHOLD: usize = 50; + +/// Sender-chain counters leased per durable reservation (see +/// `SessionRecord::reserve_sender_chain_counters`). Message keys and IVs are +/// derived deterministically from the counter, so an outbound counter must +/// never repeat across a crash; instead of persisting every advance before it +/// hits the wire, the record durably reserves this many counters ahead and a +/// reloaded snapshot fast-forwards past them. Bounds both the sync-flush +/// amortization (one per this many sends) and the worst-case counter gap a +/// receiver sees after a crash — keep it well under MAX_FORWARD_JUMPS. +pub const SENDER_CHAIN_RESERVATION_BATCH: u32 = 64; + +/// Upper bound for the reservation fast-forward on load. A legitimate lease +/// gap is < SENDER_CHAIN_RESERVATION_BATCH; anything past this ceiling means +/// a corrupt record, and refusing it keeps a bogus reserved index from +/// turning the load into an unbounded KDF loop. +pub const MAX_RESERVATION_FAST_FORWARD: u32 = MAX_FORWARD_JUMPS as u32; diff --git a/wacore/libsignal/src/protocol/group_cipher.rs b/wacore/libsignal/src/protocol/group_cipher.rs index 503701ace..c51dd7246 100644 --- a/wacore/libsignal/src/protocol/group_cipher.rs +++ b/wacore/libsignal/src/protocol/group_cipher.rs @@ -111,6 +111,11 @@ pub async fn group_encrypt( sender_key_state.set_sender_chain_key(next_sender_chain_key); + // Outbound advance: this iteration's (key, IV) must never be re-derivable, + // so the store must gate the ciphertext on durability. Decrypt-side + // advances stay ungated (they re-derive forward). + record.mark_wire_gated(); + sender_key_store .store_sender_key(sender_key_name, record) .await?; diff --git a/wacore/libsignal/src/protocol/sender_keys.rs b/wacore/libsignal/src/protocol/sender_keys.rs index 4b72fa4ff..00980b320 100644 --- a/wacore/libsignal/src/protocol/sender_keys.rs +++ b/wacore/libsignal/src/protocol/sender_keys.rs @@ -428,6 +428,12 @@ impl SenderKeyState { #[derive(Debug, Clone)] pub struct SenderKeyRecord { states: VecDeque, + /// An outbound chain advance not yet known durable. Sender-key message + /// keys/IVs derive deterministically from the iteration, so the advance + /// must reach storage before its ciphertext reaches the wire (unlike + /// decrypt advances, which re-derive forward). Transient — never + /// serialized; the store layer converts it into flush gating. + wire_gated: bool, } impl SenderKeyRecord { @@ -438,6 +444,7 @@ impl SenderKeyRecord { pub fn new_empty() -> Self { Self { states: VecDeque::with_capacity(consts::MAX_SENDER_KEY_STATES), + wire_gated: false, } } @@ -456,7 +463,24 @@ impl SenderKeyRecord { } states.push_back(SenderKeyState::from_protobuf(state)); } - Ok(Self { states }) + Ok(Self { + states, + wire_gated: false, + }) + } + + /// Flag an outbound chain advance; cleared by the store layer once it + /// owns the durability gate. + pub fn mark_wire_gated(&mut self) { + self.wire_gated = true; + } + + pub fn is_wire_gated(&self) -> bool { + self.wire_gated + } + + pub fn clear_wire_gated(&mut self) { + self.wire_gated = false; } pub fn sender_key_state(&self) -> Result<&SenderKeyState, InvalidSenderKeySessionError> { diff --git a/wacore/libsignal/src/protocol/session.rs b/wacore/libsignal/src/protocol/session.rs index 70d6a851f..04e95e054 100644 --- a/wacore/libsignal/src/protocol/session.rs +++ b/wacore/libsignal/src/protocol/session.rs @@ -136,7 +136,9 @@ async fn process_prekey_impl( new_session.set_local_registration_id(identity_store.get_local_registration_id().await?); new_session.set_remote_registration_id(message.registration_id()); - session_record.promote_state(new_session); + // Fresh random ratchet: no counter on this chain can have been spent, so + // the record's inherited lease must not be burned into it. + session_record.promote_fresh_state(new_session); let pre_keys_used = PreKeysUsed { pre_key_id: message.pre_key_id(), @@ -244,7 +246,8 @@ async fn process_prekey_bundle_inner( .save_identity(remote_address, their_identity_key) .await?; - session_record.promote_state(session); + // Fresh random ratchet: see promote_fresh_state. + session_record.promote_fresh_state(session); Ok(identity_change) } diff --git a/wacore/libsignal/src/protocol/session_cipher.rs b/wacore/libsignal/src/protocol/session_cipher.rs index a5de86ebd..d40ac43fa 100644 --- a/wacore/libsignal/src/protocol/session_cipher.rs +++ b/wacore/libsignal/src/protocol/session_cipher.rs @@ -242,6 +242,13 @@ async fn message_encrypt_inner( session_state.set_sender_chain_key(&next_chain_key)?; + // Counters are leased in batches so the send path only needs a durable + // flush when the lease runs out; a reload fast-forwards past the whole + // lease, so this counter can never be re-derived after a crash. + if chain_key.index() >= session_record.reserved_sender_chain_index() { + session_record.reserve_sender_chain_counters(chain_key.index()); + } + Ok(message) } diff --git a/wacore/libsignal/src/protocol/state/session.rs b/wacore/libsignal/src/protocol/state/session.rs index 03ddc7272..ce59bc395 100644 --- a/wacore/libsignal/src/protocol/state/session.rs +++ b/wacore/libsignal/src/protocol/state/session.rs @@ -438,6 +438,34 @@ impl SessionState { Ok(()) } + /// Advance the sender chain until its next counter is at least `target`, + /// discarding the intermediate message keys. Restores the no-counter-reuse + /// invariant when this state re-enters service under a durable reservation + /// (snapshot reload, archived-state promotion): every counter the lease may + /// have already spent becomes underivable. A state without a usable sender + /// chain is left untouched — it cannot encrypt, so it has nothing to burn. + pub(crate) fn fast_forward_sender_chain( + &mut self, + target: u32, + ) -> Result<(), SignalProtocolError> { + let Ok(mut chain_key) = self.get_sender_chain_key() else { + return Ok(()); + }; + if target.saturating_sub(chain_key.index()) > consts::MAX_RESERVATION_FAST_FORWARD { + return Err(SignalProtocolError::InvalidSessionStructure( + "reserved sender chain index implausibly far ahead", + )); + } + if chain_key.index() >= target { + return Ok(()); + } + while chain_key.index() < target { + chain_key = chain_key.next_chain_key()?; + } + self.set_sender_chain_key(&chain_key)?; + Ok(()) + } + pub fn get_message_keys( &mut self, sender: &PublicKey, @@ -632,10 +660,29 @@ impl From<&SessionState> for SessionStructure { } } +/// Record-level field number carrying the sender-chain counter reservation in +/// the serialized `RecordStructure`. The upstream (whatspec) proto cannot be +/// edited to add local fields, so the record encoder — already hand-rolled in +/// [`SessionRecord::serialize_into`] — writes it directly; the number sits far +/// above RecordStructure's fields (1, 2) so a future upstream addition cannot +/// collide. Standard unknown-field skipping keeps old readers compatible. +const RESERVED_SENDER_CHAIN_INDEX_FIELD: u32 = 100; + #[derive(Clone)] pub struct SessionRecord { current_session: Option, previous_sessions: Arc>, + /// Durability lease: ceiling (exclusive) of sender-chain counters this + /// record's durable snapshots may already have spent on the wire. Any + /// state entering service from such a snapshot must fast-forward its + /// sender chain here first — message keys and IVs are derived + /// deterministically from the counter, so re-deriving a spent counter + /// reuses a (key, IV) pair. + reserved_sender_chain_index: u32, + /// A reservation was raised but not yet durably flushed. While set, the + /// owning ciphertext must not reach the wire; the store layer transfers + /// this into its flush gating. Transient — never serialized. + pending_reservation: bool, } impl SessionRecord { @@ -643,6 +690,8 @@ impl SessionRecord { Self { current_session: None, previous_sessions: Arc::new(Vec::new()), + reserved_sender_chain_index: 0, + pending_reservation: false, } } @@ -650,9 +699,35 @@ impl SessionRecord { Self { current_session: Some(state), previous_sessions: Arc::new(Vec::new()), + reserved_sender_chain_index: 0, + pending_reservation: false, } } + pub fn reserved_sender_chain_index(&self) -> u32 { + self.reserved_sender_chain_index + } + + /// Lease a fresh batch of sender-chain counters after `spent_counter` was + /// issued past the current reservation. Marks the record pending: the + /// caller's ciphertext must not hit the wire until a flush persists the + /// raised ceiling. + pub fn reserve_sender_chain_counters(&mut self, spent_counter: u32) { + self.reserved_sender_chain_index = + spent_counter.saturating_add(consts::SENDER_CHAIN_RESERVATION_BATCH); + self.pending_reservation = true; + } + + pub fn has_pending_reservation(&self) -> bool { + self.pending_reservation + } + + /// The store layer takes ownership of the wire gate (it tracks the address + /// until a successful flush), so the transient flag is dropped here. + pub fn clear_pending_reservation(&mut self) { + self.pending_reservation = false; + } + pub fn deserialize(bytes: &[u8]) -> Result { use waproto::whatsapp::RecordStructureView; @@ -671,7 +746,7 @@ impl SessionRecord { .collect::>() .map_err(|_| InvalidSessionError("failed to decode archived session protobuf"))?; - Ok(Self { + let mut record = Self { current_session: view .current_session .as_option() @@ -680,7 +755,45 @@ impl SessionRecord { .map_err(|_| InvalidSessionError("failed to decode current session protobuf"))? .map(Into::into), previous_sessions: Arc::new(previous_sessions), - }) + reserved_sender_chain_index: Self::decode_reserved_index(bytes)?, + pending_reservation: false, + }; + + // This snapshot may predate sends its lease already covered; burn the + // leased counters before the chain can issue keys again. + if record.reserved_sender_chain_index > 0 + && let Some(state) = record.current_session.as_mut() + { + state.fast_forward_sender_chain(record.reserved_sender_chain_index)?; + } + + Ok(record) + } + + /// Extract the record-level reservation field from the serialized bytes. + /// The generated `RecordStructureView` skips fields it does not know, so + /// the local-only field is scanned out of the raw top-level stream here. + fn decode_reserved_index(bytes: &[u8]) -> Result { + use buffa::encoding::{Tag, WireType, decode_varint, skip_field}; + + let mut buf = bytes; + let mut reserved = 0u32; + while !buf.is_empty() { + let tag = Tag::decode(&mut buf) + .map_err(|_| InvalidSessionError("failed to decode session record protobuf"))?; + if tag.field_number() == RESERVED_SENDER_CHAIN_INDEX_FIELD + && tag.wire_type() == WireType::Varint + { + let value = decode_varint(&mut buf) + .map_err(|_| InvalidSessionError("invalid reserved sender chain index"))?; + reserved = u32::try_from(value) + .map_err(|_| InvalidSessionError("invalid reserved sender chain index"))?; + } else { + skip_field(tag, &mut buf) + .map_err(|_| InvalidSessionError("failed to decode session record protobuf"))?; + } + } + Ok(reserved) } /// If there's a session with a matching version and `alice_base_key`, ensures that it is the @@ -815,9 +928,34 @@ impl SessionRecord { self.promote_state(updated_session) } + /// Make `new_state` current. A promoted state may have been current under + /// this record's lease before (archived → promoted round trip), so its + /// sender chain is fast-forwarded past every counter the lease may have + /// spent. States whose chain was never current here (fresh ratchets) go + /// through [`Self::promote_fresh_state`] instead, which skips the burn. pub fn promote_state(&mut self, new_state: SessionState) { + self.archive_current_state_inner(); + let mut state = new_state; + if self.reserved_sender_chain_index > 0 + && let Err(e) = state.fast_forward_sender_chain(self.reserved_sender_chain_index) + { + // Only reachable with a corrupt reservation (gap beyond the + // ceiling); refuse to expose the chain rather than risk reuse. + log::error!("dropping promoted sender chain: {e}"); + state.session.sender_chain = None.into(); + } + self.current_session = Some(state); + } + + /// Make a freshly ratcheted state current. Its sender chain key material + /// was just generated from a fresh random ephemeral, so no counter on it + /// can ever have been spent: the inherited lease is meaningless for it and + /// is reset instead of burned. The first send re-reserves durably before + /// hitting the wire. + pub fn promote_fresh_state(&mut self, new_state: SessionState) { self.archive_current_state_inner(); self.current_session = Some(new_state); + self.reserved_sender_chain_index = 0; } fn archive_current_state_inner(&mut self) -> bool { @@ -883,8 +1021,15 @@ impl SessionRecord { }) .sum(); + let reserved = self.reserved_sender_chain_index; + let reserved_len = if reserved > 0 { + 2 + varint_len(reserved as u64) + } else { + 0 + }; + buf.clear(); - buf.reserve(current_len + previous_len); + buf.reserve(current_len + previous_len + reserved_len); if let Some(state) = &self.current_session && let Some(msg_len) = current_msg_len @@ -894,6 +1039,10 @@ impl SessionRecord { for (session, msg_len) in self.previous_sessions.iter().zip(previous_msg_lens) { write_len_delimited(2, session, msg_len, &mut cache, buf); } + if reserved > 0 { + Tag::new(RESERVED_SENDER_CHAIN_INDEX_FIELD, WireType::Varint).encode(buf); + encode_varint(reserved as u64, buf); + } } /// Estimated in-memory footprint proxy: the protobuf-encoded size of the @@ -1075,6 +1224,78 @@ mod tests { assert!(!state.has_usable_sender_chain().unwrap()); } + /// An archived state promoted back to current may have spent counters + /// under the record's lease while it was current before; the promotion + /// must burn the whole lease into its chain. + #[test] + fn promote_state_fast_forwards_past_the_lease() { + let mut csprng = rng(); + let base_key = KeyPair::generate(&mut csprng).public_key; + let state = create_test_session_state(3, &base_key); + + let mut record = SessionRecord::new_fresh(); + record.reserve_sender_chain_counters(0); + let reserved = record.reserved_sender_chain_index(); + assert_eq!(reserved, consts::SENDER_CHAIN_RESERVATION_BATCH); + + record.promote_state(state); + let chain = record + .session_state() + .unwrap() + .get_sender_chain_key() + .unwrap(); + assert_eq!( + chain.index(), + reserved, + "promotion must make every leased counter underivable" + ); + } + + /// A freshly ratcheted chain has never spent a counter, so promoting it + /// resets the lease instead of burning it (the first send re-reserves). + #[test] + fn promote_fresh_state_resets_the_lease() { + let mut csprng = rng(); + let base_key = KeyPair::generate(&mut csprng).public_key; + let state = create_test_session_state(3, &base_key); + + let mut record = SessionRecord::new_fresh(); + record.reserve_sender_chain_counters(500); + + record.promote_fresh_state(state); + assert_eq!(record.reserved_sender_chain_index(), 0); + let chain = record + .session_state() + .unwrap() + .get_sender_chain_key() + .unwrap(); + assert_eq!(chain.index(), 0, "a fresh chain must not be burned"); + } + + /// A corrupt reservation absurdly far ahead of the chain must be refused + /// instead of turning the load into an unbounded KDF loop. + #[test] + fn deserialize_rejects_an_implausible_reservation() { + let mut csprng = rng(); + let base_key = KeyPair::generate(&mut csprng).public_key; + let record = SessionRecord::new(create_test_session_state(3, &base_key)); + + let mut bytes = record.serialize().unwrap(); + { + use buffa::encoding::{Tag, WireType, encode_varint}; + Tag::new(RESERVED_SENDER_CHAIN_INDEX_FIELD, WireType::Varint).encode(&mut bytes); + encode_varint( + (consts::MAX_RESERVATION_FAST_FORWARD as u64) + 1, + &mut bytes, + ); + } + + assert!( + SessionRecord::deserialize(&bytes).is_err(), + "an implausible lease must fail the load, not fast-forward" + ); + } + /// Creates a SessionRecord with N previous sessions for testing. fn create_record_with_previous_sessions(count: usize) -> SessionRecord { let mut csprng = rng(); @@ -1489,6 +1710,8 @@ mod tests { let record = SessionRecord { current_session: Some(SessionState::from_session_structure(current.clone())), previous_sessions: Arc::new(previous_sessions.clone()), + reserved_sender_chain_index: 0, + pending_reservation: false, }; let expected = waproto::whatsapp::RecordStructure { current_session: MessageField::some(current), diff --git a/wacore/libsignal/tests/counter_lease.rs b/wacore/libsignal/tests/counter_lease.rs new file mode 100644 index 000000000..bd2afcf6c --- /dev/null +++ b/wacore/libsignal/tests/counter_lease.rs @@ -0,0 +1,496 @@ +//! Sender-chain counter lease: outbound message keys/IVs derive +//! deterministically from the chain counter, so a counter must never be +//! re-derivable after a crash. Instead of persisting every advance before the +//! wire, the record durably reserves counters in batches +//! (`SENDER_CHAIN_RESERVATION_BATCH`) and a reloaded snapshot fast-forwards +//! past the whole lease. These tests drive the crash/reload interleavings +//! end-to-end against a real peer. +//! Async I/O uses `futures::executor::block_on` (no tokio in this crate). + +use async_trait::async_trait; +use std::collections::HashMap; +use wacore_libsignal::protocol::consts::SENDER_CHAIN_RESERVATION_BATCH; +use wacore_libsignal::protocol::{ + CiphertextMessage, Direction, GenericSignedPreKey, IdentityChange, IdentityKey, + IdentityKeyPair, IdentityKeyStore, KeyPair, PreKeyBundle, PreKeyId, PreKeyRecord, PreKeyStore, + ProtocolAddress, SessionRecord, SessionStore, SignalProtocolError, SignedPreKeyId, + SignedPreKeyRecord, SignedPreKeyStore, Timestamp, UsePQRatchet, message_decrypt, + message_encrypt, process_prekey_bundle, +}; + +// ---- in-memory store impls (clones of the session_divergence fixtures, +// kept local so this test file is self-contained) ----------------------------- + +#[derive(Clone)] +struct InMemoryIdentityKeyStore { + identity_key_pair: IdentityKeyPair, + registration_id: u32, + identities: HashMap, +} + +#[async_trait] +impl IdentityKeyStore for InMemoryIdentityKeyStore { + async fn get_identity_key_pair( + &self, + ) -> wacore_libsignal::protocol::error::Result { + Ok(self.identity_key_pair.clone()) + } + async fn get_local_registration_id(&self) -> wacore_libsignal::protocol::error::Result { + Ok(self.registration_id) + } + async fn save_identity( + &mut self, + address: &ProtocolAddress, + identity: &IdentityKey, + ) -> wacore_libsignal::protocol::error::Result { + let changed = self + .identities + .get(address) + .is_some_and(|prev| prev != identity); + self.identities.insert(address.clone(), *identity); + Ok(IdentityChange::from_changed(changed)) + } + async fn is_trusted_identity( + &self, + _: &ProtocolAddress, + _: &IdentityKey, + _: Direction, + ) -> wacore_libsignal::protocol::error::Result { + Ok(true) + } + async fn get_identity( + &self, + address: &ProtocolAddress, + ) -> wacore_libsignal::protocol::error::Result> { + Ok(self.identities.get(address).cloned()) + } +} + +#[derive(Default, Clone)] +struct InMemoryPreKeyStore(HashMap); + +#[async_trait] +impl PreKeyStore for InMemoryPreKeyStore { + async fn get_pre_key( + &self, + id: PreKeyId, + ) -> wacore_libsignal::protocol::error::Result { + self.0 + .get(&id) + .cloned() + .ok_or(SignalProtocolError::InvalidPreKeyId) + } + async fn save_pre_key( + &mut self, + id: PreKeyId, + record: &PreKeyRecord, + ) -> wacore_libsignal::protocol::error::Result<()> { + self.0.insert(id, record.clone()); + Ok(()) + } + async fn remove_pre_key( + &mut self, + id: PreKeyId, + ) -> wacore_libsignal::protocol::error::Result<()> { + self.0.remove(&id); + Ok(()) + } +} + +#[derive(Default, Clone)] +struct InMemorySignedPreKeyStore(HashMap); + +#[async_trait] +impl SignedPreKeyStore for InMemorySignedPreKeyStore { + async fn get_signed_pre_key( + &self, + id: SignedPreKeyId, + ) -> wacore_libsignal::protocol::error::Result { + self.0 + .get(&id) + .cloned() + .ok_or(SignalProtocolError::InvalidSignedPreKeyId) + } + async fn save_signed_pre_key( + &mut self, + id: SignedPreKeyId, + record: &SignedPreKeyRecord, + ) -> wacore_libsignal::protocol::error::Result<()> { + self.0.insert(id, record.clone()); + Ok(()) + } +} + +#[derive(Default, Clone)] +struct InMemorySessionStore(HashMap); + +#[async_trait] +impl SessionStore for InMemorySessionStore { + async fn load_session( + &self, + address: &ProtocolAddress, + ) -> wacore_libsignal::protocol::error::Result> { + Ok(self.0.get(address).cloned()) + } + async fn has_session( + &self, + address: &ProtocolAddress, + ) -> wacore_libsignal::protocol::error::Result { + Ok(self.0.contains_key(address)) + } + async fn store_session( + &mut self, + address: &ProtocolAddress, + record: SessionRecord, + ) -> wacore_libsignal::protocol::error::Result<()> { + self.0.insert(address.clone(), record); + Ok(()) + } +} + +// ---- peer fixture ----------------------------------------------------------- + +struct Peer { + address: ProtocolAddress, + identity_store: InMemoryIdentityKeyStore, + prekey_store: InMemoryPreKeyStore, + signed_prekey_store: InMemorySignedPreKeyStore, + session_store: InMemorySessionStore, +} + +impl Peer { + fn new(name: &str) -> Self { + let mut rng = rand::make_rng::(); + + let identity_key_pair = IdentityKeyPair::generate(&mut rng); + let registration_id = rand::random::() & 0x3FFF; + + let prekey_id: PreKeyId = 1u32.into(); + let prekey_pair = KeyPair::generate(&mut rng); + + let signed_prekey_id: SignedPreKeyId = 1u32.into(); + let signed_prekey_pair = KeyPair::generate(&mut rng); + let signed_prekey_signature = identity_key_pair + .private_key() + .calculate_signature(&signed_prekey_pair.public_key.serialize(), &mut rng) + .expect("sign"); + + let mut prekey_store = InMemoryPreKeyStore::default(); + let mut signed_prekey_store = InMemorySignedPreKeyStore::default(); + futures::executor::block_on(async { + prekey_store + .save_pre_key(prekey_id, &PreKeyRecord::new(prekey_id, &prekey_pair)) + .await + .unwrap(); + signed_prekey_store + .save_signed_pre_key( + signed_prekey_id, + &SignedPreKeyRecord::new( + signed_prekey_id, + Timestamp::from_epoch_millis(0), + &signed_prekey_pair, + &signed_prekey_signature, + ), + ) + .await + .unwrap(); + }); + + let bundle = PreKeyBundle::new( + registration_id, + 1u32.into(), + Some((prekey_id, prekey_pair.public_key)), + signed_prekey_id, + signed_prekey_pair.public_key, + signed_prekey_signature.to_vec(), + *identity_key_pair.identity_key(), + ) + .expect("valid bundle"); + + let peer = Self { + address: ProtocolAddress::new(name.to_string(), 1u32.into()), + identity_store: InMemoryIdentityKeyStore { + identity_key_pair, + registration_id, + identities: HashMap::new(), + }, + prekey_store, + signed_prekey_store, + session_store: InMemorySessionStore::default(), + }; + BUNDLES.with(|b| b.borrow_mut().insert(peer.address.clone(), bundle)); + peer + } +} + +thread_local! { + /// Bundle published by each peer at construction, keyed by address. + static BUNDLES: std::cell::RefCell> = + std::cell::RefCell::new(HashMap::new()); +} + +// ---- helpers ---------------------------------------------------------------- + +fn process_bundle(initiator: &mut Peer, target: &ProtocolAddress) { + let bundle = BUNDLES.with(|b| b.borrow().get(target).cloned().expect("bundle published")); + let mut rng = rand::make_rng::(); + futures::executor::block_on(async { + process_prekey_bundle( + target, + &mut initiator.session_store, + &mut initiator.identity_store, + &bundle, + &mut rng, + UsePQRatchet::No, + ) + .await + .expect("prekey bundle accepted"); + }); +} + +fn send(from: &mut Peer, to: &ProtocolAddress, plaintext: &[u8]) -> CiphertextMessage { + futures::executor::block_on(async { + message_encrypt( + plaintext, + to, + &mut from.session_store, + &mut from.identity_store, + ) + .await + .expect("encrypt") + }) +} + +fn receive( + to: &mut Peer, + from: &ProtocolAddress, + ct: &CiphertextMessage, +) -> Result, SignalProtocolError> { + let mut rng = rand::make_rng::(); + futures::executor::block_on(async { + message_decrypt( + ct, + from, + &mut to.session_store, + &mut to.identity_store, + &mut to.prekey_store, + &to.signed_prekey_store, + &mut rng, + UsePQRatchet::No, + ) + .await + .map(|d| d.plaintext) + }) +} + +fn establish(alice: &mut Peer, bob: &mut Peer) { + let bob_address = bob.address.clone(); + process_bundle(alice, &bob_address); + let ct = send(alice, &bob_address, b"hello bob"); + let plaintext = receive(bob, &alice.address.clone(), &ct).expect("first pkmsg decrypts"); + assert_eq!(&plaintext[..], b"hello bob"); +} + +/// The wire counter of an outbound message (pkmsg or msg). +fn wire_counter(ct: &CiphertextMessage) -> u32 { + match ct { + CiphertextMessage::SignalMessage(m) => m.counter(), + CiphertextMessage::PreKeySignalMessage(m) => m.message().counter(), + other => panic!("unexpected message type {:?}", other.message_type()), + } +} + +fn record_of(peer: &Peer, remote: &ProtocolAddress) -> SessionRecord { + peer.session_store + .0 + .get(remote) + .expect("session exists") + .clone() +} + +/// Simulate the store layer acknowledging a durable flush: take over the wire +/// gate and return the serialized snapshot that "reached storage". +fn ack_flush(peer: &mut Peer, remote: &ProtocolAddress) -> Vec { + let record = peer.session_store.0.get_mut(remote).expect("session"); + record.clear_pending_reservation(); + record.serialize().expect("serialize") +} + +/// Simulate a crash: whatever was in memory is gone, the last durable +/// snapshot is what comes back. +fn crash_reload(peer: &mut Peer, remote: &ProtocolAddress, snapshot: &[u8]) { + let restored = SessionRecord::deserialize(snapshot).expect("snapshot deserializes"); + peer.session_store.0.insert(remote.clone(), restored); +} + +// ---- scenarios -------------------------------------------------------------- + +/// The very first send on a fresh session must raise a lease and gate the +/// wire on its durability. +#[test] +fn first_send_raises_the_lease() { + let mut alice = Peer::new("alice"); + let mut bob = Peer::new("bob"); + establish(&mut alice, &mut bob); + + let record = record_of(&alice, &bob.address); + assert!( + record.has_pending_reservation(), + "first send must gate the wire on the raised lease" + ); + assert_eq!( + record.reserved_sender_chain_index(), + SENDER_CHAIN_RESERVATION_BATCH, + "counter 0 leases one full batch" + ); +} + +/// Steady-state ping-pong (every send is counter 0 of a freshly ratcheted +/// chain) must never re-raise the lease: this is what removes the per-message +/// synchronous flush from the hot send path. +#[test] +fn ping_pong_sends_stay_covered_by_the_lease() { + let mut alice = Peer::new("alice"); + let mut bob = Peer::new("bob"); + establish(&mut alice, &mut bob); + ack_flush(&mut alice, &bob.address); + + for i in 0..20 { + let reply = format!("b→a #{i}"); + let ct = send(&mut bob, &alice.address, reply.as_bytes()); + receive(&mut alice, &bob.address, &ct).expect("decrypt reply"); + + let msg = format!("a→b #{i}"); + let ct = send(&mut alice, &bob.address, msg.as_bytes()); + assert!( + !record_of(&alice, &bob.address).has_pending_reservation(), + "ping-pong send #{i} is covered by the durable lease and must not re-flush" + ); + let pt = receive(&mut bob, &alice.address, &ct).expect("decrypt"); + assert_eq!(&pt[..], msg.as_bytes()); + } +} + +/// A monologue re-raises the lease exactly when it runs out, one batch at a +/// time. +#[test] +fn monologue_re_raises_the_lease_at_the_batch_boundary() { + let batch = SENDER_CHAIN_RESERVATION_BATCH; + let mut alice = Peer::new("alice"); + let mut bob = Peer::new("bob"); + establish(&mut alice, &mut bob); // counter 0, lease -> batch + ack_flush(&mut alice, &bob.address); + + // Counters 1..batch-1 ride the existing lease. + for i in 1..batch { + let ct = send(&mut alice, &bob.address, b"streak"); + assert_eq!(wire_counter(&ct), i); + assert!( + !record_of(&alice, &bob.address).has_pending_reservation(), + "counter {i} is inside the lease" + ); + } + + // Counter `batch` exhausts it: the lease must be re-raised. + let ct = send(&mut alice, &bob.address, b"boundary"); + assert_eq!(wire_counter(&ct), batch); + let record = record_of(&alice, &bob.address); + assert!(record.has_pending_reservation()); + assert_eq!(record.reserved_sender_chain_index(), batch * 2); +} + +/// The core no-reuse guarantee: sends past the durable snapshot are covered +/// by its lease, so a crash/reload can never re-derive their counters — and +/// the peer keeps decrypting across the gap. +#[test] +fn crash_reload_never_reuses_a_counter_and_peer_decrypts_across_the_gap() { + let mut alice = Peer::new("alice"); + let mut bob = Peer::new("bob"); + establish(&mut alice, &mut bob); // counter 0 + let snapshot = ack_flush(&mut alice, &bob.address); + + // Five more sends after the snapshot; the durable state now trails. + let mut spent = vec![0u32]; + for _ in 0..5 { + let ct = send(&mut alice, &bob.address, b"unflushed"); + spent.push(wire_counter(&ct)); + receive(&mut bob, &alice.address, &ct).expect("decrypt"); + } + + crash_reload(&mut alice, &bob.address, &snapshot); + + // The reloaded chain resumes past the whole lease... + let ct = send(&mut alice, &bob.address, b"after crash"); + let resumed = wire_counter(&ct); + assert_eq!( + resumed, SENDER_CHAIN_RESERVATION_BATCH, + "reload must fast-forward to the leased ceiling" + ); + assert!( + !spent.contains(&resumed), + "a wire counter must never repeat across a crash" + ); + // ...the resumed counter exhausts the old lease, so it re-raises... + assert!(record_of(&alice, &bob.address).has_pending_reservation()); + // ...and Bob decrypts across the gap (skipped keys for the burned range). + let pt = receive(&mut bob, &alice.address, &ct).expect("decrypt across the gap"); + assert_eq!(&pt[..], b"after crash"); +} + +/// Crash after a DH ratchet whose new chain never reached storage: the +/// snapshot's OLD chain resumes past its lease, the lost chain's keys are +/// unrecoverable (fresh random ephemeral), and the peer still decrypts via +/// its retained old receiver chain. +#[test] +fn crash_reload_after_unflushed_ratchet_resumes_the_old_chain_safely() { + let mut alice = Peer::new("alice"); + let mut bob = Peer::new("bob"); + establish(&mut alice, &mut bob); + let snapshot = ack_flush(&mut alice, &bob.address); + + // Bob's reply DH-ratchets Alice onto a brand-new sender chain; her send + // on it is lease-covered (no flush) and the chain never gets persisted. + let ct = send(&mut bob, &alice.address, b"reply"); + receive(&mut alice, &bob.address, &ct).expect("decrypt reply"); + let ct = send(&mut alice, &bob.address, b"on the lost chain"); + assert_eq!(wire_counter(&ct), 0, "fresh chain starts at 0"); + assert!( + !record_of(&alice, &bob.address).has_pending_reservation(), + "the ratcheted chain send rides the record lease" + ); + receive(&mut bob, &alice.address, &ct).expect("decrypt"); + + crash_reload(&mut alice, &bob.address, &snapshot); + + // Alice resumes on the old chain, past its lease; Bob retained the old + // receiver chain and decrypts. + let ct = send(&mut alice, &bob.address, b"back on the old chain"); + assert_eq!(wire_counter(&ct), SENDER_CHAIN_RESERVATION_BATCH); + let pt = receive(&mut bob, &alice.address, &ct).expect("old receiver chain still works"); + assert_eq!(&pt[..], b"back on the old chain"); +} + +/// Serialize/deserialize round-trip: the lease survives storage, and a +/// snapshot with no lease (legacy format) loads with a zero reservation and +/// an untouched chain. +#[test] +fn lease_round_trips_through_storage_and_legacy_records_load_untouched() { + let mut alice = Peer::new("alice"); + let mut bob = Peer::new("bob"); + establish(&mut alice, &mut bob); + + let bytes = ack_flush(&mut alice, &bob.address); + let reloaded = SessionRecord::deserialize(&bytes).expect("deserialize"); + assert_eq!( + reloaded.reserved_sender_chain_index(), + SENDER_CHAIN_RESERVATION_BATCH + ); + assert!(!reloaded.has_pending_reservation(), "the gate is transient"); + + // A legacy record (serialized before the lease existed) must load with a + // zero reservation. `new_fresh` never leases, so its encoding matches the + // legacy layout exactly. + let legacy = SessionRecord::new_fresh().serialize().expect("serialize"); + let reloaded = SessionRecord::deserialize(&legacy).expect("legacy deserializes"); + assert_eq!(reloaded.reserved_sender_chain_index(), 0); +} diff --git a/wacore/src/store/signal_cache.rs b/wacore/src/store/signal_cache.rs index aefa7a62d..f7b326e64 100644 --- a/wacore/src/store/signal_cache.rs +++ b/wacore/src/store/signal_cache.rs @@ -103,6 +103,13 @@ struct SessionStoreState { cache: HashMap, SessionEntry>, dirty: HashSet>, deleted: HashSet>, + /// Sessions whose raised counter reservation has not reached the backend + /// yet. While any address is here, an outbound ciphertext may be relying + /// on a lease that only exists in memory, so the send path must flush + /// before the wire. Entries leave only when a flush actually persists + /// them (or the session is deleted/cleared). Always a subset of `dirty`, + /// so eviction can never drop a pending entry. + reservation_pending: HashSet>, } impl SessionStoreState { @@ -111,6 +118,7 @@ impl SessionStoreState { cache: HashMap::new(), dirty: HashSet::new(), deleted: HashSet::new(), + reservation_pending: HashSet::new(), } } @@ -123,8 +131,14 @@ impl SessionStoreState { } } - fn put(&mut self, address: &str, record: SessionRecord) { + fn put(&mut self, address: &str, mut record: SessionRecord) { let addr = self.key_for(address); + // Take over the record's wire gate: the address stays pending until a + // flush persists it, regardless of later checkout/put round trips. + if record.has_pending_reservation() { + record.clear_pending_reservation(); + self.reservation_pending.insert(addr.clone()); + } self.cache .insert(addr.clone(), SessionEntry::Present(Arc::new(record))); self.dirty.insert(addr.clone()); @@ -136,12 +150,16 @@ impl SessionStoreState { self.cache.insert(addr.clone(), SessionEntry::Absent); self.deleted.insert(addr.clone()); self.dirty.remove(&addr); + self.reservation_pending.remove(&addr); } fn clear(&mut self) { self.cache.clear(); self.dirty.clear(); self.deleted.clear(); + // Cleared sessions reload from the backend, whose snapshot still + // carries its own (older) lease; the first send re-reserves. + self.reservation_pending.clear(); } fn evict_if_needed(&mut self, max_entries: usize) { @@ -175,6 +193,11 @@ struct SenderKeyStoreState { // `VecDeque` with up to `MAX_MESSAGE_KEYS` message keys each. cache: HashMap, Option>>, dirty: HashSet>, + /// Chains advanced by an outbound encrypt and not yet persisted; the send + /// path must flush before the wire while any entry is here. Decrypt-side + /// dirtiness deliberately does NOT enter this set (it re-derives forward), + /// so unrelated group receives never force a sync flush onto a DM send. + wire_gate_pending: HashSet>, } impl SenderKeyStoreState { @@ -182,6 +205,7 @@ impl SenderKeyStoreState { Self { cache: HashMap::new(), dirty: HashSet::new(), + wire_gate_pending: HashSet::new(), } } @@ -192,8 +216,12 @@ impl SenderKeyStoreState { } } - fn put(&mut self, address: &str, record: SenderKeyRecord) { + fn put(&mut self, address: &str, mut record: SenderKeyRecord) { let addr = self.key_for(address); + if record.is_wire_gated() { + record.clear_wire_gated(); + self.wire_gate_pending.insert(addr.clone()); + } self.cache.insert(addr.clone(), Some(Arc::new(record))); self.dirty.insert(addr.clone()); } @@ -201,12 +229,14 @@ impl SenderKeyStoreState { fn delete(&mut self, address: &str) { let addr = self.key_for(address); self.cache.insert(addr.clone(), None); - self.dirty.insert(addr); + self.dirty.insert(addr.clone()); + self.wire_gate_pending.remove(&addr); } fn clear(&mut self) { self.cache.clear(); self.dirty.clear(); + self.wire_gate_pending.clear(); } fn evict_if_needed(&mut self, max_entries: usize) { @@ -677,6 +707,11 @@ impl SignalStoreCache { } if !batch.is_empty() { backend.put_sessions_batch(&batch).await?; + // These leases are durable now; only the written addresses + // leave the pending set (a CheckedOut session stays gated). + for (address, _) in &batch { + state.reservation_pending.remove(address); + } } for address in &deleted_keys { backend.delete_session(address).await?; @@ -797,6 +832,9 @@ impl SignalStoreCache { } if !batch.is_empty() { backend.put_sender_keys_batch(&batch).await?; + for (name, _) in &batch { + state.wire_gate_pending.remove(name); + } } for key in &dirty_keys { @@ -808,6 +846,19 @@ impl SignalStoreCache { Ok(()) } + /// Whether an outbound ciphertext produced since the last flush is still + /// gated on durability: a session counter lease was raised, or a + /// sender-key chain advanced via encrypt, and neither has reached the + /// backend. The send path flushes synchronously only while this holds; + /// everything else (decrypt advances, identities) safely rides the + /// coalesced write-behind. + pub async fn needs_pre_wire_flush(&self) -> bool { + if !self.sessions.lock().await.reservation_pending.is_empty() { + return true; + } + !self.sender_keys.lock().await.wire_gate_pending.is_empty() + } + /// Entry counts and estimated retained bytes for each store /// (sessions, identities, sender_keys). Sizes use the records' encoded-size /// proxy (see `SessionRecord::estimated_size`); on-demand only — walks the @@ -1900,3 +1951,134 @@ mod eviction_tests { } } } + +#[cfg(test)] +mod pre_wire_gate_tests { + use super::*; + use crate::libsignal::store::sender_key_name::SenderKeyName; + use crate::store::in_memory::InMemoryBackend; + + fn addr(user: &str) -> ProtocolAddress { + ProtocolAddress::new(user.to_string(), 1.into()) + } + + fn leased_record() -> SessionRecord { + let mut record = SessionRecord::new_fresh(); + record.reserve_sender_chain_counters(0); + record + } + + /// A raised lease gates the wire until a flush actually persists it; a + /// plain (decrypt-style) session write never does. + #[tokio::test] + async fn session_lease_gates_until_a_successful_flush() { + let backend = InMemoryBackend::new(); + let cache = SignalStoreCache::new(); + + cache + .put_session(&addr("15550000001"), SessionRecord::new_fresh()) + .await; + assert!( + !cache.needs_pre_wire_flush().await, + "a dirty session without a raised lease must not gate the wire" + ); + + cache + .put_session(&addr("15550000002"), leased_record()) + .await; + assert!(cache.needs_pre_wire_flush().await); + + cache.flush(&backend).await.unwrap(); + assert!( + !cache.needs_pre_wire_flush().await, + "a persisted lease releases the gate" + ); + } + + /// A failed flush must keep the gate closed — the lease never reached + /// storage, so the ciphertext must keep waiting. + #[tokio::test] + async fn failed_flush_keeps_the_gate_closed() { + let backend = InMemoryBackend::new(); + let cache = SignalStoreCache::new(); + + cache + .put_session(&addr("15550000003"), leased_record()) + .await; + backend.set_fail_session_writes(true); + assert!(cache.flush(&backend).await.is_err()); + assert!( + cache.needs_pre_wire_flush().await, + "an unpersisted lease must keep gating the wire" + ); + + backend.set_fail_session_writes(false); + cache.flush(&backend).await.unwrap(); + assert!(!cache.needs_pre_wire_flush().await); + } + + /// A checked-out session cannot be persisted by a flush, so its pending + /// lease must survive that flush and release only once the returned + /// record is actually written. + #[tokio::test] + async fn checked_out_session_keeps_its_lease_pending_across_a_flush() { + let backend = InMemoryBackend::new(); + let cache = SignalStoreCache::new(); + let a = addr("15550000004"); + + cache.put_session(&a, leased_record()).await; + let taken = cache.get_session(&a, &backend).await.unwrap().unwrap(); + + cache.flush(&backend).await.unwrap(); + assert!( + cache.needs_pre_wire_flush().await, + "a checked-out lease was not persisted and must keep the gate closed" + ); + + cache.put_session(&a, taken).await; + cache.flush(&backend).await.unwrap(); + assert!(!cache.needs_pre_wire_flush().await); + } + + /// Outbound sender-key advances gate the wire; decrypt-side dirtiness + /// (no wire gate mark) must not, so group receives never force a sync + /// flush onto an unrelated DM send. + #[tokio::test] + async fn only_encrypt_marked_sender_keys_gate_the_wire() { + let backend = InMemoryBackend::new(); + let cache = SignalStoreCache::new(); + let name = SenderKeyName::from_parts("g@g.us", "u@s.whatsapp.net:0"); + + cache + .put_sender_key(&name, SenderKeyRecord::new_empty()) + .await; + assert!( + !cache.needs_pre_wire_flush().await, + "a decrypt-side sender-key write must not gate the wire" + ); + + let mut outbound = SenderKeyRecord::new_empty(); + outbound.mark_wire_gated(); + cache.put_sender_key(&name, outbound).await; + assert!(cache.needs_pre_wire_flush().await); + + cache.flush(&backend).await.unwrap(); + assert!(!cache.needs_pre_wire_flush().await); + } + + /// Deleting or clearing drops the pending gate together with the state it + /// guarded (the reloaded snapshot re-reserves on its first send). + #[tokio::test] + async fn delete_and_clear_drop_the_pending_gate() { + let cache = SignalStoreCache::new(); + let a = addr("15550000005"); + + cache.put_session(&a, leased_record()).await; + cache.delete_session(&a).await; + assert!(!cache.needs_pre_wire_flush().await); + + cache.put_session(&a, leased_record()).await; + cache.clear().await; + assert!(!cache.needs_pre_wire_flush().await); + } +} From 6bdd2e45c2a5a06e850eb077392f68e18b2f65bb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:05:27 -0300 Subject: [PATCH 2/3] test(signal): cover the sender-key wire gate surviving a failed flush Mirrors the session-lease failed-flush test: a flush that errors writing the chain advance must leave the gate closed. Adds the matching fail_sender_key_writes hook to the in-memory backend. --- wacore/src/store/in_memory.rs | 19 +++++++++++++++++++ wacore/src/store/signal_cache.rs | 24 ++++++++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index dac9139da..4b1467dd5 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -104,6 +104,11 @@ pub struct InMemoryBackend { /// ratchet advance cannot be persisted. #[cfg(any(test, feature = "test-util"))] fail_session_writes: AtomicBool, + /// When set, `put_sender_keys_batch` fails. Test hook: the sender-key + /// counterpart of `fail_session_writes` (wire gate must survive a failed + /// flush). + #[cfg(any(test, feature = "test-util"))] + fail_sender_key_writes: AtomicBool, } impl InMemoryBackend { @@ -116,6 +121,8 @@ impl InMemoryBackend { session_batch_writes: AtomicU32::new(0), #[cfg(any(test, feature = "test-util"))] fail_session_writes: AtomicBool::new(false), + #[cfg(any(test, feature = "test-util"))] + fail_sender_key_writes: AtomicBool::new(false), } } @@ -130,6 +137,12 @@ impl InMemoryBackend { pub fn set_fail_session_writes(&self, fail: bool) { self.fail_session_writes.store(fail, Ordering::Relaxed); } + + /// Make every subsequent `put_sender_keys_batch` fail (or stop failing). + #[cfg(any(test, feature = "test-util"))] + pub fn set_fail_sender_key_writes(&self, fail: bool) { + self.fail_sender_key_writes.store(fail, Ordering::Relaxed); + } } impl Default for InMemoryBackend { @@ -315,6 +328,12 @@ impl SignalStore for InMemoryBackend { } async fn put_sender_key(&self, address: &str, record: &[u8]) -> Result<()> { + #[cfg(any(test, feature = "test-util"))] + if self.fail_sender_key_writes.load(Ordering::Relaxed) { + return Err(crate::store::error::StoreError::Io(std::io::Error::other( + "put_sender_key failing (test hook)", + ))); + } self.state .lock() .await diff --git a/wacore/src/store/signal_cache.rs b/wacore/src/store/signal_cache.rs index f7b326e64..341650cd6 100644 --- a/wacore/src/store/signal_cache.rs +++ b/wacore/src/store/signal_cache.rs @@ -2066,6 +2066,30 @@ mod pre_wire_gate_tests { assert!(!cache.needs_pre_wire_flush().await); } + /// The sender-key counterpart of `failed_flush_keeps_the_gate_closed`: a + /// flush that fails writing the chain advance must keep the wire gated. + #[tokio::test] + async fn failed_flush_keeps_the_sender_key_gate_closed() { + let backend = InMemoryBackend::new(); + let cache = SignalStoreCache::new(); + let name = SenderKeyName::from_parts("g@g.us", "u@s.whatsapp.net:0"); + + let mut outbound = SenderKeyRecord::new_empty(); + outbound.mark_wire_gated(); + cache.put_sender_key(&name, outbound).await; + + backend.set_fail_sender_key_writes(true); + assert!(cache.flush(&backend).await.is_err()); + assert!( + cache.needs_pre_wire_flush().await, + "an unpersisted sender-key advance must keep gating the wire" + ); + + backend.set_fail_sender_key_writes(false); + cache.flush(&backend).await.unwrap(); + assert!(!cache.needs_pre_wire_flush().await); + } + /// Deleting or clearing drops the pending gate together with the state it /// guarded (the reloaded snapshot re-reserves on its first send). #[tokio::test] From d30eaab0087b66495c1a39bb1a8fdf83c1dd7147 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:39:29 -0300 Subject: [PATCH 3/3] fix(signal): fail closed on an unreadable lease field; document the gate's edges MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups on the counter-lease PR. Real fix — decode_reserved_index skipped a field-100 tag it could not read (wrong wire type), and let a duplicate win under last-one-wins. Both paths silently yield reservation 0, which disables the load-time fast-forward and re-enables every counter the lease had already spent: fail-open in exactly the place the mechanism exists to prevent reuse. Both now fail the load, which is recoverable where a reused (key, IV) pair is not. Documented, not changed: - clear() dropping a pending gate is safe only because no clear can race a live wire: every caller runs before the transport exists (connect) or after cleanup_connection_state took the noise socket, and a send resolves that socket after its pre-wire gate — so a send that sees the cleared set fails NotConnected instead of reaching a peer. That ordering is load-bearing and was undocumented; moving a clear ahead of the socket teardown would reintroduce reuse for lease-raising sends. - the pre-wire gate is a global predicate, so another session's unpersisted lease can still make this send flush (and its failure abort this send). Erring toward an extra flush is the safe direction; scoping it to the stanza's addresses would need them threaded back up the send path. - downgrading the library across a crash resumes a leased chain from its stale snapshot: already-released readers skip unknown fields by definition, so no code here can gate it. Release-note constraint. set_states_for_testing now resets wire_gated with the states it replaces. Adds guards for both fixture and format: peers_generate_independent_keys pins that make_rng seeds from OS entropy (a deterministic fixture would make the crash assertions pass for the wrong reason), and the decode test covers the wrong-wire-type and duplicate-field cases. --- src/client/adapters.rs | 15 +++- wacore/libsignal/src/protocol/sender_keys.rs | 3 + .../libsignal/src/protocol/state/session.rs | 70 +++++++++++++++++-- wacore/libsignal/tests/counter_lease.rs | 39 +++++++++++ wacore/src/store/signal_cache.rs | 12 +++- 5 files changed, 128 insertions(+), 11 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 10ae56ca2..3aa39023e 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -72,8 +72,9 @@ impl Client { /// returning once the flush completes (or fails). /// /// The live receive path schedules a coalesced flush (see `signal_flush.rs`) - /// instead of writing through, and lease-covered sends do the same (only a - /// send that raises its counter lease flushes synchronously). On success the + /// instead of writing through, and lease-covered sends do the same. Only a + /// send that raises its counter lease — or advances a group sender-key + /// chain, which has no lease — flushes synchronously. On success the /// backend normally trails the cache by about the coalescing window, but /// that is not a hard wall-clock bound — the timer can slip under runtime /// starvation and the flush can wait on locks or slow/failing storage (a @@ -170,6 +171,16 @@ impl Client { /// path instead of paying a serialize + storage transaction per message. /// A failure must abort the send — transmitting a ciphertext whose lease /// could not be persisted reintroduces the counter-reuse window. + /// + /// The gate is deliberately a GLOBAL predicate, not one scoped to the + /// addresses this stanza encrypted for: the send path does not carry them + /// back up, and erring toward an extra flush is the safe direction (it can + /// over-flush, never under-flush). The cost is a sharp edge under a failing + /// backend: another session's pending lease makes this send flush, and that + /// flush's failure aborts this send too — the same collateral every send + /// took before leases existed, now limited to the window where some lease + /// is actually unpersisted. So "a lease-covered send needs no flush" is a + /// statement about what this send REQUIRES, not a promise it never flushes. pub(crate) async fn persist_signal_state_pre_wire(&self) -> Result<(), anyhow::Error> { if self.signal_cache.needs_pre_wire_flush().await { return self.flush_signal_cache_batch_safe().await; diff --git a/wacore/libsignal/src/protocol/sender_keys.rs b/wacore/libsignal/src/protocol/sender_keys.rs index 00980b320..c324c5a7b 100644 --- a/wacore/libsignal/src/protocol/sender_keys.rs +++ b/wacore/libsignal/src/protocol/sender_keys.rs @@ -437,8 +437,11 @@ pub struct SenderKeyRecord { } impl SenderKeyRecord { + /// Replaces the states wholesale, so the wire gate — which belongs to the + /// advance being replaced — resets with them. pub fn set_states_for_testing(&mut self, states: std::collections::VecDeque) { self.states = states; + self.wire_gated = false; } pub fn new_empty() -> Self { diff --git a/wacore/libsignal/src/protocol/state/session.rs b/wacore/libsignal/src/protocol/state/session.rs index ce59bc395..e8b46762a 100644 --- a/wacore/libsignal/src/protocol/state/session.rs +++ b/wacore/libsignal/src/protocol/state/session.rs @@ -666,6 +666,13 @@ impl From<&SessionState> for SessionStructure { /// [`SessionRecord::serialize_into`] — writes it directly; the number sits far /// above RecordStructure's fields (1, 2) so a future upstream addition cannot /// collide. Standard unknown-field skipping keeps old readers compatible. +/// +/// That compatibility is one-way: a build without lease support silently +/// ignores this field, so downgrading the library across a crash can resume a +/// leased chain from its stale snapshot and reuse a counter. Nothing here can +/// prevent that — already-released readers skip unknown fields by definition — +/// so it is a release-note constraint, not a code one. Never lower this +/// number into a range an older reader might interpret. const RESERVED_SENDER_CHAIN_INDEX_FIELD: u32 = 100; #[derive(Clone)] @@ -773,27 +780,37 @@ impl SessionRecord { /// Extract the record-level reservation field from the serialized bytes. /// The generated `RecordStructureView` skips fields it does not know, so /// the local-only field is scanned out of the raw top-level stream here. + /// + /// Fails closed: anything unreadable under our own field number is an + /// error rather than a skip. Skipping would silently yield reservation 0, + /// which disables the load-time fast-forward and re-enables already-spent + /// counters — the one outcome this whole mechanism exists to prevent. A + /// load error is recoverable (the session rebuilds); a reused (key, IV) + /// pair is not. fn decode_reserved_index(bytes: &[u8]) -> Result { use buffa::encoding::{Tag, WireType, decode_varint, skip_field}; let mut buf = bytes; - let mut reserved = 0u32; + let mut reserved = None; while !buf.is_empty() { let tag = Tag::decode(&mut buf) .map_err(|_| InvalidSessionError("failed to decode session record protobuf"))?; - if tag.field_number() == RESERVED_SENDER_CHAIN_INDEX_FIELD - && tag.wire_type() == WireType::Varint - { + if tag.field_number() == RESERVED_SENDER_CHAIN_INDEX_FIELD { + if tag.wire_type() != WireType::Varint || reserved.is_some() { + return Err(InvalidSessionError("invalid reserved sender chain index").into()); + } let value = decode_varint(&mut buf) .map_err(|_| InvalidSessionError("invalid reserved sender chain index"))?; - reserved = u32::try_from(value) - .map_err(|_| InvalidSessionError("invalid reserved sender chain index"))?; + reserved = Some( + u32::try_from(value) + .map_err(|_| InvalidSessionError("invalid reserved sender chain index"))?, + ); } else { skip_field(tag, &mut buf) .map_err(|_| InvalidSessionError("failed to decode session record protobuf"))?; } } - Ok(reserved) + Ok(reserved.unwrap_or(0)) } /// If there's a session with a matching version and `alice_base_key`, ensures that it is the @@ -1296,6 +1313,45 @@ mod tests { ); } + /// A lease field that is unreadable — wrong wire type, or duplicated so + /// "last one wins" could pick a lower ceiling — must fail the load. Were + /// it skipped, the record would load with reservation 0, silently + /// re-enabling every counter the real lease had already spent. + #[test] + fn deserialize_rejects_a_malformed_or_duplicated_lease_field() { + use buffa::encoding::{Tag, WireType, encode_varint}; + + let mut csprng = rng(); + let base_key = KeyPair::generate(&mut csprng).public_key; + let mut record = SessionRecord::new(create_test_session_state(3, &base_key)); + record.reserve_sender_chain_counters(0); + let good = record.serialize().unwrap(); + assert!( + SessionRecord::deserialize(&good).is_ok(), + "the well-formed record must still load" + ); + + // Same field number, non-varint wire type. + let mut wrong_type = record.serialize().unwrap(); + Tag::new(RESERVED_SENDER_CHAIN_INDEX_FIELD, WireType::LengthDelimited) + .encode(&mut wrong_type); + encode_varint(0, &mut wrong_type); // zero-length payload + assert!( + SessionRecord::deserialize(&wrong_type).is_err(), + "a non-varint lease field must fail the load, not be skipped" + ); + + // Duplicated field: a trailing 0 would win under last-one-wins and + // wipe the ceiling. + let mut duplicated = record.serialize().unwrap(); + Tag::new(RESERVED_SENDER_CHAIN_INDEX_FIELD, WireType::Varint).encode(&mut duplicated); + encode_varint(0, &mut duplicated); + assert!( + SessionRecord::deserialize(&duplicated).is_err(), + "a duplicated lease field must fail the load rather than lower the ceiling" + ); + } + /// Creates a SessionRecord with N previous sessions for testing. fn create_record_with_previous_sessions(count: usize) -> SessionRecord { let mut csprng = rng(); diff --git a/wacore/libsignal/tests/counter_lease.rs b/wacore/libsignal/tests/counter_lease.rs index bd2afcf6c..1769469ec 100644 --- a/wacore/libsignal/tests/counter_lease.rs +++ b/wacore/libsignal/tests/counter_lease.rs @@ -325,6 +325,45 @@ fn crash_reload(peer: &mut Peer, remote: &ProtocolAddress, snapshot: &[u8]) { // ---- scenarios -------------------------------------------------------------- +/// Guard for the fixture itself: `make_rng` seeds a StdRng from the OS/thread +/// entropy source, so each peer must get independent key material. Were it +/// deterministic, both peers would agree on keys by accident and every +/// assertion below would pass for the wrong reason. +#[test] +fn peers_generate_independent_keys() { + let alice = Peer::new("alice"); + let bob = Peer::new("bob"); + + assert_ne!( + alice + .identity_store + .identity_key_pair + .identity_key() + .serialize(), + bob.identity_store + .identity_key_pair + .identity_key() + .serialize(), + "peers must not share an identity key" + ); + let key_of = |p: &Peer| { + BUNDLES.with(|b| { + b.borrow() + .get(&p.address) + .expect("bundle") + .pre_key_public() + .expect("bundle read") + .expect("one-time prekey") + .serialize() + }) + }; + assert_ne!( + key_of(&alice), + key_of(&bob), + "peers must not share a one-time prekey" + ); +} + /// The very first send on a fresh session must raise a lease and gate the /// wire on its durability. #[test] diff --git a/wacore/src/store/signal_cache.rs b/wacore/src/store/signal_cache.rs index 341650cd6..09a0020b7 100644 --- a/wacore/src/store/signal_cache.rs +++ b/wacore/src/store/signal_cache.rs @@ -157,8 +157,16 @@ impl SessionStoreState { self.cache.clear(); self.dirty.clear(); self.deleted.clear(); - // Cleared sessions reload from the backend, whose snapshot still - // carries its own (older) lease; the first send re-reserves. + // Dropping a pending gate here is only safe because a clear can never + // race a live wire: every caller runs either before the transport + // exists (connect) or after `cleanup_connection_state` has already + // taken the noise socket, and a send resolves that socket AFTER its + // pre-wire gate. So a send that observes this cleared set cannot reach + // send_node — it fails NotConnected — and its unpersisted lease dies + // with the ciphertext instead of reaching a peer. Moving a clear ahead + // of the socket teardown would silently reintroduce counter reuse for + // lease-raising sends. The reloaded snapshot carries its own older + // lease, and the first send re-reserves from there. self.reservation_pending.clear(); }