diff --git a/src/client/tests.rs b/src/client/tests.rs index 6935ad6da..83f2d80ad 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -1510,6 +1510,49 @@ async fn cleanup_connection_state_flushes_dirty_sender_key() { ); } +#[tokio::test] +async fn cleanup_connection_state_does_not_burn_a_clean_sender_key_lease() { + use wacore::libsignal::protocol::{KeyPair, SenderKeyRecord}; + use wacore::libsignal::store::sender_key_name::SenderKeyName; + let client = create_offline_sync_test_client().await; + let name = SenderKeyName::from_parts("group@g.us", "5550001001@s.whatsapp.net:1"); + let mut rng = rand::make_rng::(); + let signing_key = KeyPair::generate(&mut rng); + let mut record = SenderKeyRecord::new_empty(); + record + .add_sender_key_state( + 3, + 12345, + 0, + &[0x42; 32], + signing_key.public_key, + Some(signing_key.private_key), + ) + .expect("sender key state"); + record.reserve_iterations(0); + client.signal_cache.put_sender_key(&name, record).await; + + client.cleanup_connection_state().await; + + let device = client.persistence_manager.get_device_arc().await; + let guard = device.read().await; + let reloaded = client + .signal_cache + .get_sender_key(&name, &*guard.backend) + .await + .expect("sender key load") + .expect("sender key"); + assert_eq!( + reloaded + .sender_key_state() + .expect("sender key state") + .sender_chain_key() + .expect("sender chain") + .iteration(), + 0 + ); +} + /// When the flush itself fails, cleanup must NOT clear the cache, or it would /// drop the very state the flush was meant to persist. #[tokio::test] diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 06732c398..89a5b7d6c 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -477,7 +477,7 @@ impl Client { // drop with them instead of persisting rowless. if durable && !self.inbound_commit_batch.has_entries() { match self.flush_signal_cache().await { - Ok(()) => self.signal_cache.clear().await, + Ok(()) => self.signal_cache.clear_after_flush().await, // Committed/acked state the server never redelivers: keep // it resident so the next successful flush persists it. // Safe to carry across the reconnect — the teardown diff --git a/src/store/signal.rs b/src/store/signal.rs index 3e7c19527..4514d2ebc 100644 --- a/src/store/signal.rs +++ b/src/store/signal.rs @@ -1,7 +1,8 @@ use crate::store::Device; use async_lock::Mutex; use async_trait::async_trait; -use std::sync::Arc; +use rand::RngExt; +use std::sync::{Arc, OnceLock}; use wacore::libsignal::protocol::error::Result as SignalResult; use wacore::libsignal::protocol::{ Direction, IdentityChange, IdentityKey, IdentityKeyPair, IdentityKeyStore, PrivateKey, @@ -14,6 +15,18 @@ use waproto::whatsapp::{PreKeyRecordStructure, SignedPreKeyRecordStructure}; type StoreError = Box; +type DirectStoreIncarnation = [u8; 16]; + +// Synchronous direct writes make process restart the only unsafe reload boundary. +fn direct_store_incarnation() -> &'static DirectStoreIncarnation { + static INCARNATION: OnceLock = OnceLock::new(); + INCARNATION.get_or_init(|| { + let mut incarnation = [0; 16]; + rand::make_rng::().fill(&mut incarnation); + incarnation + }) +} + macro_rules! impl_store_wrapper { ($wrapper_ty:ty, $read_lock:ident, $write_lock:ident) => { #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] @@ -407,7 +420,8 @@ impl SessionStore for Device { let address_str = address.as_str(); match self.backend.get_session(address_str).await { Ok(Some(session_data)) => { - SessionRecord::deserialize(&session_data).map_err(|e| Box::new(e) as StoreError) + SessionRecord::deserialize_for_store(&session_data, direct_store_incarnation()) + .map_err(|e| Box::new(e) as StoreError) } Ok(None) => Ok(SessionRecord::new_fresh()), Err(e) => Err(Box::new(e) as StoreError), @@ -425,7 +439,8 @@ impl SessionStore for Device { record: &SessionRecord, ) -> Result<(), StoreError> { let address_str = address.as_str(); - let session_data = record.serialize().map_err(|e| Box::new(e) as StoreError)?; + let mut session_data = Vec::new(); + record.serialize_into_for_store(&mut session_data, direct_store_incarnation()); self.backend .put_session(address_str, &session_data) @@ -497,7 +512,7 @@ impl SenderKeyStore for Device { sender_key_name: &SenderKeyName, record: SenderKeyRecord, ) -> SignalResult<()> { - let serialized_record = record.serialize()?; + let serialized_record = record.serialize_for_store(direct_store_incarnation())?; self.backend .put_sender_key(sender_key_name.cache_key(), &serialized_record) .await @@ -515,8 +530,9 @@ impl SenderKeyStore for Device { .map_err(|e| SignalProtocolError::BackendError("load_sender_key", Box::new(e)))? { Some(data) => { - let record = SenderKeyRecord::deserialize(&data)?; - if record.serialize()?.is_empty() { + let record = + SenderKeyRecord::deserialize_for_store(&data, direct_store_incarnation())?; + if record.is_empty() { Ok(None) } else { Ok(Some(record)) @@ -531,6 +547,121 @@ impl SenderKeyStore for Device { mod tests { use super::*; + fn leased_session() -> SessionRecord { + use wacore::libsignal::protocol::{ChainKey, KeyPair, RootKey, SessionState}; + + let mut rng = rand::make_rng::(); + let local = IdentityKey::new(KeyPair::generate(&mut rng).public_key); + let remote = IdentityKey::new(KeyPair::generate(&mut rng).public_key); + let base_key = KeyPair::generate(&mut rng).public_key; + let mut state = SessionState::new(3, &local, &remote, &RootKey::new([0; 32]), &base_key); + state.set_sender_chain(&KeyPair::generate(&mut rng), &ChainKey::new([1; 32], 0)); + let mut record = SessionRecord::new(state); + record.reserve_sender_chain_counters(0); + record + } + + fn session_chain_index(record: &SessionRecord) -> u32 { + record + .session_state() + .expect("session") + .get_sender_chain_key() + .expect("sender chain") + .index() + } + + #[tokio::test] + async fn direct_session_store_preserves_clean_reload_and_recovery_ceiling() { + let backend = crate::test_utils::create_test_backend().await; + let device = Device::new(backend.clone()); + let address = ProtocolAddress::new("15550001001".to_string(), 1.into()); + + SessionStore::store_session(&device, &address, &leased_session()) + .await + .expect("store session"); + let clean = SessionStore::load_session(&device, &address) + .await + .expect("clean reload"); + assert_eq!(session_chain_index(&clean), 0); + + let replacement = Device::new(backend.clone()); + let same_process = SessionStore::load_session(&replacement, &address) + .await + .expect("same-process reload"); + assert_eq!(session_chain_index(&same_process), 0); + + let durable = backend + .get_session(address.as_str()) + .await + .expect("read durable session") + .expect("durable session"); + let recovered = SessionRecord::deserialize(&durable).expect("recovery reload"); + assert_eq!( + session_chain_index(&recovered), + wacore::libsignal::protocol::consts::SENDER_CHAIN_RESERVATION_BATCH + ); + } + + #[tokio::test] + async fn direct_sender_key_store_preserves_clean_reloads_and_recovery_ceiling() { + use wacore::libsignal::protocol::{ + create_sender_key_distribution_message, group_decrypt, group_encrypt, + process_sender_key_distribution_message, + }; + + let sender_backend = crate::test_utils::create_test_backend().await; + let mut sender = Device::new(sender_backend.clone()); + let mut receiver = Device::new(crate::test_utils::create_test_backend().await); + let name = SenderKeyName::from_parts("1234567890@g.us", "15550001000@s.whatsapp.net:0"); + let mut rng = rand::make_rng::(); + let distribution = create_sender_key_distribution_message(&name, &mut sender, &mut rng) + .await + .expect("sender setup"); + process_sender_key_distribution_message(&name, &distribution, &mut receiver) + .await + .expect("receiver setup"); + + let mut last = None; + for expected_iteration in 0..=32 { + let message = group_encrypt(&mut sender, &name, b"payload", &mut rng) + .await + .expect("group encrypt"); + assert_eq!(message.iteration(), expected_iteration); + last = Some(message); + } + + let plaintext = group_decrypt( + last.expect("last message").serialized(), + &mut receiver, + &name, + ) + .await + .expect("receiver decrypts after missed messages"); + assert_eq!(plaintext, b"payload"); + + let mut replacement = Device::new(sender_backend.clone()); + let same_process = group_encrypt(&mut replacement, &name, b"same-process", &mut rng) + .await + .expect("encrypt after same-process replacement"); + assert_eq!(same_process.iteration(), 33); + + let durable = sender_backend + .get_sender_key(name.cache_key()) + .await + .expect("read durable sender key") + .expect("durable sender key"); + let recovered = SenderKeyRecord::deserialize(&durable).expect("recovery reload"); + assert_eq!( + recovered + .sender_key_state() + .expect("sender-key state") + .sender_chain_key() + .expect("sender chain") + .iteration(), + wacore::libsignal::protocol::consts::SENDER_CHAIN_RESERVATION_BATCH + ); + } + // A rotated-out signed pre-key (id != current field) must load from the // backend table so delayed prekey messages naming the old id still decrypt. #[tokio::test] diff --git a/wacore/libsignal/src/protocol/group_cipher.rs b/wacore/libsignal/src/protocol/group_cipher.rs index c51dd7246..6579f3f69 100644 --- a/wacore/libsignal/src/protocol/group_cipher.rs +++ b/wacore/libsignal/src/protocol/group_cipher.rs @@ -25,6 +25,12 @@ struct CryptoBuffer { impl CryptoBuffer { const INITIAL_CAPACITY: usize = 1024; + /// Reuse the buffer across sends, but do not let a one-off large message pin + /// an oversized allocation on this thread for the rest of the process. Small + /// messages (the common case) fit under this and reuse the buffer with no + /// reallocation; anything larger is released back after its result is copied + /// out. Sized well above a typical message so normal traffic never churns. + const MAX_RETAINED_CAPACITY: usize = 16 * 1024; fn new() -> Self { Self { @@ -43,6 +49,18 @@ impl CryptoBuffer { fn take_buffer(&mut self) -> Vec { std::mem::replace(&mut self.buffer, Vec::with_capacity(Self::INITIAL_CAPACITY)) } + + /// Copy the written bytes into a right-sized box, keeping the buffer for the + /// next (small) message instead of handing away its capacity. A one-off + /// large message's capacity is released here so it is not retained on this + /// thread-local for the process lifetime. + fn copy_out(&mut self) -> Box<[u8]> { + let out: Box<[u8]> = self.buffer.as_slice().into(); + if self.buffer.capacity() > Self::MAX_RETAINED_CAPACITY { + self.buffer = Vec::with_capacity(Self::INITIAL_CAPACITY); + } + out + } } thread_local! { @@ -53,8 +71,8 @@ thread_local! { /// Caller must hold `SenderKeyStore::sender_key_lock` for `sender_key_name` /// across this call (and any paired SKDM creation) so the load/advance/store /// of the chain is atomic against concurrent encrypts. -pub async fn group_encrypt( - sender_key_store: &mut dyn SenderKeyStore, +pub async fn group_encrypt( + sender_key_store: &mut S, sender_key_name: &SenderKeyName, plaintext: &[u8], csprng: &mut R, @@ -87,13 +105,21 @@ pub async fn group_encrypt( let ciphertext = ENCRYPTION_BUFFER.with(|buffer| { let mut buf_wrapper = buffer.borrow_mut(); - let buf = buf_wrapper.get_buffer(); - aes_256_cbc_encrypt_into(plaintext, message_keys.cipher_key(), message_keys.iv(), buf) - .map_err(|_| { - log::error!("outgoing sender key state corrupt for distribution"); - SignalProtocolError::InvalidSenderKeySession - })?; - Ok::, SignalProtocolError>(buf_wrapper.take_buffer()) + { + let buf = buf_wrapper.get_buffer(); + aes_256_cbc_encrypt_into(plaintext, message_keys.cipher_key(), message_keys.iv(), buf) + .map_err(|_| { + log::error!("outgoing sender key state corrupt for distribution"); + SignalProtocolError::InvalidSenderKeySession + })?; + } + // Copy out a right-sized ciphertext and keep the reusable buffer, rather + // than handing away its (1 KiB) capacity and re-allocating one per send. + // A group skmsg is small, so `take_buffer` here otherwise cost a fresh + // 1 KiB per send once the production path started sharing this primitive. + // `copy_out` releases the capacity of a one-off large message so it is + // not pinned on this thread-local afterward. + Ok::, SignalProtocolError>(buf_wrapper.copy_out()) })?; let signing_key = sender_key_state @@ -104,17 +130,24 @@ pub async fn group_encrypt( message_version, sender_key_state.chain_id(), message_keys.iteration(), - ciphertext.into_boxed_slice(), + ciphertext, csprng, &signing_key, )?; 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(); + // Outbound advance: this iteration's (key, IV) must never be re-derivable. + // Iterations below the durable reservation are already covered, so their + // 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); + } sender_key_store .store_sender_key(sender_key_name, record) @@ -359,6 +392,28 @@ mod tests { use async_trait::async_trait; use std::collections::HashMap; + #[test] + fn crypto_buffer_reuses_small_and_releases_oversized() { + let mut b = CryptoBuffer::new(); + + // A small message keeps its (reusable) capacity: no release, no churn. + b.get_buffer().extend_from_slice(&[0u8; 512]); + let _ = b.copy_out(); + assert!(b.buffer.capacity() <= CryptoBuffer::MAX_RETAINED_CAPACITY); + assert!(b.buffer.capacity() >= CryptoBuffer::INITIAL_CAPACITY); + + // A one-off large message grows the buffer, but copy_out releases the + // excess so it is not pinned on the thread-local afterward. + b.get_buffer() + .resize(CryptoBuffer::MAX_RETAINED_CAPACITY * 4, 0); + assert!(b.buffer.capacity() > CryptoBuffer::MAX_RETAINED_CAPACITY); + let _ = b.copy_out(); + assert!( + b.buffer.capacity() <= CryptoBuffer::MAX_RETAINED_CAPACITY, + "oversized capacity must be released after copy_out" + ); + } + struct InMemorySenderKeyStore { keys: HashMap, } @@ -425,4 +480,136 @@ mod tests { other => panic!("expected NoSenderKeyState, got {other:?}"), } } + + use crate::protocol::consts::SENDER_CHAIN_RESERVATION_BATCH; + use futures::executor::block_on; + + /// THE crash-safety invariant: a reload mid-lease must never re-derive an + /// iteration that a lost send may already have put on the wire, and a peer + /// must still decrypt across the resulting gap. Bob sends iteration 0 + /// (reserving a batch), his record is snapshotted, a further send is lost in + /// the crash, and the reload fast-forwards past the whole reservation. + #[test] + fn crash_mid_lease_skips_spent_iterations_and_peer_decrypts() { + 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(), + }; + let skdm = block_on(create_sender_key_distribution_message( + &name, &mut bob, &mut rng, + )) + .expect("bob creates his sender key distribution message"); + let mut alice = InMemorySenderKeyStore { + keys: HashMap::new(), + }; + block_on(process_sender_key_distribution_message( + &name, &skdm, &mut alice, + )) + .expect("alice processes bob's distribution message"); + + // Bob's first send: iteration 0, reserves the first batch. + let m0 = block_on(group_encrypt(&mut bob, &name, b"m0", &mut rng)) + .expect("bob encrypts m0 at iteration 0"); + assert_eq!(m0.iteration(), 0); + assert_eq!( + block_on(group_decrypt(m0.serialized(), &mut alice, &name)) + .expect("alice decrypts m0 at iteration 0"), + b"m0" + ); + + // Snapshot Bob's durable state (chain at iteration 1, reservation persisted). + let snapshot = bob + .keys + .get(&name) + .expect("bob's record is stored after his first send") + .serialize() + .expect("bob's record serializes for the durable snapshot"); + + // A send whose advance never becomes durable (crash before flush). + let lost = block_on(group_encrypt(&mut bob, &name, b"lost", &mut rng)) + .expect("bob encrypts the send that the crash loses"); + assert_eq!(lost.iteration(), 1); + + // Reload from the snapshot: the current chain fast-forwards past the lease. + let reloaded = SenderKeyRecord::deserialize(&snapshot) + .expect("the durable snapshot deserializes on reload"); + bob.keys.insert(name.clone(), reloaded); + + // The next send must NOT reuse iterations 1..batch-1 (the lost one included). + let after = block_on(group_encrypt(&mut bob, &name, b"after", &mut rng)) + .expect("bob encrypts the first send after the reload"); + assert!( + after.iteration() >= SENDER_CHAIN_RESERVATION_BATCH, + "reload reused a possibly-spent iteration: {} < {}", + after.iteration(), + SENDER_CHAIN_RESERVATION_BATCH + ); + + // Alice, who last saw iteration 0, still decrypts across the gap + // (a forward jump well under MAX_FORWARD_JUMPS). + assert_eq!( + block_on(group_decrypt(after.serialized(), &mut alice, &name)) + .expect("alice decrypts across the fast-forwarded iteration gap"), + b"after" + ); + } + + /// 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. + struct GateCountingStore { + keys: HashMap, + gated: usize, + } + #[async_trait] + impl SenderKeyStore for GateCountingStore { + async fn store_sender_key( + &mut self, + name: &SenderKeyName, + mut record: SenderKeyRecord, + ) -> Result<()> { + if record.is_wire_gated() { + self.gated += 1; + record.clear_wire_gated(); + } + self.keys.insert(name.clone(), record); + Ok(()) + } + async fn load_sender_key(&self, name: &SenderKeyName) -> Result> { + Ok(self.keys.get(name).cloned()) + } + } + + /// The lease must amortize the synchronous flush to one per batch: over + /// `2*batch + 2` sends, exactly the sends at iterations 0, batch and 2*batch + /// gate; the rest ride the coalesced write-behind. + #[test] + fn lease_amortizes_the_wire_gate() { + let mut rng = rand::rng(); + let name = SenderKeyName::new("group@g.us".to_string(), "bob.0".to_string()); + let mut bob = GateCountingStore { + keys: HashMap::new(), + gated: 0, + }; + block_on(create_sender_key_distribution_message( + &name, &mut bob, &mut rng, + )) + .expect("bob creates his sender key distribution message"); + + let sends = 2 * SENDER_CHAIN_RESERVATION_BATCH + 2; + for i in 0..sends { + block_on(group_encrypt(&mut bob, &name, b"x", &mut rng)).unwrap_or_else(|e| { + panic!("bob's send {i} of {sends} under the lease failed: {e}") + }); + } + assert_eq!( + bob.gated, + 3, + "expected one gate per batch (iterations 0, {batch}, {two_batch}), got {got}", + batch = SENDER_CHAIN_RESERVATION_BATCH, + two_batch = 2 * SENDER_CHAIN_RESERVATION_BATCH, + got = bob.gated + ); + } } diff --git a/wacore/libsignal/src/protocol/local_field.rs b/wacore/libsignal/src/protocol/local_field.rs new file mode 100644 index 000000000..552ac8172 --- /dev/null +++ b/wacore/libsignal/src/protocol/local_field.rs @@ -0,0 +1,67 @@ +//! Local record metadata fails closed so a corrupt lease cannot re-enable spent +//! counters. + +use buffa::encoding::{Tag, WireType, decode_varint, encode_varint, skip_field}; + +const STORE_INCARNATION_FIELD: u32 = 101; +const STORE_INCARNATION_LEN: usize = 16; +pub(crate) const STORE_INCARNATION_ENCODED_LEN: usize = 19; + +pub(crate) struct LocalRecordFields { + pub(crate) reservation: u32, + pub(crate) incarnation: Option<[u8; STORE_INCARNATION_LEN]>, +} + +/// Malformed or duplicate metadata cannot be trusted to preserve a lease. +pub(crate) fn decode_local_record_fields( + mut bytes: &[u8], + reservation_field: u32, + on_err: impl Fn() -> E, +) -> Result { + let mut reservation = None; + let mut incarnation = None; + let mut incarnation_valid = true; + while !bytes.is_empty() { + let tag = Tag::decode(&mut bytes).map_err(|_| on_err())?; + if tag.field_number() == reservation_field { + if tag.wire_type() != WireType::Varint || reservation.is_some() { + return Err(on_err()); + } + let raw = decode_varint(&mut bytes).map_err(|_| on_err())?; + reservation = Some(u32::try_from(raw).map_err(|_| on_err())?); + } else if tag.field_number() == STORE_INCARNATION_FIELD { + if tag.wire_type() == WireType::LengthDelimited { + let len = usize::try_from(decode_varint(&mut bytes).map_err(|_| on_err())?) + .map_err(|_| on_err())?; + if bytes.len() < len { + return Err(on_err()); + } + if len == STORE_INCARNATION_LEN && incarnation.is_none() && incarnation_valid { + let mut value = [0; STORE_INCARNATION_LEN]; + value.copy_from_slice(&bytes[..len]); + incarnation = Some(value); + } else { + incarnation = None; + incarnation_valid = false; + } + bytes = &bytes[len..]; + } else { + skip_field(tag, &mut bytes).map_err(|_| on_err())?; + incarnation = None; + incarnation_valid = false; + } + } else { + skip_field(tag, &mut bytes).map_err(|_| on_err())?; + } + } + Ok(LocalRecordFields { + reservation: reservation.unwrap_or(0), + incarnation: if incarnation_valid { incarnation } else { None }, + }) +} + +pub(crate) fn encode_store_incarnation(bytes: &mut Vec, incarnation: &[u8; 16]) { + Tag::new(STORE_INCARNATION_FIELD, WireType::LengthDelimited).encode(bytes); + encode_varint(STORE_INCARNATION_LEN as u64, bytes); + bytes.extend_from_slice(incarnation); +} diff --git a/wacore/libsignal/src/protocol/mod.rs b/wacore/libsignal/src/protocol/mod.rs index 18a2e6526..3e32b4857 100644 --- a/wacore/libsignal/src/protocol/mod.rs +++ b/wacore/libsignal/src/protocol/mod.rs @@ -22,6 +22,7 @@ mod crypto; pub mod error; mod group_cipher; mod identity_key; +mod local_field; #[allow(clippy::module_inception)] mod protocol; mod ratchet; diff --git a/wacore/libsignal/src/protocol/sender_keys.rs b/wacore/libsignal/src/protocol/sender_keys.rs index c324c5a7b..b71c27aa7 100644 --- a/wacore/libsignal/src/protocol/sender_keys.rs +++ b/wacore/libsignal/src/protocol/sender_keys.rs @@ -203,6 +203,14 @@ pub struct SenderKeyState { /// `state.sender_message_keys` is kept empty in memory; this is the source of /// truth, reassembled into the protobuf only at `as_protobuf` (serialization). message_keys: std::sync::Arc>, + /// The current sender chain key, held as a `Copy` value instead of in the + /// protobuf. The chain seed is a `Bytes` in the generated structure, so + /// keeping it there made every record clone (and the copy-on-write on every + /// encrypt/decrypt advance) promote that `Bytes` to a shared allocation. + /// Same source-of-truth-outside-the-protobuf trick as `message_keys`: + /// `state.sender_chain_key` stays empty in memory, reassembled only at + /// `as_protobuf`. `None` only for a structurally invalid state. + sender_chain: Option, /// Parsed signing key with its XEdDSA cache pre-derived, memoized so the /// per-send signature skips a basepoint multiplication (~18% of a warm /// group send when re-derived from bytes every message). Clones carry the @@ -247,11 +255,12 @@ impl SenderKeyState { let chain_key_arr: [u8; 32] = chain_key .try_into() .map_err(|_| SignalProtocolError::InvalidProtobufEncoding)?; + let sender_chain = Some(SenderChainKey::new(iteration, chain_key_arr)); let state = SenderKeyStateStructure { sender_key_id: Some(chain_id), - sender_chain_key: MessageField::some( - SenderChainKey::new(iteration, chain_key_arr).as_protobuf(), - ), + // Source of truth is `sender_chain`; the protobuf field stays empty + // in memory and is reassembled at as_protobuf. + sender_chain_key: MessageField::none(), sender_signing_key: MessageField::some(sender_key_state_structure::SenderSigningKey { public: Some(Bytes::copy_from_slice(&signature_key.serialize())), private: signature_private_key @@ -278,6 +287,7 @@ impl SenderKeyState { Ok(Self { state, message_keys: std::sync::Arc::new(Vec::new()), + sender_chain, signing_key_memo, verifying_key_memo, }) @@ -292,9 +302,16 @@ impl SenderKeyState { .map(StoredMessageKey::from_protobuf) .collect::>(), ); + // Likewise move the chain key out into the Copy field; the seed was + // validated at deserialize before this runs. + let sender_chain = state.sender_chain_key.take().and_then(|sc| { + let seed: [u8; 32] = sc.seed.as_deref()?.try_into().ok()?; + Some(SenderChainKey::new(sc.iteration.unwrap_or_default(), seed)) + }); Self { state, message_keys, + sender_chain, signing_key_memo: std::sync::OnceLock::new(), verifying_key_memo: std::sync::OnceLock::new(), } @@ -309,21 +326,39 @@ impl SenderKeyState { } pub fn sender_chain_key(&self) -> Option { - let sender_chain = self.state.sender_chain_key.as_option()?; - let seed: [u8; 32] = sender_chain - .seed - .as_deref() - .unwrap_or_default() - .try_into() - .ok()?; - Some(SenderChainKey::new( - sender_chain.iteration.unwrap_or_default(), - seed, - )) + self.sender_chain } pub fn set_sender_chain_key(&mut self, chain_key: SenderChainKey) { - self.state.sender_chain_key = MessageField::some(chain_key.as_protobuf()); + self.sender_chain = Some(chain_key); + } + + /// Advance the sender chain up to a reload's reserved iteration ceiling so no + /// possibly-spent iteration below it stays derivable. Bounded by + /// `MAX_RESERVATION_FAST_FORWARD`; a target past that is a corrupt + /// reservation and errors rather than looping the KDF unboundedly. Mirrors + /// `SessionRecord::fast_forward_sender_chain`. + pub(crate) fn fast_forward_sender_chain( + &mut self, + target: u32, + ) -> Result<(), SignalProtocolError> { + let Some(mut chain_key) = self.sender_chain_key() else { + return Ok(()); + }; + if target.saturating_sub(chain_key.iteration()) > consts::MAX_RESERVATION_FAST_FORWARD { + return Err(SignalProtocolError::InvalidState( + "fast_forward_sender_chain", + "reserved sender-key iteration implausibly far ahead".into(), + )); + } + if chain_key.iteration() >= target { + return Ok(()); + } + while chain_key.iteration() < target { + chain_key = chain_key.next()?; + } + self.set_sender_chain_key(chain_key); + Ok(()) } pub fn signing_key_public(&self) -> Result { @@ -386,8 +421,9 @@ impl SenderKeyState { pub(crate) fn as_protobuf(&self) -> SenderKeyStateStructure { debug_assert!( - self.state.sender_message_keys.is_empty(), - "backlog must live only in `message_keys`; the protobuf copy stays empty" + self.state.sender_message_keys.is_empty() + && self.state.sender_chain_key.as_option().is_none(), + "backlog and chain key must live only in their Copy/Arc fields; the protobuf copies stay empty" ); let mut state = self.state.clone(); state.sender_message_keys = self @@ -395,6 +431,10 @@ impl SenderKeyState { .iter() .map(StoredMessageKey::as_protobuf) .collect(); + state.sender_chain_key = self + .sender_chain + .as_ref() + .map_or_else(MessageField::none, |c| MessageField::some(c.as_protobuf())); state } @@ -434,24 +474,77 @@ pub struct SenderKeyRecord { /// 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, } +/// Local-only field appended to the serialized record for `reserved_iteration`. +/// The vendored `SenderKeyRecordStructure` proto is untouched; the generated +/// decoder skips this unknown top-level field and `deserialize` scans it out. +/// Matches the field-number scheme `SessionRecord` uses for its DM counterpart. +const RESERVED_ITERATION_FIELD: u32 = 100; + 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; + self.reserved_iteration = 0; } pub fn new_empty() -> Self { Self { states: VecDeque::with_capacity(consts::MAX_SENDER_KEY_STATES), wire_gated: false, + reserved_iteration: 0, } } + pub fn is_empty(&self) -> bool { + self.states.is_empty() + } + + /// 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 + } + + /// Lease a fresh batch of iterations after `spent_iteration` reached the + /// current ceiling. Marks the record wire-gated: the caller's ciphertext + /// 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; + } + pub fn deserialize(buf: &[u8]) -> Result { + Self::deserialize_inner(buf, None) + } + + /// A matching live cache proves this snapshot was not recovered after a crash. + #[doc(hidden)] + pub fn deserialize_for_store( + buf: &[u8], + incarnation: &[u8; 16], + ) -> Result { + Self::deserialize_inner(buf, Some(incarnation)) + } + + fn deserialize_inner( + buf: &[u8], + incarnation: Option<&[u8; 16]>, + ) -> Result { let skr = SenderKeyRecordStructure::decode_from_slice(buf) .map_err(|_| SignalProtocolError::InvalidProtobufEncoding)?; @@ -466,9 +559,27 @@ impl SenderKeyRecord { } states.push_back(SenderKeyState::from_protobuf(state)); } + + let local_fields = + super::local_field::decode_local_record_fields(buf, RESERVED_ITERATION_FIELD, || { + SignalProtocolError::InvalidProtobufEncoding + })?; + let reserved_iteration = local_fields.reservation; + let trusted_reload = + incarnation.is_some_and(|current| local_fields.incarnation == Some(*current)); + + // Only an untrusted snapshot may have spent its still-reserved range. + if !trusted_reload + && reserved_iteration > 0 + && let Some(state) = states.front_mut() + { + state.fast_forward_sender_chain(reserved_iteration)?; + } + Ok(Self { states, wire_gated: false, + reserved_iteration, }) } @@ -552,6 +663,16 @@ impl SenderKeyRecord { } self.states.push_front(state); + // Reset the reservation unconditionally. It is a record-level ceiling for + // whatever chain is current, and this call may have replaced or reordered + // the current chain. Resetting is always safe (the next send re-reserves + // and re-gates); keeping a stale ceiling across a chain change is not, + // since a lower-iteration chain would treat already-covered iterations as + // durable and re-derive a spent (key, IV). In practice this is a no-op: + // 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; Ok(()) } @@ -587,7 +708,46 @@ impl SenderKeyRecord { } pub fn serialize(&self) -> Result, SignalProtocolError> { - Ok(self.as_protobuf().encode_to_vec()) + self.serialize_inner(None) + } + + /// The incarnation prevents clean reloads from looking like crashes. + #[doc(hidden)] + pub fn serialize_for_store( + &self, + incarnation: &[u8; 16], + ) -> Result, SignalProtocolError> { + self.serialize_inner(Some(incarnation)) + } + + fn serialize_inner( + &self, + incarnation: Option<&[u8; 16]>, + ) -> Result, SignalProtocolError> { + use buffa::encoding::{Tag, WireType, encode_varint, varint_len}; + + let mut buf = self.as_protobuf().encode_to_vec(); + let incarnation = incarnation.filter(|_| self.reserved_iteration > 0); + let reservation_len = if self.reserved_iteration > 0 { + 2 + varint_len(self.reserved_iteration as u64) + } else { + 0 + }; + let incarnation_len = incarnation + .map(|_| super::local_field::STORE_INCARNATION_ENCODED_LEN) + .unwrap_or(0); + buf.reserve(reservation_len + incarnation_len); + // 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 { + Tag::new(RESERVED_ITERATION_FIELD, WireType::Varint).encode(&mut buf); + encode_varint(self.reserved_iteration as u64, &mut buf); + } + if let Some(incarnation) = incarnation { + super::local_field::encode_store_incarnation(&mut buf, incarnation); + } + Ok(buf) } /// Estimated in-memory footprint proxy: encoded size of each state's @@ -601,6 +761,8 @@ impl SenderKeyRecord { .map(|s| { s.state.compute_size(&mut cache) as usize + s.message_keys.len() * std::mem::size_of::() + + s.sender_chain + .map_or(0, |_| std::mem::size_of::()) }) .sum() } @@ -1020,6 +1182,171 @@ mod tests { assert_eq!(state.chain_id(), 12345); } + fn record_with_state(chain_id: u32, seed: u8) -> SenderKeyRecord { + let mut rng = rand::make_rng::(); + let keypair = KeyPair::generate(&mut rng); + let mut record = SenderKeyRecord::new_empty(); + record + .add_sender_key_state( + 3, + chain_id, + 0, + &[seed; 32], + keypair.public_key, + Some(keypair.private_key), + ) + .expect("state should be valid"); + record + } + + fn current_iteration(record: &SenderKeyRecord) -> u32 { + record + .sender_key_state() + .expect("test") + .sender_chain_key() + .expect("test") + .iteration() + } + + /// A reservation survives a serialize/deserialize round-trip, and the reload + /// fast-forwards the current chain past the reserved ceiling so no + /// possibly-spent iteration below it stays derivable. + #[test] + fn reservation_survives_roundtrip_and_reload_fast_forwards() { + let mut record = record_with_state(12345, 0x42); + record.reserve_iterations(0); + assert_eq!( + record.reserved_iteration(), + consts::SENDER_CHAIN_RESERVATION_BATCH + ); + assert_eq!( + current_iteration(&record), + 0, + "reserving does not advance the chain" + ); + + let reloaded = + SenderKeyRecord::deserialize(&record.serialize().expect("test")).expect("test"); + assert_eq!( + reloaded.reserved_iteration(), + consts::SENDER_CHAIN_RESERVATION_BATCH + ); + assert_eq!( + current_iteration(&reloaded), + consts::SENDER_CHAIN_RESERVATION_BATCH, + "reload must burn the reserved iterations" + ); + } + + #[test] + fn cache_incarnation_separates_clean_reload_from_recovery() { + let mut record = record_with_state(12345, 0x42); + record.reserve_iterations(0); + let incarnation = [0xA1; 16]; + let replacement = [0xB2; 16]; + let bytes = record.serialize_for_store(&incarnation).expect("test"); + + let clean = SenderKeyRecord::deserialize_for_store(&bytes, &incarnation).expect("test"); + assert_eq!(current_iteration(&clean), 0); + + let recovered = SenderKeyRecord::deserialize_for_store(&bytes, &replacement).expect("test"); + assert_eq!( + current_iteration(&recovered), + consts::SENDER_CHAIN_RESERVATION_BATCH + ); + + let conservative = SenderKeyRecord::deserialize(&bytes).expect("test"); + assert_eq!( + current_iteration(&conservative), + consts::SENDER_CHAIN_RESERVATION_BATCH + ); + + let legacy = record.serialize().expect("test"); + let migrated = SenderKeyRecord::deserialize_for_store(&legacy, &incarnation).expect("test"); + assert_eq!( + current_iteration(&migrated), + consts::SENDER_CHAIN_RESERVATION_BATCH + ); + + let mut duplicated = bytes; + crate::protocol::local_field::encode_store_incarnation(&mut duplicated, &incarnation); + let duplicate_recovery = + SenderKeyRecord::deserialize_for_store(&duplicated, &incarnation).expect("test"); + assert_eq!( + current_iteration(&duplicate_recovery), + consts::SENDER_CHAIN_RESERVATION_BATCH + ); + } + + /// Rotating the current sender-key state (a fresh chain) resets the lease so + /// the next send re-reserves against the new chain instead of treating stale + /// iterations as durable. + #[test] + fn rotation_resets_the_lease() { + let mut record = record_with_state(111, 0x42); + record.reserve_iterations(0); + assert_eq!( + record.reserved_iteration(), + consts::SENDER_CHAIN_RESERVATION_BATCH + ); + + let mut rng = rand::make_rng::(); + let kp = KeyPair::generate(&mut rng); + record + .add_sender_key_state( + 3, + 222, + 0, + &[0x43u8; 32], + kp.public_key, + Some(kp.private_key), + ) + .expect("test"); + assert_eq!( + record.reserved_iteration(), + 0, + "rotation must reset the lease" + ); + } + + /// A record written without the local-only field (an older lib, or an + /// unreserved record) is byte-identical to the plain generated encoding and + /// loads with reservation 0; the first send re-reserves. + #[test] + fn legacy_record_without_field_loads_with_zero() { + let record = record_with_state(12345, 0x42); + let bytes = record.serialize().expect("test"); + assert_eq!( + bytes, + record.as_protobuf().encode_to_vec(), + "an unreserved record appends no field" + ); + let loaded = SenderKeyRecord::deserialize(&bytes).expect("test"); + assert_eq!(loaded.reserved_iteration(), 0); + assert_eq!( + current_iteration(&loaded), + 0, + "no reservation, no fast-forward" + ); + } + + /// A reservation implausibly far past the current chain is a corrupt record + /// and must fail closed at load rather than looping the KDF unboundedly. + #[test] + fn corrupt_reservation_is_rejected() { + use buffa::encoding::{Tag, WireType, encode_varint}; + + let record = record_with_state(12345, 0x42); + let mut bytes = record.serialize().expect("test"); + let corrupt = consts::MAX_RESERVATION_FAST_FORWARD + 1000; + Tag::new(RESERVED_ITERATION_FIELD, WireType::Varint).encode(&mut bytes); + encode_varint(corrupt as u64, &mut bytes); + assert!( + SenderKeyRecord::deserialize(&bytes).is_err(), + "an implausibly-far reservation must fail closed" + ); + } + /// Test SenderKeyRecord state limit #[test] fn test_sender_key_record_state_limit() { diff --git a/wacore/libsignal/src/protocol/state/session.rs b/wacore/libsignal/src/protocol/state/session.rs index e8b46762a..f5ef3b8d1 100644 --- a/wacore/libsignal/src/protocol/state/session.rs +++ b/wacore/libsignal/src/protocol/state/session.rs @@ -736,6 +736,22 @@ impl SessionRecord { } pub fn deserialize(bytes: &[u8]) -> Result { + Self::deserialize_inner(bytes, None) + } + + /// A matching live cache proves this snapshot was not recovered after a crash. + #[doc(hidden)] + pub fn deserialize_for_store( + bytes: &[u8], + incarnation: &[u8; 16], + ) -> Result { + Self::deserialize_inner(bytes, Some(incarnation)) + } + + fn deserialize_inner( + bytes: &[u8], + incarnation: Option<&[u8; 16]>, + ) -> Result { use waproto::whatsapp::RecordStructureView; // Decode to a zero-copy view first, then only convert sessions we @@ -753,6 +769,12 @@ impl SessionRecord { .collect::>() .map_err(|_| InvalidSessionError("failed to decode archived session protobuf"))?; + let local_fields = crate::protocol::local_field::decode_local_record_fields( + bytes, + RESERVED_SENDER_CHAIN_INDEX_FIELD, + || InvalidSessionError("invalid local session metadata"), + )?; + let mut record = Self { current_session: view .current_session @@ -762,13 +784,16 @@ 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)?, + reserved_sender_chain_index: local_fields.reservation, 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 trusted_reload = + incarnation.is_some_and(|current| local_fields.incarnation == Some(*current)); + + // An untrusted snapshot may predate sends covered by its lease. + if !trusted_reload + && record.reserved_sender_chain_index > 0 && let Some(state) = record.current_session.as_mut() { state.fast_forward_sender_chain(record.reserved_sender_chain_index)?; @@ -777,42 +802,6 @@ impl SessionRecord { 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. - /// - /// 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 = 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 { - 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 = 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.unwrap_or(0)) - } - /// If there's a session with a matching version and `alice_base_key`, ensures that it is the /// current session, promoting if necessary. /// @@ -1004,6 +993,16 @@ impl SessionRecord { /// Encode into a caller-supplied buffer (allows reuse across flushes). pub fn serialize_into(&self, buf: &mut Vec) { + self.serialize_into_inner(buf, None); + } + + /// The incarnation prevents clean reloads from looking like crashes. + #[doc(hidden)] + pub fn serialize_into_for_store(&self, buf: &mut Vec, incarnation: &[u8; 16]) { + self.serialize_into_inner(buf, Some(incarnation)); + } + + fn serialize_into_inner(&self, buf: &mut Vec, incarnation: Option<&[u8; 16]>) { use buffa::encoding::{Tag, WireType, encode_varint, varint_len}; fn write_len_delimited( @@ -1039,14 +1038,18 @@ impl SessionRecord { .sum(); let reserved = self.reserved_sender_chain_index; + let incarnation = incarnation.filter(|_| reserved > 0); let reserved_len = if reserved > 0 { 2 + varint_len(reserved as u64) } else { 0 }; + let incarnation_len = incarnation + .map(|_| crate::protocol::local_field::STORE_INCARNATION_ENCODED_LEN) + .unwrap_or(0); buf.clear(); - buf.reserve(current_len + previous_len + reserved_len); + buf.reserve(current_len + previous_len + reserved_len + incarnation_len); if let Some(state) = &self.current_session && let Some(msg_len) = current_msg_len @@ -1060,6 +1063,9 @@ impl SessionRecord { Tag::new(RESERVED_SENDER_CHAIN_INDEX_FIELD, WireType::Varint).encode(buf); encode_varint(reserved as u64, buf); } + if let Some(incarnation) = incarnation { + crate::protocol::local_field::encode_store_incarnation(buf, incarnation); + } } /// Estimated in-memory footprint proxy: the protobuf-encoded size of the @@ -1313,6 +1319,63 @@ mod tests { ); } + #[test] + fn cache_incarnation_separates_clean_reload_from_recovery() { + 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 incarnation = [0xA1; 16]; + let replacement = [0xB2; 16]; + let mut bytes = Vec::new(); + record.serialize_into_for_store(&mut bytes, &incarnation); + + let clean = SessionRecord::deserialize_for_store(&bytes, &incarnation).unwrap(); + assert_eq!( + clean + .session_state() + .unwrap() + .get_sender_chain_key() + .unwrap() + .index(), + 0 + ); + + let recovered = SessionRecord::deserialize_for_store(&bytes, &replacement).unwrap(); + assert_eq!( + recovered + .session_state() + .unwrap() + .get_sender_chain_key() + .unwrap() + .index(), + consts::SENDER_CHAIN_RESERVATION_BATCH + ); + + let conservative = SessionRecord::deserialize(&bytes).unwrap(); + assert_eq!( + conservative + .session_state() + .unwrap() + .get_sender_chain_key() + .unwrap() + .index(), + consts::SENDER_CHAIN_RESERVATION_BATCH + ); + + let legacy = record.serialize().unwrap(); + let migrated = SessionRecord::deserialize_for_store(&legacy, &incarnation).unwrap(); + assert_eq!( + migrated + .session_state() + .unwrap() + .get_sender_chain_key() + .unwrap() + .index(), + consts::SENDER_CHAIN_RESERVATION_BATCH + ); + } + /// 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 diff --git a/wacore/src/send.rs b/wacore/src/send.rs index a88ec76f6..c27708186 100644 --- a/wacore/src/send.rs +++ b/wacore/src/send.rs @@ -1,8 +1,7 @@ use crate::client::context::{GroupInfo, SendContextResolver}; use crate::libsignal::protocol::{ - CiphertextMessage, IdentityChange, ProtocolAddress, SENDERKEY_MESSAGE_CURRENT_VERSION, - SenderKeyMessage, SenderKeyStore, SignalProtocolError, UsePQRatchet, message_encrypt, - process_prekey_bundle, + CiphertextMessage, IdentityChange, ProtocolAddress, SenderKeyMessage, SenderKeyStore, + UsePQRatchet, message_encrypt, process_prekey_bundle, }; use crate::libsignal::store::sender_key_name::SenderKeyName; use crate::messages::MessageUtils; @@ -23,7 +22,6 @@ use std::future::Future; use wacore_binary::Node; use wacore_binary::builder::NodeBuilder; use wacore_binary::{CompactString, Jid, JidExt as _}; -use wacore_libsignal::crypto::aes_256_cbc_encrypt_into; use waproto::whatsapp as wa; /// Wire-format constants (MsgCreateDeviceStanza.js). diff --git a/wacore/src/send/encrypt.rs b/wacore/src/send/encrypt.rs index da40c7e33..cb0d6d3d3 100644 --- a/wacore/src/send/encrypt.rs +++ b/wacore/src/send/encrypt.rs @@ -1,6 +1,7 @@ //! Per-device Signal encryption fanout and the bounded spawn helper. use super::*; +use anyhow::Context; /// Caller must hold `SenderKeyStore::sender_key_lock` for `sender_key_name` /// across the surrounding SKDM creation + this encrypt, so a concurrent send @@ -19,62 +20,14 @@ where S: SenderKeyStore + ?Sized, R: Rng + CryptoRng, { - log::debug!( - "Attempting to load sender key for group {} sender {}", - sender_key_name.group_id(), - sender_key_name.sender_id() - ); - - let mut record = sender_key_store - .load_sender_key(sender_key_name) - .await? - .ok_or_else(|| { - SignalProtocolError::NoSenderKeyState(format!( - "no sender key record for group {} sender {}", - sender_key_name.group_id(), - sender_key_name.sender_id() - )) - })?; - - let sender_key_state = record - .sender_key_state_mut() - .map_err(|e| anyhow!("Invalid SenderKey session: {:?}", e))?; - - let sender_chain_key = sender_key_state - .sender_chain_key() - .ok_or_else(|| anyhow!("Invalid SenderKey session: missing chain key"))?; - - let message_keys = sender_chain_key.sender_message_key(); - - let mut ciphertext = Vec::new(); - aes_256_cbc_encrypt_into( - plaintext, - message_keys.cipher_key(), - message_keys.iv(), - &mut ciphertext, - ) - .map_err(|_| anyhow!("AES encryption failed"))?; - - let signing_key = sender_key_state - .signing_key_private() - .map_err(|e| anyhow!("Invalid SenderKey session: missing signing key: {:?}", e))?; - - let skm = SenderKeyMessage::new( - SENDERKEY_MESSAGE_CURRENT_VERSION, - sender_key_state.chain_id(), - message_keys.iteration(), - ciphertext.into_boxed_slice(), - csprng, - &signing_key, - )?; - - sender_key_state.set_sender_chain_key(sender_chain_key.next()?); - - sender_key_store - .store_sender_key(sender_key_name, record) - .await?; - - Ok(skm) + // Delegate to the libsignal primitive so the sender-key advance, the wire + // gate, and the iteration lease live in exactly one place. `.context` keeps + // the concrete SignalProtocolError as the source, so callers can still + // downcast NoSenderKeyState to clear stale tracking and retry with SKDM + // redistribution. + crate::libsignal::protocol::group_encrypt(sender_key_store, sender_key_name, plaintext, csprng) + .await + .context("group encrypt failed") } /// Object-safe `SessionStore` that can clone itself into an owned box. The diff --git a/wacore/src/send/tests.rs b/wacore/src/send/tests.rs index 18cc195b9..591e44802 100644 --- a/wacore/src/send/tests.rs +++ b/wacore/src/send/tests.rs @@ -2871,6 +2871,65 @@ mod mark_full_distribution_list { } } + /// An ungated sender-chain advance puts group ciphertext on the wire before + /// the advance is durable, so a reload re-derives the same iteration: one + /// (key, IV) reused toward every member. + #[tokio::test] + async fn encrypt_group_message_leases_the_sender_chain() { + use crate::libsignal::protocol::consts::SENDER_CHAIN_RESERVATION_BATCH; + use crate::libsignal::protocol::{KeyPair, SenderKeyRecord}; + + let name = SenderKeyName::new("g@g.us".to_string(), "me.0".to_string()); + let mut rng = rand::make_rng::(); + let kp = KeyPair::generate(&mut rng); + let mut record = SenderKeyRecord::new_empty(); + record + .add_sender_key_state(3, 1, 0, &[7u8; 32], kp.public_key, Some(kp.private_key)) + .expect("valid sender key state"); + + let mut sks = MemSenderKeyStore::default(); + sks.records.insert(name.clone(), record); + + crate::send::encrypt_group_message(&mut sks, &name, b"hi", &mut rng) + .await + .expect("group encrypt"); + + let stored = sks + .load_sender_key(&name) + .await + .expect("load") + .expect("record present"); + assert_eq!( + stored.reserved_iteration(), + SENDER_CHAIN_RESERVATION_BATCH, + "encrypt_group_message must lease the sender chain" + ); + } + + /// The warm-send recovery downcasts NoSenderKeyState to clear stale device + /// tracking and retry with SKDM redistribution, so erasing the concrete + /// error type here would silently cost the self-heal. + #[tokio::test] + async fn encrypt_group_message_preserves_no_sender_key_state() { + use crate::libsignal::protocol::SignalProtocolError; + + let name = SenderKeyName::new("g@g.us".to_string(), "me.0".to_string()); + let mut rng = rand::make_rng::(); + // Empty store: no local SenderKeyRecord for `name`. + let mut sks = MemSenderKeyStore::default(); + + let err = crate::send::encrypt_group_message(&mut sks, &name, b"hi", &mut rng) + .await + .expect_err("a missing sender key must error"); + assert!( + matches!( + err.downcast_ref::(), + Some(SignalProtocolError::NoSenderKeyState(_)) + ), + "NoSenderKeyState must survive the delegation for the SKDM-redistribution retry, got: {err:#}" + ); + } + // Outgoing group encryption never consumes our own prekeys, and device B // has no bundle (so no session is established for it) — these are never // called; present only to satisfy the generic bounds. diff --git a/wacore/src/store/signal_cache.rs b/wacore/src/store/signal_cache.rs index 09a0020b7..311d0054a 100644 --- a/wacore/src/store/signal_cache.rs +++ b/wacore/src/store/signal_cache.rs @@ -3,11 +3,20 @@ use std::sync::Arc; use anyhow::Result; use async_lock::Mutex; +use rand::RngExt; use crate::libsignal::protocol::{ProtocolAddress, SenderKeyRecord, SessionRecord}; use crate::libsignal::store::sender_key_name::SenderKeyName; use crate::store::traits::SignalStore; +type StoreIncarnation = [u8; 16]; + +fn new_store_incarnation() -> StoreIncarnation { + let mut incarnation = [0; 16]; + rand::make_rng::().fill(&mut incarnation); + incarnation +} + /// Evict clean (non-dirty, non-deleted) entries from a cache HashMap. /// Negative entries (None values) are evicted first. /// @@ -100,6 +109,7 @@ enum SessionEntry { } struct SessionStoreState { + incarnation: StoreIncarnation, cache: HashMap, SessionEntry>, dirty: HashSet>, deleted: HashSet>, @@ -113,8 +123,9 @@ struct SessionStoreState { } impl SessionStoreState { - fn new() -> Self { + fn new(incarnation: StoreIncarnation) -> Self { Self { + incarnation, cache: HashMap::new(), dirty: HashSet::new(), deleted: HashSet::new(), @@ -157,19 +168,16 @@ impl SessionStoreState { self.cache.clear(); self.dirty.clear(); self.deleted.clear(); - // 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. + // Lossy callers have removed the transport; clean callers require no + // pending gate before preserving exact-reload trust. self.reservation_pending.clear(); } + fn discard(&mut self, incarnation: StoreIncarnation) { + self.clear(); + self.incarnation = incarnation; + } + fn evict_if_needed(&mut self, max_entries: usize) { if self.cache.len() <= high_watermark(max_entries) { return; @@ -196,6 +204,7 @@ impl SessionStoreState { // === Sender key object cache (same pattern as sessions) === struct SenderKeyStoreState { + incarnation: StoreIncarnation, // `Arc`-wrapped so a warm `get_sender_key` (the per-send peek reads and the // per-decrypt load) bumps a refcount instead of deep-cloning the record's // `VecDeque` with up to `MAX_MESSAGE_KEYS` message keys each. @@ -209,8 +218,9 @@ struct SenderKeyStoreState { } impl SenderKeyStoreState { - fn new() -> Self { + fn new(incarnation: StoreIncarnation) -> Self { Self { + incarnation, cache: HashMap::new(), dirty: HashSet::new(), wire_gate_pending: HashSet::new(), @@ -247,6 +257,11 @@ impl SenderKeyStoreState { self.wire_gate_pending.clear(); } + fn discard(&mut self, incarnation: StoreIncarnation) { + self.clear(); + self.incarnation = incarnation; + } + fn evict_if_needed(&mut self, max_entries: usize) { evict_clean_entries(&mut self.cache, &self.dirty, None, max_entries); } @@ -334,10 +349,14 @@ impl SignalStoreCache { } pub fn with_max_entries(max_entries: usize) -> Self { + Self::with_max_entries_and_incarnation(max_entries, new_store_incarnation()) + } + + fn with_max_entries_and_incarnation(max_entries: usize, incarnation: StoreIncarnation) -> Self { Self { - sessions: Mutex::new(SessionStoreState::new()), + sessions: Mutex::new(SessionStoreState::new(incarnation)), identities: Mutex::new(ByteStoreState::new()), - sender_keys: Mutex::new(SenderKeyStoreState::new()), + sender_keys: Mutex::new(SenderKeyStoreState::new(incarnation)), removed_prekeys: Mutex::new(HashMap::new()), sender_key_locks: Mutex::new(HashMap::new()), max_entries, @@ -407,9 +426,12 @@ impl SignalStoreCache { // Another task populated this slot while we were loading; // defer to whatever they wrote (Present, CheckedOut, etc). // Deserialize and return without caching to avoid conflict. - return Ok(Some(SessionRecord::deserialize(&bytes)?)); + return Ok(Some(SessionRecord::deserialize_for_store( + &bytes, + &state.incarnation, + )?)); } - let record = SessionRecord::deserialize(&bytes)?; + let record = SessionRecord::deserialize_for_store(&bytes, &state.incarnation)?; state.cache.insert(Arc::from(key), SessionEntry::CheckedOut); state.evict_if_needed(self.max_entries); Ok(Some(record)) @@ -446,7 +468,10 @@ impl SignalStoreCache { let mut state = self.sessions.lock().await; match backend_result { Some(bytes) => { - let record = Arc::new(SessionRecord::deserialize(&bytes)?); + let record = Arc::new(SessionRecord::deserialize_for_store( + &bytes, + &state.incarnation, + )?); if !state.cache.contains_key(key) { state .cache @@ -615,7 +640,10 @@ impl SignalStoreCache { return Ok(cached.clone()); } let record = match backend.get_sender_key(key).await? { - Some(bytes) => Some(Arc::new(SenderKeyRecord::deserialize(&bytes)?)), + Some(bytes) => Some(Arc::new(SenderKeyRecord::deserialize_for_store( + &bytes, + &state.incarnation, + )?)), None => None, }; state.cache.insert(Arc::from(key), record.clone()); @@ -698,6 +726,7 @@ impl SignalStoreCache { // backend call (and one SQLite transaction) per session. { let mut state = self.sessions.lock().await; + let incarnation = state.incarnation; let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect(); let deleted_keys: Vec<_> = state.deleted.iter().cloned().collect(); @@ -709,7 +738,7 @@ impl SignalStoreCache { // deferred below until a later flush sees it durable. if let Some(SessionEntry::Present(record)) = state.cache.get(address.as_ref()) { let mut buf = Vec::new(); - record.serialize_into(&mut buf); + record.serialize_into_for_store(&mut buf, &incarnation); batch.push((address.clone(), bytes::Bytes::from(buf))); } } @@ -821,6 +850,7 @@ impl SignalStoreCache { // Flush sender keys { let mut state = self.sender_keys.lock().await; + let incarnation = state.incarnation; let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect(); let mut batch: Vec<(Arc, bytes::Bytes)> = Vec::new(); @@ -828,7 +858,7 @@ impl SignalStoreCache { match state.cache.get(name.as_ref()) { Some(Some(record)) => { let bytes = record - .serialize() + .serialize_for_store(&incarnation) .map_err(|e| anyhow::anyhow!("sender key serialize for {name}: {e}"))?; batch.push((name.clone(), bytes::Bytes::from(bytes))); } @@ -942,17 +972,45 @@ impl SignalStoreCache { (sessions, identities, sender_keys) } - /// Clear all cached state (used on disconnect/reconnect). - /// Retains allocated capacity for reuse on reconnect. + /// A lossy discard must invalidate exact-reload trust. pub async fn clear(&self) { - self.sessions.lock().await.clear(); + self.clear_with_incarnation(new_store_incarnation()).await; + } + + async fn clear_with_incarnation(&self, incarnation: StoreIncarnation) { + self.sessions.lock().await.discard(incarnation); self.identities.lock().await.clear(); - self.sender_keys.lock().await.clear(); + self.sender_keys.lock().await.discard(incarnation); // Drop buffered prekey removals together with the volatile sessions they // belong to: the promoted session is gone, so the still-durable prekey // must stay so a redelivered pkmsg can rebuild the session. self.removed_prekeys.lock().await.clear(); } + + /// Only a discard can make a post-flush write's stale snapshot reloadable. + #[doc(hidden)] + pub async fn clear_after_flush(&self) { + let mut sessions = self.sessions.lock().await; + if sessions.dirty.is_empty() + && sessions.deleted.is_empty() + && sessions.reservation_pending.is_empty() + { + sessions.clear(); + self.removed_prekeys.lock().await.clear(); + } + drop(sessions); + + let mut identities = self.identities.lock().await; + if identities.dirty.is_empty() && identities.deleted.is_empty() { + identities.clear(); + } + drop(identities); + + let mut sender_keys = self.sender_keys.lock().await; + if sender_keys.dirty.is_empty() && sender_keys.wire_gate_pending.is_empty() { + sender_keys.clear(); + } + } } #[cfg(test)] @@ -1960,6 +2018,318 @@ mod eviction_tests { } } +#[cfg(test)] +mod lease_reload_tests { + use super::*; + use crate::libsignal::protocol::{ + ChainKey, IdentityKey, KeyPair, RootKey, SenderKeyStore, SessionState, + create_sender_key_distribution_message, group_decrypt, group_encrypt, + process_sender_key_distribution_message, + }; + use crate::store::in_memory::InMemoryBackend; + + struct CachedSenderKeyStore<'a> { + cache: &'a SignalStoreCache, + backend: &'a InMemoryBackend, + } + + #[async_trait::async_trait] + impl SenderKeyStore for CachedSenderKeyStore<'_> { + async fn store_sender_key( + &mut self, + name: &SenderKeyName, + record: SenderKeyRecord, + ) -> crate::libsignal::protocol::error::Result<()> { + self.cache.put_sender_key(name, record).await; + Ok(()) + } + + async fn load_sender_key( + &self, + name: &SenderKeyName, + ) -> crate::libsignal::protocol::error::Result> { + Ok(self + .cache + .get_sender_key(name, self.backend) + .await + .expect("test backend") + .map(|record| (*record).clone())) + } + } + + fn sender_key_name() -> SenderKeyName { + SenderKeyName::from_parts("group@g.us", "15550001000@s.whatsapp.net:0") + } + + fn leased_session() -> SessionRecord { + let mut rng = rand::make_rng::(); + let local = IdentityKey::new(KeyPair::generate(&mut rng).public_key); + let remote = IdentityKey::new(KeyPair::generate(&mut rng).public_key); + let base_key = KeyPair::generate(&mut rng).public_key; + let mut state = SessionState::new(3, &local, &remote, &RootKey::new([0; 32]), &base_key); + state.set_sender_chain(&KeyPair::generate(&mut rng), &ChainKey::new([1; 32], 0)); + let mut record = SessionRecord::new(state); + record.reserve_sender_chain_counters(0); + record + } + + fn session_chain_index(record: &SessionRecord) -> u32 { + record + .session_state() + .expect("session") + .get_sender_chain_key() + .expect("sender chain") + .index() + } + + #[tokio::test] + async fn dm_clean_reload_is_exact_but_new_cache_burns_the_lease() { + let backend = InMemoryBackend::new(); + let cache = SignalStoreCache::with_max_entries_and_incarnation( + DEFAULT_MAX_CACHE_ENTRIES, + [0xA1; 16], + ); + let address = ProtocolAddress::new("15550001001".to_string(), 1.into()); + cache.put_session(&address, leased_session()).await; + cache.flush(&backend).await.expect("flush"); + cache.clear_after_flush().await; + + let clean = cache + .get_session(&address, &backend) + .await + .expect("cache load") + .expect("session"); + assert_eq!(session_chain_index(&clean), 0); + + let replacement = SignalStoreCache::with_max_entries_and_incarnation( + DEFAULT_MAX_CACHE_ENTRIES, + [0xB2; 16], + ); + let recovered = replacement + .get_session(&address, &backend) + .await + .expect("recovery load") + .expect("session"); + assert_eq!( + session_chain_index(&recovered), + crate::libsignal::protocol::consts::SENDER_CHAIN_RESERVATION_BATCH + ); + } + + #[tokio::test] + async fn incomplete_session_flush_retains_newer_state_and_fails_closed_on_recovery() { + let backend = InMemoryBackend::new(); + let cache = SignalStoreCache::with_max_entries_and_incarnation( + DEFAULT_MAX_CACHE_ENTRIES, + [0xA1; 16], + ); + let address = ProtocolAddress::new("15550001002".to_string(), 1.into()); + cache.put_session(&address, leased_session()).await; + cache.flush(&backend).await.expect("initial flush"); + + let mut advanced = cache + .get_session(&address, &backend) + .await + .expect("cache load") + .expect("session"); + let next = advanced + .session_state() + .expect("session") + .get_sender_chain_key() + .expect("sender chain") + .next_chain_key() + .expect("chain advance"); + advanced + .session_state_mut() + .expect("session") + .set_sender_chain_key(&next) + .expect("chain update"); + cache.put_session(&address, advanced).await; + + let checked_out = cache + .get_session(&address, &backend) + .await + .expect("cache checkout") + .expect("session"); + cache.flush(&backend).await.expect("skipped flush"); + cache.clear_after_flush().await; + + { + let state = cache.sessions.lock().await; + assert_eq!(state.incarnation, [0xA1; 16]); + assert!(state.dirty.contains(address.as_str())); + assert!(matches!( + state.cache.get(address.as_str()), + Some(SessionEntry::CheckedOut) + )); + } + + let replacement = SignalStoreCache::with_max_entries_and_incarnation( + DEFAULT_MAX_CACHE_ENTRIES, + [0xB2; 16], + ); + let recovered = replacement + .get_session(&address, &backend) + .await + .expect("recovery load") + .expect("session"); + assert_eq!( + session_chain_index(&recovered), + crate::libsignal::protocol::consts::SENDER_CHAIN_RESERVATION_BATCH + ); + + cache.put_session(&address, checked_out).await; + cache.flush(&backend).await.expect("retry flush"); + cache.clear_after_flush().await; + let exact = cache + .get_session(&address, &backend) + .await + .expect("exact reload") + .expect("session"); + assert_eq!(session_chain_index(&exact), 1); + } + + #[tokio::test] + async fn repeated_clean_reloads_keep_group_messages_within_forward_jump_limit() { + let sender_backend = InMemoryBackend::new(); + let sender_cache = SignalStoreCache::new(); + let mut sender = CachedSenderKeyStore { + cache: &sender_cache, + backend: &sender_backend, + }; + let receiver_backend = InMemoryBackend::new(); + let receiver_cache = SignalStoreCache::new(); + let mut receiver = CachedSenderKeyStore { + cache: &receiver_cache, + backend: &receiver_backend, + }; + let name = sender_key_name(); + let mut rng = rand::make_rng::(); + let skdm = create_sender_key_distribution_message(&name, &mut sender, &mut rng) + .await + .expect("sender setup"); + process_sender_key_distribution_message(&name, &skdm, &mut receiver) + .await + .expect("receiver setup"); + + let mut last = None; + for expected_iteration in 0..=32 { + let message = group_encrypt(&mut sender, &name, b"payload", &mut rng) + .await + .expect("group encrypt"); + assert_eq!(message.iteration(), expected_iteration); + last = Some(message); + sender_cache.flush(&sender_backend).await.expect("flush"); + sender_cache.clear_after_flush().await; + } + + let plaintext = group_decrypt(last.expect("message").serialized(), &mut receiver, &name) + .await + .expect("a peer may miss every preceding message"); + assert_eq!(plaintext, b"payload"); + } + + #[tokio::test] + async fn clean_sender_key_eviction_does_not_burn_a_lease() { + let backend = InMemoryBackend::new(); + let cache = SignalStoreCache::new(); + let mut store = CachedSenderKeyStore { + cache: &cache, + backend: &backend, + }; + let name = sender_key_name(); + let mut rng = rand::make_rng::(); + create_sender_key_distribution_message(&name, &mut store, &mut rng) + .await + .expect("sender setup"); + + let first = group_encrypt(&mut store, &name, b"first", &mut rng) + .await + .expect("first send"); + assert_eq!(first.iteration(), 0); + cache.flush(&backend).await.expect("flush"); + assert!( + cache + .sender_keys + .lock() + .await + .cache + .remove(name.cache_key()) + .is_some() + ); + + let second = group_encrypt(&mut store, &name, b"second", &mut rng) + .await + .expect("send after eviction"); + assert_eq!(second.iteration(), 1); + } + + #[tokio::test] + async fn dirty_sender_key_stays_resident_while_recovery_fails_closed() { + let backend = InMemoryBackend::new(); + let cache = SignalStoreCache::with_max_entries_and_incarnation( + DEFAULT_MAX_CACHE_ENTRIES, + [0xA1; 16], + ); + let mut store = CachedSenderKeyStore { + cache: &cache, + backend: &backend, + }; + let name = sender_key_name(); + let mut rng = rand::make_rng::(); + create_sender_key_distribution_message(&name, &mut store, &mut rng) + .await + .expect("sender setup"); + + let first = group_encrypt(&mut store, &name, b"first", &mut rng) + .await + .expect("first send"); + assert_eq!(first.iteration(), 0); + cache.flush(&backend).await.expect("flush"); + + let unflushed = group_encrypt(&mut store, &name, b"unflushed", &mut rng) + .await + .expect("unflushed send"); + assert_eq!(unflushed.iteration(), 1); + cache.clear_after_flush().await; + + { + let state = cache.sender_keys.lock().await; + assert_eq!(state.incarnation, [0xA1; 16]); + assert!(state.dirty.contains(name.cache_key())); + assert!(state.cache.contains_key(name.cache_key())); + } + + let resumed = group_encrypt(&mut store, &name, b"resumed", &mut rng) + .await + .expect("resident send"); + assert_eq!(resumed.iteration(), 2); + + let replacement = SignalStoreCache::with_max_entries_and_incarnation( + DEFAULT_MAX_CACHE_ENTRIES, + [0xB2; 16], + ); + let mut recovered_store = CachedSenderKeyStore { + cache: &replacement, + backend: &backend, + }; + let recovered = group_encrypt(&mut recovered_store, &name, b"recovered", &mut rng) + .await + .expect("recovery send"); + assert_eq!( + recovered.iteration(), + crate::libsignal::protocol::consts::SENDER_CHAIN_RESERVATION_BATCH + ); + + cache.flush(&backend).await.expect("retry flush"); + cache.clear_after_flush().await; + let exact = group_encrypt(&mut store, &name, b"exact", &mut rng) + .await + .expect("exact reload"); + assert_eq!(exact.iteration(), 3); + } +} + #[cfg(test)] mod pre_wire_gate_tests { use super::*; @@ -2098,12 +2468,82 @@ mod pre_wire_gate_tests { assert!(!cache.needs_pre_wire_flush().await); } + /// Cleanup racing a post-flush write must not release its durability gate. + #[tokio::test] + async fn clear_after_flush_retains_every_post_flush_write_and_wire_gate() { + const PREKEY_ID: u32 = 7001; + + let backend = InMemoryBackend::new(); + let cache = SignalStoreCache::with_max_entries_and_incarnation( + DEFAULT_MAX_CACHE_ENTRIES, + [0xA1; 16], + ); + let address = addr("15550000005"); + let name = SenderKeyName::from_parts("g@g.us", "u@s.whatsapp.net:0"); + + cache.flush(&backend).await.unwrap(); + cache.put_session(&address, leased_record()).await; + cache.put_identity(&address, &[7; 32]).await; + backend + .store_prekey(PREKEY_ID, b"prekey", false) + .await + .unwrap(); + cache.remove_prekey(PREKEY_ID, address.as_str()).await; + let mut outbound = SenderKeyRecord::new_empty(); + outbound.mark_wire_gated(); + cache.put_sender_key(&name, outbound).await; + + cache.clear_after_flush().await; + + assert!(cache.needs_pre_wire_flush().await); + { + let sessions = cache.sessions.lock().await; + assert_eq!(sessions.incarnation, [0xA1; 16]); + assert!(sessions.dirty.contains(address.as_str())); + assert!(sessions.reservation_pending.contains(address.as_str())); + } + { + let identities = cache.identities.lock().await; + assert!(identities.dirty.contains(address.as_str())); + } + { + let sender_keys = cache.sender_keys.lock().await; + assert_eq!(sender_keys.incarnation, [0xA1; 16]); + assert!(sender_keys.dirty.contains(name.cache_key())); + assert!(sender_keys.wire_gate_pending.contains(name.cache_key())); + } + assert!(cache.removed_prekeys.lock().await.contains_key(&PREKEY_ID)); + + cache.flush(&backend).await.unwrap(); + + assert!(!cache.needs_pre_wire_flush().await); + assert!( + backend + .get_session(address.as_str()) + .await + .unwrap() + .is_some() + ); + assert_eq!( + backend.load_identity(address.as_str()).await.unwrap(), + Some([7; 32]) + ); + assert!( + backend + .get_sender_key(name.cache_key()) + .await + .unwrap() + .is_some() + ); + assert!(backend.load_prekey(PREKEY_ID).await.unwrap().is_none()); + } + /// 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"); + let a = addr("15550000006"); cache.put_session(&a, leased_record()).await; cache.delete_session(&a).await;