feat(libsignal): let a consumer opt out of counter leasing

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

wacore/libsignal/src/protocol/counter_lease.rs

@@ -0,0 +1,168 @@
+//! Whether a record leases outbound counters ahead of durability.
+//!
+//! 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;
+
+/// 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);
+        }
+    }
+
+    /// Drop the lease.
+    ///
+    /// 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;
+    }
+}
+
+#[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_reports_nothing_left_to_materialize() {
+        let mut lease = CounterLease::from_persisted_ceiling(512);
+        assert_eq!(lease.ceiling(), 512);
+
+        lease.waive();
+        assert_eq!(lease, CounterLease::Waived);
+        // Converged: nothing left for a later load to advance past.
+        assert_eq!(lease.ceiling(), 0);
+    }
+
+    #[test]
+    fn rebasing_a_waived_lease_is_inert() {
+        let mut lease = CounterLease::Waived;
+        lease.rebase();
+        assert_eq!(lease, CounterLease::Waived);
+    }
+}

wacore/libsignal/src/protocol/group_cipher.rs

@@ -142,12 +142,10 @@ pub async fn group_encrypt<S: SenderKeyStore + ?Sized, R: Rng + CryptoRng>(
     // 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,165 @@ 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<SenderKeyName, crate::protocol::SenderKeyRecordComponents>,
+        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<Option<SenderKeyRecord>> {
+            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<u32>, 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::<Vec<u32>>());
+        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<u32> = (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 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.

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;