From 696d0986e01fc92aedcb8d0aaa0dabdb74753f80 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:07:00 -0300 Subject: [PATCH 1/2] feat(libsignal): let a consumer opt out of counter leasing Counters are leased in batches so the send path needs one durable flush per batch, and a reload fast-forwards past the whole lease. That only works because there is somewhere to persist the ceiling. A consumer whose storage is a component export has nowhere: into_components has to materialize the reservation before handing the record over, so the whole batch burns on every export rather than once per batch. Four consecutive DM sends land on the wire at counters 0, 64, 128, 192, and the peer buffers 63 skipped keys for each of them. A consumer whose persistence is already synchronous and durable before the ciphertext reaches the wire gets nothing from the lease and pays all of that. SessionRecord::waive_counter_lease and its SenderKeyRecord counterpart are how it says so. Nothing is inferred: the same record shape can be persisted by a consumer that wants the lease and by one that does not, so deriving the policy from the representation would turn a storage change into a silent change of guarantee. The two records now share a CounterLease enum whose Waived variant carries neither ceiling nor pending flag, so a record cannot hold a reservation the send path would gate on while having waived the lease that reservation implements. That covers creation, not just export: a waived record reserves nothing, so no ciphertext is gated on a flush that no longer protects anything. Waiving gives up a real guarantee. Without the lease, a crash between the encrypt and the write can reissue a counter and with it the (key, IV) pair. A record loaded from a snapshot written under the lease still carries a reservation that may already have been published, so waiving materializes it once, across archived states too, and then runs consecutively. The default is untouched. --- wacore/libsignal/src/protocol/consts.rs | 3 + .../libsignal/src/protocol/counter_lease.rs | 182 +++++++++++++++++ wacore/libsignal/src/protocol/group_cipher.rs | 142 ++++++++++++- wacore/libsignal/src/protocol/mod.rs | 1 + wacore/libsignal/src/protocol/sender_keys.rs | 83 ++++---- .../libsignal/src/protocol/session_cipher.rs | 7 +- .../libsignal/src/protocol/state/session.rs | 99 +++++---- wacore/libsignal/tests/counter_lease.rs | 190 ++++++++++++++++++ 8 files changed, 620 insertions(+), 87 deletions(-) create mode 100644 wacore/libsignal/src/protocol/counter_lease.rs diff --git a/wacore/libsignal/src/protocol/consts.rs b/wacore/libsignal/src/protocol/consts.rs index 92d118930..69fedcfb7 100644 --- a/wacore/libsignal/src/protocol/consts.rs +++ b/wacore/libsignal/src/protocol/consts.rs @@ -33,6 +33,9 @@ pub const MESSAGE_KEY_PRUNE_THRESHOLD: usize = 50; /// 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. +/// +/// None of this applies to a record whose consumer waived the lease: it +/// reserves nothing, so there is no batch to fast-forward past and no gap. pub const SENDER_CHAIN_RESERVATION_BATCH: u32 = 64; /// Upper bound for the reservation fast-forward on load. A legitimate lease diff --git a/wacore/libsignal/src/protocol/counter_lease.rs b/wacore/libsignal/src/protocol/counter_lease.rs new file mode 100644 index 000000000..3c1f00a8f --- /dev/null +++ b/wacore/libsignal/src/protocol/counter_lease.rs @@ -0,0 +1,182 @@ +//! Whether a record leases outbound counters ahead of durability. +//! +//! Message keys and IVs are derived deterministically from an outbound +//! counter, so republishing one after a crash reuses a (key, IV) pair. The +//! default guards that by leasing counters in batches: the send path needs a +//! durable flush only when a batch runs out, and any reload fast-forwards past +//! the whole lease. +//! +//! A consumer whose persistence is already synchronous and durable before the +//! ciphertext reaches the wire gets nothing from the lease and pays for it, +//! because every export has to burn the reserved range. [`CounterLease::Waived`] +//! is that consumer's declaration, and it is never inferred: the same record +//! shape can be persisted by a consumer that wants the lease and by one that +//! does not, so tying the policy to the representation would turn a storage +//! change into a silent change of guarantee. +//! +//! Waiving gives up a real guarantee. Without the lease, a crash between the +//! encrypt and the write can reissue a counter and with it the (key, IV) pair. +//! Only a consumer that persists before the wire can make that trade. + +use crate::protocol::consts; + +/// Reservation state, or the consumer's declaration that it needs none. +/// +/// `Waived` carries no ceiling and no pending flag, so a record cannot hold a +/// reservation the send path would gate on while also having waived the lease +/// that reservation implements. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum CounterLease { + Leased { + /// Exclusive ceiling of counters a durable snapshot may already have + /// published. + ceiling: u32, + /// A reservation was raised but not yet durably flushed. While set, + /// the owning ciphertext must not reach the wire. + pending_flush: bool, + }, + Waived, +} + +impl Default for CounterLease { + fn default() -> Self { + Self::Leased { + ceiling: 0, + pending_flush: false, + } + } +} + +impl CounterLease { + pub(crate) fn from_persisted_ceiling(ceiling: u32) -> Self { + Self::Leased { + ceiling, + pending_flush: false, + } + } + + /// Exclusive ceiling of the current reservation; zero when nothing is + /// reserved, which is also what a waived lease reports, since there is + /// nothing for a reload or an export to advance past. + pub(crate) fn ceiling(self) -> u32 { + match self { + Self::Leased { ceiling, .. } => ceiling, + Self::Waived => 0, + } + } + + pub(crate) fn is_pending_flush(self) -> bool { + matches!( + self, + Self::Leased { + pending_flush: true, + .. + } + ) + } + + /// Lease a fresh batch after `spent_counter` was issued. + /// + /// Reservations only ever rise: a counter must not be published under a + /// ceiling lower than one a durable snapshot already carries, so a spent + /// counter still inside the current batch changes nothing. + pub(crate) fn reserve(&mut self, spent_counter: u32) { + if let Self::Leased { + ceiling, + pending_flush, + } = self + && spent_counter >= *ceiling + { + *ceiling = spent_counter.saturating_add(consts::SENDER_CHAIN_RESERVATION_BATCH); + *pending_flush = true; + } + } + + pub(crate) fn set_pending_flush(&mut self, pending: bool) { + if let Self::Leased { pending_flush, .. } = self { + *pending_flush = pending; + } + } + + /// Drop the reservation without changing whether the lease is in force. + pub(crate) fn clear_reservation(&mut self) { + if let Self::Leased { ceiling, .. } = self { + *ceiling = 0; + } + } + + /// Cap the ceiling at one batch after a ratchet replaced the leased chain. + pub(crate) fn rebase(&mut self) { + if let Self::Leased { ceiling, .. } = self { + *ceiling = (*ceiling).min(consts::SENDER_CHAIN_RESERVATION_BATCH); + } + } + + /// Waive the lease, returning any ceiling the caller must still materialize. + /// + /// A record loaded from a snapshot written while the lease was in force + /// carries a reservation that may already have been published. Waiving does + /// not make that untrue, so the caller advances past it once; from then on + /// counters are consecutive. Refusing such a record instead would strand + /// the address for good. + pub(crate) fn waive(&mut self) -> u32 { + let ceiling = self.ceiling(); + *self = Self::Waived; + ceiling + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_waived_lease_reserves_nothing_and_never_gates_the_wire() { + let mut lease = CounterLease::Waived; + lease.reserve(0); + lease.reserve(u32::MAX); + + assert_eq!(lease, CounterLease::Waived); + assert_eq!(lease.ceiling(), 0); + assert!(!lease.is_pending_flush()); + } + + #[test] + fn a_leased_reservation_rises_and_gates_the_wire() { + let mut lease = CounterLease::default(); + assert_eq!(lease.ceiling(), 0); + assert!(!lease.is_pending_flush()); + + lease.reserve(0); + assert_eq!(lease.ceiling(), consts::SENDER_CHAIN_RESERVATION_BATCH); + assert!(lease.is_pending_flush()); + + // Inside the batch: nothing to raise, so the send stays uncovered by a + // fresh gate. + lease.set_pending_flush(false); + lease.reserve(1); + assert_eq!(lease.ceiling(), consts::SENDER_CHAIN_RESERVATION_BATCH); + assert!(!lease.is_pending_flush()); + + lease.reserve(consts::SENDER_CHAIN_RESERVATION_BATCH); + assert_eq!(lease.ceiling(), consts::SENDER_CHAIN_RESERVATION_BATCH * 2); + assert!(lease.is_pending_flush()); + } + + #[test] + fn waiving_hands_back_the_ceiling_to_materialize_once() { + let mut lease = CounterLease::from_persisted_ceiling(512); + + assert_eq!(lease.waive(), 512); + assert_eq!(lease, CounterLease::Waived); + // Converged: a second waive has nothing left to materialize. + assert_eq!(lease.waive(), 0); + } + + #[test] + fn rebasing_a_waived_lease_is_inert() { + let mut lease = CounterLease::Waived; + lease.rebase(); + assert_eq!(lease, CounterLease::Waived); + } +} diff --git a/wacore/libsignal/src/protocol/group_cipher.rs b/wacore/libsignal/src/protocol/group_cipher.rs index 90c1729e5..5f2b929af 100644 --- a/wacore/libsignal/src/protocol/group_cipher.rs +++ b/wacore/libsignal/src/protocol/group_cipher.rs @@ -142,12 +142,10 @@ pub async fn group_encrypt( // sends ride the coalesced write-behind; only the send that reaches the // ceiling re-reserves and gates the ciphertext on a synchronous flush (which // fast-forwards past the reservation after any reload). Decrypt-side advances - // stay ungated (they re-derive forward). Mirrors the DM counter lease in - // SessionRecord. - let spent_iteration = message_keys.iteration(); - if spent_iteration >= record.reserved_iteration() { - record.reserve_iterations(spent_iteration); - } + // stay ungated (they re-derive forward). A consumer that waived the lease + // persists before the wire and reserves nothing. Mirrors the DM counter + // lease in SessionRecord. + record.reserve_iterations(message_keys.iteration()); sender_key_store .store_sender_key(sender_key_name, record) @@ -560,6 +558,138 @@ mod tests { ); } + /// A store whose persistence is a component export, the group counterpart + /// of the DM case in `tests/counter_lease.rs`: it rebuilds the record from + /// components on every load, which is exactly what materializes a + /// reservation and burns the batch. + struct ComponentStore { + states: HashMap, + waive: bool, + } + #[async_trait] + impl SenderKeyStore for ComponentStore { + async fn store_sender_key( + &mut self, + name: &SenderKeyName, + record: SenderKeyRecord, + ) -> Result<()> { + self.states.insert(name.clone(), record.into_components()?); + Ok(()) + } + async fn load_sender_key(&self, name: &SenderKeyName) -> Result> { + let Some(components) = self.states.get(name) else { + return Ok(None); + }; + let mut record = SenderKeyRecord::from_components(components.clone())?; + if self.waive { + record.waive_counter_lease()?; + } + Ok(Some(record)) + } + } + + /// Iterations the group sender put on the wire over `count` sends, and the + /// skipped keys the receiver had to buffer for them. + fn exported_group_run(count: usize, waive: bool) -> (Vec, usize) { + let mut rng = rand::rng(); + let name = SenderKeyName::new("group@g.us".to_string(), "bob.0".to_string()); + let mut bob = ComponentStore { + states: HashMap::new(), + waive, + }; + let skdm = block_on(create_sender_key_distribution_message( + &name, &mut bob, &mut rng, + )) + .expect("bob creates his distribution message"); + let mut alice = InMemorySenderKeyStore { + keys: HashMap::new(), + }; + block_on(process_sender_key_distribution_message( + &name, &skdm, &mut alice, + )) + .expect("alice processes it"); + + let iterations = (0..count) + .map(|_| { + let msg = + block_on(group_encrypt(&mut bob, &name, b"m", &mut rng)).expect("bob encrypts"); + let iteration = msg.iteration(); + block_on(group_decrypt(msg.serialized(), &mut alice, &name)) + .expect("alice decrypts"); + iteration + }) + .collect(); + + let skipped = alice + .keys + .remove(&name) + .expect("alice has a record") + .into_components() + .expect("components") + .states + .iter() + .map(|state| state.message_keys.len()) + .sum(); + (iterations, skipped) + } + + /// The group symptom, and its fix: exporting components burns a batch per + /// send, so iterations stride by 64 and the receiver buffers the gap. + #[test] + fn a_waived_lease_keeps_group_iterations_consecutive() { + let (iterations, skipped) = exported_group_run(8, true); + + assert_eq!(iterations, (0..8).collect::>()); + assert_eq!(skipped, 0); + } + + /// The default is untouched: same run, same batch stride, same backlog. + #[test] + fn the_default_group_lease_still_burns_a_batch_per_export() { + let (iterations, skipped) = exported_group_run(8, false); + + let expected: Vec = (0..8) + .map(|i| i as u32 * SENDER_CHAIN_RESERVATION_BATCH) + .collect(); + assert_eq!(iterations, expected); + assert!( + skipped > 0, + "the leased run must leave the receiver with skipped keys" + ); + } + + /// A record written under the lease may already have published iterations + /// below its ceiling; waiving materializes that ceiling once, then runs + /// consecutively. + #[test] + fn waiving_materializes_a_previously_reserved_group_ceiling_once() { + let mut rng = rand::rng(); + let name = SenderKeyName::new("group@g.us".to_string(), "bob.0".to_string()); + let mut bob = InMemorySenderKeyStore { + keys: HashMap::new(), + }; + block_on(create_sender_key_distribution_message( + &name, &mut bob, &mut rng, + )) + .expect("distribution message"); + block_on(group_encrypt(&mut bob, &name, b"m0", &mut rng)).expect("first send reserves"); + + let record = bob.keys.get_mut(&name).expect("record"); + let ceiling = record.reserved_iteration(); + assert_eq!(ceiling, SENDER_CHAIN_RESERVATION_BATCH); + record + .waive_counter_lease() + .expect("waive materializes once"); + assert_eq!(record.reserved_iteration(), 0); + + let first = block_on(group_encrypt(&mut bob, &name, b"after", &mut rng)) + .expect("send after waiving"); + assert_eq!(first.iteration(), ceiling); + let second = + block_on(group_encrypt(&mut bob, &name, b"next", &mut rng)).expect("the next send"); + assert_eq!(second.iteration(), ceiling + 1); + } + /// A store that emulates the real signal-cache gate: it records whether each /// stored advance was wire-gated and clears the transient flag, so a run of /// sends can be counted for gate frequency. diff --git a/wacore/libsignal/src/protocol/mod.rs b/wacore/libsignal/src/protocol/mod.rs index 133659784..92c9f33e0 100644 --- a/wacore/libsignal/src/protocol/mod.rs +++ b/wacore/libsignal/src/protocol/mod.rs @@ -18,6 +18,7 @@ #![deny(unsafe_code)] pub mod consts; +mod counter_lease; mod crypto; pub mod error; mod group_cipher; diff --git a/wacore/libsignal/src/protocol/sender_keys.rs b/wacore/libsignal/src/protocol/sender_keys.rs index 9856e0c13..9d52e5b1e 100644 --- a/wacore/libsignal/src/protocol/sender_keys.rs +++ b/wacore/libsignal/src/protocol/sender_keys.rs @@ -10,6 +10,7 @@ use buffa::{Message, MessageField}; use hmac::{HmacReset, KeyInit, Mac}; use sha2::Sha256; +use crate::protocol::counter_lease::CounterLease; use crate::protocol::crypto::hmac_sha256; use crate::protocol::record_components::{ SenderKeyRecordComponents, sender_state_components_from_structure, @@ -493,21 +494,16 @@ 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, - /// Durably-reserved iteration ceiling for the current state's sender chain, - /// mirroring `SessionRecord::reserved_sender_chain_index` for DM. Iterations - /// below this ceiling are covered by a persisted reservation, so their sends - /// skip the synchronous pre-wire flush and ride the coalesced write-behind; - /// only the send that raises the ceiling gates. A reload fast-forwards the - /// current chain past this ceiling so no possibly-spent iteration is - /// re-derivable. Reset to 0 on any state change (rotation/promotion), which - /// forces the next send to re-reserve and gate; never reuses an iteration. - reserved_iteration: u32, + /// Durability lease over sender-chain iterations, mirroring + /// `SessionRecord`'s for DM, or the consumer's declaration that it needs + /// none. Iterations below the reserved ceiling ride the coalesced + /// write-behind; only the send that raises it gates the wire, and a reload + /// fast-forwards past it so no possibly-spent iteration is re-derivable. + /// Reset to 0 on any state change (rotation/promotion), which forces the + /// next send to re-reserve and gate. The wire-gate flag it carries is + /// transient and never serialized; the store layer converts it into flush + /// gating. + lease: CounterLease, } /// Local-only field appended to the serialized record for `reserved_iteration`. @@ -520,8 +516,7 @@ impl SenderKeyRecord { pub fn new_empty() -> Self { Self { states: VecDeque::with_capacity(consts::MAX_SENDER_KEY_STATES), - wire_gated: false, - reserved_iteration: 0, + lease: CounterLease::default(), } } @@ -542,8 +537,7 @@ impl SenderKeyRecord { Ok(Self { states, - wire_gated: false, - reserved_iteration: 0, + lease: CounterLease::default(), }) } @@ -553,10 +547,11 @@ impl SenderKeyRecord { /// before export so rebuilding the record cannot derive a possibly spent /// message key again. pub fn into_components(mut self) -> Result { - if self.reserved_iteration > 0 + let reserved_iteration = self.lease.ceiling(); + if reserved_iteration > 0 && let Some(state) = self.states.front_mut() { - state.fast_forward_sender_chain(self.reserved_iteration)?; + state.fast_forward_sender_chain(reserved_iteration)?; } let states = self .states @@ -574,7 +569,25 @@ impl SenderKeyRecord { /// Iterations strictly below this ceiling are covered by a durable /// reservation and their sends need no synchronous flush. pub fn reserved_iteration(&self) -> u32 { - self.reserved_iteration + self.lease.ceiling() + } + + /// Waive counter leasing on this record, declaring that the consumer's + /// persistence is synchronous and durable before the ciphertext reaches + /// the wire. + /// + /// The group counterpart of `SessionRecord::waive_counter_lease`, applied + /// per load and never inferred. A snapshot written under the lease is + /// materialized once here, since its reservation may already have been + /// published. The guarantee being given up is stated on `CounterLease`. + pub fn waive_counter_lease(&mut self) -> Result<(), SignalProtocolError> { + let ceiling = self.lease.waive(); + if ceiling > 0 + && let Some(state) = self.states.front_mut() + { + state.fast_forward_sender_chain(ceiling)?; + } + Ok(()) } /// Lease a fresh batch of iterations after `spent_iteration` reached the @@ -582,9 +595,7 @@ impl SenderKeyRecord { /// must not hit the wire until a flush persists the raised ceiling. Mirrors /// `SessionRecord::reserve_sender_chain_counters`. pub fn reserve_iterations(&mut self, spent_iteration: u32) { - self.reserved_iteration = - spent_iteration.saturating_add(consts::SENDER_CHAIN_RESERVATION_BATCH); - self.wire_gated = true; + self.lease.reserve(spent_iteration); } pub fn deserialize(buf: &[u8]) -> Result { @@ -645,23 +656,22 @@ impl SenderKeyRecord { Ok(Self { states, - wire_gated: false, - reserved_iteration, + lease: CounterLease::from_persisted_ceiling(reserved_iteration), }) } /// 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; + self.lease.set_pending_flush(true); } pub fn is_wire_gated(&self) -> bool { - self.wire_gated + self.lease.is_pending_flush() } pub fn clear_wire_gated(&mut self) { - self.wire_gated = false; + self.lease.set_pending_flush(false); } pub fn sender_key_state(&self) -> Result<&SenderKeyState, InvalidSenderKeySessionError> { @@ -739,7 +749,7 @@ impl SenderKeyRecord { // the sending record only reaches here on first creation (reservation // already 0, warm sends reuse the record without re-adding), and receiver // records never carry a reservation. - self.reserved_iteration = 0; + self.lease.clear_reservation(); Ok(()) } @@ -794,9 +804,10 @@ impl SenderKeyRecord { use buffa::encoding::{Tag, WireType, encode_varint, varint_len}; let mut buf = waproto::codec::sender_key_record_to_vec(&self.as_protobuf()); - let incarnation = incarnation.filter(|_| self.reserved_iteration > 0); - let reservation_len = if self.reserved_iteration > 0 { - 2 + varint_len(self.reserved_iteration as u64) + let reserved_iteration = self.lease.ceiling(); + let incarnation = incarnation.filter(|_| reserved_iteration > 0); + let reservation_len = if reserved_iteration > 0 { + 2 + varint_len(reserved_iteration as u64) } else { 0 }; @@ -807,9 +818,9 @@ impl SenderKeyRecord { // Append the local-only reservation as a top-level field the generated // decoder skips. Emitted only when non-zero, so legacy/unreserved records // stay byte-identical. Mirrors SessionRecord::serialize_into. - if self.reserved_iteration > 0 { + if reserved_iteration > 0 { Tag::new(RESERVED_ITERATION_FIELD, WireType::Varint).encode(&mut buf); - encode_varint(self.reserved_iteration as u64, &mut buf); + encode_varint(reserved_iteration as u64, &mut buf); } if let Some(incarnation) = incarnation { super::local_field::encode_store_incarnation(&mut buf, incarnation); diff --git a/wacore/libsignal/src/protocol/session_cipher.rs b/wacore/libsignal/src/protocol/session_cipher.rs index dcf8b9dba..ef31edccd 100644 --- a/wacore/libsignal/src/protocol/session_cipher.rs +++ b/wacore/libsignal/src/protocol/session_cipher.rs @@ -401,10 +401,9 @@ async fn message_encrypt_inner( // 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()); - } + // lease, so this counter can never be re-derived after a crash. A consumer + // that waived the lease persists before the wire and reserves nothing. + 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 18dc28b12..0ae1a9d2d 100644 --- a/wacore/libsignal/src/protocol/state/session.rs +++ b/wacore/libsignal/src/protocol/state/session.rs @@ -11,6 +11,7 @@ use buffa::{Message, MessageField}; use subtle::ConstantTimeEq; use crate::core::curve::KeyType; +use crate::protocol::counter_lease::CounterLease; use crate::protocol::ratchet::keys::MessageKeyGenerator; use crate::protocol::ratchet::{ChainKey, RootKey}; use crate::protocol::record_components::{ @@ -752,17 +753,13 @@ const RESERVED_SENDER_CHAIN_INDEX_FIELD: u32 = 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, + /// Durability lease over sender-chain counters, or the consumer's + /// declaration that it persists before the wire and needs none. Any state + /// entering service from a snapshot must fast-forward past a live lease + /// first — message keys and IVs are derived deterministically from the + /// counter, so re-deriving a spent counter reuses a (key, IV) pair. The + /// pending flag it carries is transient and never serialized. + lease: CounterLease, } impl SessionRecord { @@ -770,8 +767,7 @@ impl SessionRecord { Self { current_session: None, previous_sessions: Arc::new(Vec::new()), - reserved_sender_chain_index: 0, - pending_reservation: false, + lease: CounterLease::default(), } } @@ -779,8 +775,7 @@ impl SessionRecord { Self { current_session: Some(state), previous_sessions: Arc::new(Vec::new()), - reserved_sender_chain_index: 0, - pending_reservation: false, + lease: CounterLease::default(), } } @@ -805,8 +800,7 @@ impl SessionRecord { Ok(Self { current_session, previous_sessions: Arc::new(previous_sessions), - reserved_sender_chain_index: 0, - pending_reservation: false, + lease: CounterLease::default(), }) } @@ -815,9 +809,10 @@ impl SessionRecord { /// Any durably reserved sender range is advanced to its exclusive ceiling /// before export so rebuilding the record cannot derive a possibly spent /// message key again. A chain too stale to advance is dropped fail-closed - /// without discarding the rest of the record. + /// without discarding the rest of the record. A waived lease reserves + /// nothing, so there is nothing to advance past. pub fn into_components(mut self) -> Result { - let reserved_sender_chain_index = self.reserved_sender_chain_index; + let reserved_sender_chain_index = self.lease.ceiling(); if reserved_sender_chain_index > 0 && let Some(state) = self.current_session.as_mut() { @@ -847,17 +842,43 @@ impl SessionRecord { } pub fn reserved_sender_chain_index(&self) -> u32 { - self.reserved_sender_chain_index + self.lease.ceiling() + } + + /// Waive counter leasing on this record, declaring that the consumer's + /// persistence is synchronous and durable before the ciphertext reaches + /// the wire. + /// + /// The consumer applies this to every record it loads; nothing is inferred + /// from the stored representation. A snapshot written while the lease was + /// in force still carries a reservation that may already have been + /// published, so it is materialized once here before the lease goes away. + /// The guarantee being given up is stated on [`CounterLease`]. + pub fn waive_counter_lease(&mut self) { + let ceiling = self.lease.waive(); + if ceiling == 0 { + return; + } + // Every state the lease covered has to burn it, archived ones included: + // once the ceiling is gone, a later promotion has nothing left to tell + // it those counters may already be on the wire. + if let Some(state) = self.current_session.as_mut() { + state.fast_forward_sender_chain_or_drop(ceiling); + } + for session in Arc::make_mut(&mut self.previous_sessions) { + let mut state = SessionState::from_session_structure(std::mem::take(session)); + state.fast_forward_sender_chain_or_drop(ceiling); + *session = state.session; + } } /// 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. + /// raised ceiling. A counter still inside the batch, or a waived lease, + /// leaves both untouched. 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; + self.lease.reserve(spent_counter); } /// Rebase the lease after a DH ratchet replaced the leased sender chain @@ -883,19 +904,17 @@ impl SessionRecord { // Never raises: a counter must not be published under a ceiling that // is not yet durable. An in-chain lease is always within one batch of // the live index, so this is a no-op outside a chain replacement. - self.reserved_sender_chain_index = self - .reserved_sender_chain_index - .min(consts::SENDER_CHAIN_RESERVATION_BATCH); + self.lease.rebase(); } pub fn has_pending_reservation(&self) -> bool { - self.pending_reservation + self.lease.is_pending_flush() } /// 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; + self.lease.set_pending_flush(false); } pub fn deserialize(bytes: &[u8]) -> Result { @@ -947,8 +966,7 @@ impl SessionRecord { .map_err(|_| InvalidSessionError("failed to decode current session protobuf"))? .map(Into::into), previous_sessions: Arc::new(previous_sessions), - reserved_sender_chain_index: local_fields.reservation, - pending_reservation: false, + lease: CounterLease::from_persisted_ceiling(local_fields.reservation), }; let trusted_reload = @@ -956,10 +974,10 @@ impl SessionRecord { // An untrusted snapshot may predate sends covered by its lease. if !trusted_reload - && record.reserved_sender_chain_index > 0 + && record.lease.ceiling() > 0 && let Some(state) = record.current_session.as_mut() { - state.fast_forward_sender_chain(record.reserved_sender_chain_index)?; + state.fast_forward_sender_chain(record.lease.ceiling())?; } Ok(record) @@ -1105,8 +1123,8 @@ impl SessionRecord { 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 { - state.fast_forward_sender_chain_or_drop(self.reserved_sender_chain_index); + if self.lease.ceiling() > 0 { + state.fast_forward_sender_chain_or_drop(self.lease.ceiling()); } self.current_session = Some(state); } @@ -1117,14 +1135,14 @@ impl SessionRecord { /// archived before resetting it for the fresh chain; otherwise the archive /// could later reissue a counter covered by the discarded lease. pub fn promote_fresh_state(&mut self, new_state: SessionState) { - if self.reserved_sender_chain_index > 0 + if self.lease.ceiling() > 0 && let Some(state) = self.current_session.as_mut() { - state.fast_forward_sender_chain_or_drop(self.reserved_sender_chain_index); + state.fast_forward_sender_chain_or_drop(self.lease.ceiling()); } self.archive_current_state_inner(); self.current_session = Some(new_state); - self.reserved_sender_chain_index = 0; + self.lease.clear_reservation(); } fn archive_current_state_inner(&mut self) -> bool { @@ -1203,7 +1221,7 @@ impl SessionRecord { }) .sum(); - let reserved = self.reserved_sender_chain_index; + let reserved = self.lease.ceiling(); let incarnation = incarnation.filter(|_| reserved > 0); let reserved_len = if reserved > 0 { 2 + varint_len(reserved as u64) @@ -2075,8 +2093,7 @@ 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, + lease: CounterLease::default(), }; 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 index 73c8f01ef..c14387d05 100644 --- a/wacore/libsignal/tests/counter_lease.rs +++ b/wacore/libsignal/tests/counter_lease.rs @@ -710,3 +710,193 @@ fn a_rebased_lease_never_republishes_a_counter_across_a_crash() { let pt = receive(&mut bob, &alice.address, &ct).expect("bob decrypts across the burned gap"); assert_eq!(&pt[..], b"after crash"); } + +// ---- waiver ----------------------------------------------------------------- + +/// Stand in for a consumer whose persistence is a component export: it rebuilds +/// the record from components on every load, and waives the lease because its +/// writes are durable before the ciphertext reaches the wire. +fn components_roundtrip(peer: &mut Peer, remote: &ProtocolAddress, waive: bool) { + let record = peer.session_store.0.remove(remote).expect("session exists"); + let mut rebuilt = + SessionRecord::from_components(record.into_components().expect("export")).expect("import"); + if waive { + rebuilt.waive_counter_lease(); + } + peer.session_store.0.insert(remote.clone(), rebuilt); +} + +/// Skipped message keys the peer had to buffer, across every receiver chain. +fn skipped_keys(peer: &Peer, remote: &ProtocolAddress) -> usize { + record_of(peer, remote) + .into_components() + .expect("components") + .current_session + .expect("current session") + .receiver_chains + .iter() + .map(|chain| chain.message_keys.len()) + .sum() +} + +/// Drive `count` sends, exporting and reimporting components between each, and +/// return the wire counters the peer saw. +fn exported_send_run(alice: &mut Peer, bob: &mut Peer, count: u32, waive: bool) -> Vec { + let bob_address = bob.address.clone(); + let alice_address = alice.address.clone(); + process_bundle(alice, &bob_address); + components_roundtrip(alice, &bob_address, waive); + + let mut counters = Vec::with_capacity(count as usize); + for _ in 0..count { + let ct = send(alice, &bob_address, b"m"); + counters.push(wire_counter(&ct)); + receive(bob, &alice_address, &ct).expect("peer decrypts"); + components_roundtrip(alice, &bob_address, waive); + } + counters +} + +/// The symptom: a consumer that exports components burns a whole batch per +/// export, so consecutive sends land 64 apart and the peer buffers 63 skipped +/// keys for each one. Under a waived lease the counters are consecutive and +/// nothing is skipped. +#[test] +fn a_waived_lease_keeps_exported_counters_consecutive() { + let mut alice = Peer::new("alice-waived"); + let mut bob = Peer::new("bob-waived"); + + let counters = exported_send_run(&mut alice, &mut bob, 8, true); + + assert_eq!(counters, (0..8).collect::>()); + assert_eq!(skipped_keys(&bob, &alice.address.clone()), 0); +} + +/// The default is untouched: the reservation is still created, so an export +/// still materializes it and the counters still stride by a batch. +#[test] +fn the_default_lease_still_burns_a_batch_per_export() { + let mut alice = Peer::new("alice-leased"); + let mut bob = Peer::new("bob-leased"); + + let counters = exported_send_run(&mut alice, &mut bob, 4, false); + + let batch = SENDER_CHAIN_RESERVATION_BATCH; + assert_eq!(counters, vec![0, batch, batch * 2, batch * 3]); + assert!( + skipped_keys(&bob, &alice.address.clone()) > 0, + "the leased run must leave the peer with skipped keys" + ); +} + +/// The waiver removes the reservation, not just its materialization: a send +/// under it must not gate the ciphertext on a flush that no longer protects +/// anything. +#[test] +fn a_waived_lease_never_gates_the_wire() { + let mut alice = Peer::new("alice-ungated"); + let mut bob = Peer::new("bob-ungated"); + let bob_address = bob.address.clone(); + + process_bundle(&mut alice, &bob_address); + components_roundtrip(&mut alice, &bob_address, true); + for _ in 0..3 { + let ct = send(&mut alice, &bob_address, b"m"); + receive(&mut bob, &alice.address.clone(), &ct).expect("peer decrypts"); + let record = record_of(&alice, &bob_address); + assert_eq!(record.reserved_sender_chain_index(), 0); + assert!(!record.has_pending_reservation()); + } +} + +/// Under the default, the same run keeps gating and reserving. +#[test] +fn the_default_lease_still_gates_the_wire() { + let mut alice = Peer::new("alice-gated"); + let mut bob = Peer::new("bob-gated"); + let bob_address = bob.address.clone(); + + process_bundle(&mut alice, &bob_address); + let ct = send(&mut alice, &bob_address, b"m"); + receive(&mut bob, &alice.address.clone(), &ct).expect("peer decrypts"); + + let record = record_of(&alice, &bob_address); + assert_eq!( + record.reserved_sender_chain_index(), + SENDER_CHAIN_RESERVATION_BATCH + ); + assert!(record.has_pending_reservation()); +} + +/// A record written while the lease was in force may already have published +/// counters below its ceiling, and waiving does not make that untrue. The +/// ceiling is materialized once, then counters run consecutively from there. +#[test] +fn waiving_materializes_a_previously_reserved_ceiling_once() { + let mut alice = Peer::new("alice-preexisting"); + let mut bob = Peer::new("bob-preexisting"); + let bob_address = bob.address.clone(); + let alice_address = alice.address.clone(); + + // Build up a real reservation under the default lease. + establish(&mut alice, &mut bob); + let ceiling = record_of(&alice, &bob_address).reserved_sender_chain_index(); + assert_eq!(ceiling, SENDER_CHAIN_RESERVATION_BATCH); + + // The consumer turns the waiver on and reloads the record it already had. + let record = alice + .session_store + .0 + .get_mut(&bob_address) + .expect("session"); + record.waive_counter_lease(); + assert_eq!(record.reserved_sender_chain_index(), 0); + + let first = send(&mut alice, &bob_address, b"after waiving"); + assert_eq!(wire_counter(&first), ceiling); + receive(&mut bob, &alice_address, &first).expect("peer decrypts across the burn"); + + let second = send(&mut alice, &bob_address, b"and the next"); + assert_eq!(wire_counter(&second), ceiling + 1); + receive(&mut bob, &alice_address, &second).expect("peer decrypts"); +} + +/// Archived states were covered by the same ceiling, so waiving has to burn +/// them too: once the lease is gone, a promotion has nothing left telling it +/// those counters may already be on the wire. +#[test] +fn waiving_burns_the_ceiling_into_archived_states_too() { + let mut alice = Peer::new("alice-archived"); + let mut bob = Peer::new("bob-archived"); + let bob_address = bob.address.clone(); + + establish(&mut alice, &mut bob); + let ceiling = record_of(&alice, &bob_address).reserved_sender_chain_index(); + assert_eq!(ceiling, SENDER_CHAIN_RESERVATION_BATCH); + + // Archive the leased state, then waive. + let record = alice + .session_store + .0 + .get_mut(&bob_address) + .expect("session"); + record + .archive_current_state() + .expect("archive the leased state"); + record.waive_counter_lease(); + + let archived_index = record_of(&alice, &bob_address) + .into_components() + .expect("components") + .previous_sessions + .first() + .expect("archived state") + .sender_chain + .as_ref() + .expect("sender chain") + .chain_key + .as_ref() + .expect("chain key") + .index; + assert_eq!(archived_index, Some(ceiling)); +} From 40e401dec6dbdd151ea69cc92d1eb541350aff48 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:23:00 -0300 Subject: [PATCH 2/2] fix(libsignal): keep the lease when a waiver cannot materialize it Two problems on the same path. The doc link from the public waive_counter_lease pointed at the crate-private CounterLease, which fails the rustdoc gate. The guarantee a consumer gives up belongs on the public surface that offers the choice anyway, so it moved to SessionRecord::waive_counter_lease and the group counterpart links there. More importantly, the sender-key waiver dropped the ceiling before advancing past it. A chain too stale to advance returned the error with the lease already gone, leaving the record free to reissue exactly the iterations that ceiling said may already be on the wire. Materialize first, waive only on success. --- .../libsignal/src/protocol/counter_lease.rs | 46 +++++++------------ wacore/libsignal/src/protocol/group_cipher.rs | 27 +++++++++++ wacore/libsignal/src/protocol/sender_keys.rs | 17 +++---- .../libsignal/src/protocol/state/session.rs | 27 ++++++++--- 4 files changed, 73 insertions(+), 44 deletions(-) diff --git a/wacore/libsignal/src/protocol/counter_lease.rs b/wacore/libsignal/src/protocol/counter_lease.rs index 3c1f00a8f..169b17186 100644 --- a/wacore/libsignal/src/protocol/counter_lease.rs +++ b/wacore/libsignal/src/protocol/counter_lease.rs @@ -1,22 +1,9 @@ //! Whether a record leases outbound counters ahead of durability. //! -//! Message keys and IVs are derived deterministically from an outbound -//! counter, so republishing one after a crash reuses a (key, IV) pair. The -//! default guards that by leasing counters in batches: the send path needs a -//! durable flush only when a batch runs out, and any reload fast-forwards past -//! the whole lease. -//! -//! A consumer whose persistence is already synchronous and durable before the -//! ciphertext reaches the wire gets nothing from the lease and pays for it, -//! because every export has to burn the reserved range. [`CounterLease::Waived`] -//! is that consumer's declaration, and it is never inferred: the same record -//! shape can be persisted by a consumer that wants the lease and by one that -//! does not, so tying the policy to the representation would turn a storage -//! change into a silent change of guarantee. -//! -//! Waiving gives up a real guarantee. Without the lease, a crash between the -//! encrypt and the write can reissue a counter and with it the (key, IV) pair. -//! Only a consumer that persists before the wire can make that trade. +//! Shared by `SessionRecord` and `SenderKeyRecord`, which lease the same way +//! over different units. What the policy means to a consumer, and what waiving +//! costs it, is documented on `SessionRecord::waive_counter_lease`, the public +//! surface that offers the choice. use crate::protocol::consts; @@ -112,17 +99,15 @@ impl CounterLease { } } - /// Waive the lease, returning any ceiling the caller must still materialize. + /// Drop the lease. /// - /// A record loaded from a snapshot written while the lease was in force - /// carries a reservation that may already have been published. Waiving does - /// not make that untrue, so the caller advances past it once; from then on - /// counters are consecutive. Refusing such a record instead would strand - /// the address for good. - pub(crate) fn waive(&mut self) -> u32 { - let ceiling = self.ceiling(); + /// The caller materializes [`Self::ceiling`] first and only gets here once + /// that succeeded: a record loaded from a snapshot written while the lease + /// was in force carries a reservation that may already have been published, + /// and dropping the ceiling before advancing past it would leave the record + /// free to reissue those counters. + pub(crate) fn waive(&mut self) { *self = Self::Waived; - ceiling } } @@ -164,13 +149,14 @@ mod tests { } #[test] - fn waiving_hands_back_the_ceiling_to_materialize_once() { + fn waiving_reports_nothing_left_to_materialize() { let mut lease = CounterLease::from_persisted_ceiling(512); + assert_eq!(lease.ceiling(), 512); - assert_eq!(lease.waive(), 512); + lease.waive(); assert_eq!(lease, CounterLease::Waived); - // Converged: a second waive has nothing left to materialize. - assert_eq!(lease.waive(), 0); + // Converged: nothing left for a later load to advance past. + assert_eq!(lease.ceiling(), 0); } #[test] diff --git a/wacore/libsignal/src/protocol/group_cipher.rs b/wacore/libsignal/src/protocol/group_cipher.rs index 5f2b929af..72992170b 100644 --- a/wacore/libsignal/src/protocol/group_cipher.rs +++ b/wacore/libsignal/src/protocol/group_cipher.rs @@ -690,6 +690,33 @@ mod tests { assert_eq!(second.iteration(), ceiling + 1); } + /// A ceiling too far ahead to advance past must leave the record on its + /// lease: dropping it there would free the record to reissue exactly the + /// iterations the ceiling says may already be on the wire. + #[test] + fn a_waiver_that_cannot_materialize_keeps_the_lease() { + let mut rng = rand::rng(); + let name = SenderKeyName::new("group@g.us".to_string(), "bob.0".to_string()); + let mut bob = InMemorySenderKeyStore { + keys: HashMap::new(), + }; + block_on(create_sender_key_distribution_message( + &name, &mut bob, &mut rng, + )) + .expect("distribution message"); + + let record = bob.keys.get_mut(&name).expect("record"); + record.reserve_iterations(consts::MAX_RESERVATION_FAST_FORWARD + 1); + let ceiling = record.reserved_iteration(); + + assert!(record.waive_counter_lease().is_err()); + assert_eq!( + record.reserved_iteration(), + ceiling, + "a failed waiver must not drop the ceiling" + ); + } + /// A store that emulates the real signal-cache gate: it records whether each /// stored advance was wire-gated and clears the transient flag, so a run of /// sends can be counted for gate frequency. diff --git a/wacore/libsignal/src/protocol/sender_keys.rs b/wacore/libsignal/src/protocol/sender_keys.rs index 9d52e5b1e..c294cd487 100644 --- a/wacore/libsignal/src/protocol/sender_keys.rs +++ b/wacore/libsignal/src/protocol/sender_keys.rs @@ -572,21 +572,22 @@ impl SenderKeyRecord { self.lease.ceiling() } - /// Waive counter leasing on this record, declaring that the consumer's - /// persistence is synchronous and durable before the ciphertext reaches - /// the wire. + /// Waive counter leasing on this record. /// - /// The group counterpart of `SessionRecord::waive_counter_lease`, applied - /// per load and never inferred. A snapshot written under the lease is - /// materialized once here, since its reservation may already have been - /// published. The guarantee being given up is stated on `CounterLease`. + /// The group counterpart of + /// [`SessionRecord::waive_counter_lease`](crate::protocol::SessionRecord::waive_counter_lease), + /// including the guarantee it gives up. pub fn waive_counter_lease(&mut self) -> Result<(), SignalProtocolError> { - let ceiling = self.lease.waive(); + // Materialize before dropping the ceiling: a chain too stale to advance + // leaves the record on its lease rather than free to reissue the + // iterations that ceiling covers. + let ceiling = self.lease.ceiling(); if ceiling > 0 && let Some(state) = self.states.front_mut() { state.fast_forward_sender_chain(ceiling)?; } + self.lease.waive(); Ok(()) } diff --git a/wacore/libsignal/src/protocol/state/session.rs b/wacore/libsignal/src/protocol/state/session.rs index 0ae1a9d2d..7e90a9e21 100644 --- a/wacore/libsignal/src/protocol/state/session.rs +++ b/wacore/libsignal/src/protocol/state/session.rs @@ -849,14 +849,28 @@ impl SessionRecord { /// persistence is synchronous and durable before the ciphertext reaches /// the wire. /// - /// The consumer applies this to every record it loads; nothing is inferred - /// from the stored representation. A snapshot written while the lease was - /// in force still carries a reservation that may already have been - /// published, so it is materialized once here before the lease goes away. - /// The guarantee being given up is stated on [`CounterLease`]. + /// By default the record leases outbound counters in batches, so the send + /// path needs a durable flush only when a batch runs out and any reload + /// fast-forwards past the whole lease. A consumer that persists before the + /// wire gets nothing from that and pays for it, since every export has to + /// burn the reserved range. + /// + /// **This gives up a real guarantee.** Message keys and IVs are derived + /// deterministically from the counter, so without the lease a crash between + /// the encrypt and the write can reissue a counter and with it the + /// (key, IV) pair. Only a consumer whose writes are durable before the wire + /// can make that trade. + /// + /// Apply this to every record loaded; nothing is inferred from the stored + /// representation, because the same representation can be persisted by a + /// consumer that wants the lease and by one that does not. A snapshot + /// written while the lease was in force still carries a reservation that + /// may already have been published, so it is materialized once here before + /// the lease goes away. pub fn waive_counter_lease(&mut self) { - let ceiling = self.lease.waive(); + let ceiling = self.lease.ceiling(); if ceiling == 0 { + self.lease.waive(); return; } // Every state the lease covered has to burn it, archived ones included: @@ -870,6 +884,7 @@ impl SessionRecord { state.fast_forward_sender_chain_or_drop(ceiling); *session = state.session; } + self.lease.waive(); } /// Lease a fresh batch of sender-chain counters after `spent_counter` was