Skip to main content

libsignal_protocol/
session_management.rs

1//
2// Copyright 2026 Signal Messenger, LLC.
3// SPDX-License-Identifier: AGPL-3.0-only
4//
5
6//! Session management and public encrypt/decrypt API for Signal 1:1 messaging.
7//!
8//! This module owns two things:
9//!
10//! 1. **The public API** — [`message_encrypt`], [`message_decrypt`],
11//!    [`message_decrypt_signal`], [`message_decrypt_prekey`]. These are the
12//!    entry points used by the bridge layer and `sealed_sender`.
13//!
14//! 2. **Sesame session management** — the "which session do we use?" logic:
15//!    trial-decryption across current and previous sessions, session promotion
16//!    on success, and session selection for encryption.
17//!
18//! All cryptographic ratchet operations are delegated to
19//! [`TripleRatchet`]. This module has no knowledge of chain keys,
20//! root keys, or SPQR internals.
21
22use std::time::SystemTime;
23
24use displaydoc::Display;
25use libsignal_core::try_scoped;
26use rand::{CryptoRng, Rng};
27
28use crate::consts::MAX_UNACKNOWLEDGED_SESSION_AGE;
29use crate::state::{InvalidSessionError, SessionState};
30use crate::triple_ratchet::{OutgoingTripleRatchet, TripleRatchet};
31use crate::{
32    CiphertextMessage, CiphertextMessageType, Direction, IdentityKeyStore, KyberPayload,
33    KyberPreKeyStore, PreKeySignalMessage, PreKeyStore, ProtocolAddress, Result, SessionRecord,
34    SessionStore, SignalMessage, SignalProtocolError, SignedPreKeyStore, session,
35};
36
37// ── Public API ───────────────────────────────────────────────────────────────
38
39/// Encrypt `ptext` for `remote_address`, loading and storing session state.
40///
41/// If the session is unacknowledged (a locally-initiated session that has not
42/// yet received a response), wraps the [`SignalMessage`] in a
43/// [`PreKeySignalMessage`] containing the original pre-key material.
44pub async fn message_encrypt<R: Rng + CryptoRng>(
45    ptext: &[u8],
46    remote_address: &ProtocolAddress,
47    local_address: &ProtocolAddress,
48    session_store: &mut dyn SessionStore,
49    identity_store: &mut dyn IdentityKeyStore,
50    now: SystemTime,
51    csprng: &mut R,
52) -> Result<CiphertextMessage> {
53    let mut session_record = session_store
54        .load_session(remote_address)
55        .await?
56        .ok_or_else(|| SignalProtocolError::SessionNotFound(remote_address.clone()))?;
57    let session_state = session_record
58        .session_state_mut()
59        .ok_or_else(|| SignalProtocolError::SessionNotFound(remote_address.clone()))?;
60
61    let mut session = OutgoingTripleRatchet::from_session_state(session_state).map_err(|e| {
62        log::error!("session state corrupt for {remote_address}: {e}");
63        e
64    })?;
65
66    let their_identity_key = session_state
67        .remote_identity_key()?
68        .expect("session was valid; must have remote identity key");
69
70    // Pre-key wrapping — session management concern.
71    let message = if let Some(items) = session_state.unacknowledged_pre_key_message_items()? {
72        let timestamp_as_unix_time = items
73            .timestamp()
74            .duration_since(SystemTime::UNIX_EPOCH)
75            .unwrap_or_default()
76            .as_secs();
77        if items.timestamp() + MAX_UNACKNOWLEDGED_SESSION_AGE < now {
78            log::warn!(
79                "stale unacknowledged session for {remote_address} (created at {timestamp_as_unix_time})"
80            );
81            return Err(SignalProtocolError::SessionNotFound(remote_address.clone()));
82        }
83
84        let local_registration_id = session_state.local_registration_id();
85
86        log::info!(
87            "Building PreKeyWhisperMessage for: {} with preKeyId: {} (session created at {})",
88            remote_address,
89            items
90                .pre_key_id()
91                .map_or_else(|| "<none>".to_string(), |id| id.to_string()),
92            timestamp_as_unix_time,
93        );
94
95        let kyber_payload = items
96            .kyber_pre_key_id()
97            .zip(items.kyber_ciphertext())
98            .map(|(id, ciphertext)| KyberPayload::new(id, ciphertext.into()));
99        let signal_message = session.encrypt(ptext, Some(local_address), remote_address, csprng)?;
100
101        CiphertextMessage::PreKeySignalMessage(PreKeySignalMessage::new(
102            session.session_version(),
103            local_registration_id,
104            items.pre_key_id(),
105            items.signed_pre_key_id(),
106            kyber_payload,
107            *items.base_key(),
108            *session.local_identity_key(),
109            signal_message,
110        )?)
111    } else {
112        let signal_message = session.encrypt(ptext, None, remote_address, csprng)?;
113        CiphertextMessage::SignalMessage(signal_message)
114    };
115
116    // In clients, `is_trusted_identity` for the Sending direction checks
117    // whether the session's identity key matches the stored key AND whether the
118    // user has approved it (safety number changes, verification status). This
119    // prevents sending to a contact whose identity has changed without user
120    // acknowledgment.
121    if !identity_store
122        .is_trusted_identity(remote_address, &their_identity_key, Direction::Sending)
123        .await?
124    {
125        log::warn!(
126            "Identity key {} is not trusted for remote address {}",
127            hex::encode(their_identity_key.public_key().public_key_bytes()),
128            remote_address,
129        );
130        return Err(SignalProtocolError::UntrustedIdentity(
131            remote_address.clone(),
132        ));
133    }
134
135    identity_store
136        .save_identity(remote_address, &their_identity_key)
137        .await?;
138
139    // Commit and save session state changes.
140    session.apply_to_session_state(session_state);
141
142    session_store
143        .store_session(remote_address, &session_record)
144        .await?;
145    Ok(message)
146}
147
148/// Decrypt a [`CiphertextMessage`] from `remote_address`.
149///
150/// Routes to [`message_decrypt_signal`] or [`message_decrypt_prekey`] based
151/// on message type.
152#[allow(clippy::too_many_arguments)]
153pub async fn message_decrypt<R: Rng + CryptoRng>(
154    ciphertext: &CiphertextMessage,
155    remote_address: &ProtocolAddress,
156    local_address: &ProtocolAddress,
157    session_store: &mut dyn SessionStore,
158    identity_store: &mut dyn IdentityKeyStore,
159    pre_key_store: &mut dyn PreKeyStore,
160    signed_pre_key_store: &dyn SignedPreKeyStore,
161    kyber_pre_key_store: &mut dyn KyberPreKeyStore,
162    csprng: &mut R,
163) -> Result<Vec<u8>> {
164    match ciphertext {
165        CiphertextMessage::SignalMessage(m) => {
166            message_decrypt_signal(
167                m,
168                remote_address,
169                local_address,
170                session_store,
171                identity_store,
172                csprng,
173            )
174            .await
175        }
176        CiphertextMessage::PreKeySignalMessage(m) => {
177            message_decrypt_prekey(
178                m,
179                remote_address,
180                local_address,
181                session_store,
182                identity_store,
183                pre_key_store,
184                signed_pre_key_store,
185                kyber_pre_key_store,
186                csprng,
187            )
188            .await
189        }
190        _ => Err(SignalProtocolError::InvalidArgument(format!(
191            "message_decrypt cannot be used to decrypt {:?} messages",
192            ciphertext.message_type()
193        ))),
194    }
195}
196
197/// Decrypt a [`PreKeySignalMessage`] from `remote_address`.
198///
199/// Processes the pre-key material to establish a session (via
200/// [`session::process_prekey`]), then decrypts the inner [`SignalMessage`].
201#[allow(clippy::too_many_arguments)]
202pub async fn message_decrypt_prekey<R: Rng + CryptoRng>(
203    ciphertext: &PreKeySignalMessage,
204    remote_address: &ProtocolAddress,
205    local_address: &ProtocolAddress,
206    session_store: &mut dyn SessionStore,
207    identity_store: &mut dyn IdentityKeyStore,
208    pre_key_store: &mut dyn PreKeyStore,
209    signed_pre_key_store: &dyn SignedPreKeyStore,
210    kyber_pre_key_store: &mut dyn KyberPreKeyStore,
211    csprng: &mut R,
212) -> Result<Vec<u8>> {
213    let mut session_record = session_store
214        .load_session(remote_address)
215        .await?
216        .unwrap_or_else(SessionRecord::new_fresh);
217
218    // Make sure we log the session state if we fail to process the pre-key.
219    let process_prekey_result = session::process_prekey(
220        ciphertext,
221        remote_address,
222        local_address,
223        &mut session_record,
224        identity_store,
225        pre_key_store,
226        signed_pre_key_store,
227        kyber_pre_key_store,
228    )
229    .await;
230
231    let (pre_key_used, identity_to_save) = match process_prekey_result {
232        Ok(result) => result,
233        Err(e) => {
234            let errs = [e];
235            log::error!(
236                "{}",
237                format_decryption_failure_log(
238                    remote_address,
239                    &errs,
240                    &session_record,
241                    ciphertext.message()
242                )?
243            );
244            let [e] = errs;
245            return Err(e);
246        }
247    };
248
249    let ptext = try_decrypt_from_record(
250        &mut session_record,
251        remote_address,
252        local_address,
253        ciphertext.message(),
254        CiphertextMessageType::PreKey,
255        csprng,
256    )?;
257
258    identity_store
259        .save_identity(
260            identity_to_save.remote_address,
261            identity_to_save.their_identity_key,
262        )
263        .await?;
264
265    if let Some(pre_key_used) = pre_key_used {
266        if let Some(kyber_pre_key_id) = pre_key_used.kyber_pre_key_id {
267            kyber_pre_key_store
268                .mark_kyber_pre_key_used(
269                    kyber_pre_key_id,
270                    pre_key_used.signed_ec_pre_key_id,
271                    ciphertext.base_key(),
272                )
273                .await?;
274        }
275
276        if let Some(pre_key_id) = pre_key_used.one_time_ec_pre_key_id {
277            pre_key_store.remove_pre_key(pre_key_id).await?;
278        }
279    }
280
281    session_store
282        .store_session(remote_address, &session_record)
283        .await?;
284
285    Ok(ptext)
286}
287
288/// Decrypt a [`SignalMessage`] from `remote_address`.
289///
290/// Tries all sessions in the session record. Checks identity key trust
291/// after decryption.
292pub async fn message_decrypt_signal<R: Rng + CryptoRng>(
293    ciphertext: &SignalMessage,
294    remote_address: &ProtocolAddress,
295    local_address: &ProtocolAddress,
296    session_store: &mut dyn SessionStore,
297    identity_store: &mut dyn IdentityKeyStore,
298    csprng: &mut R,
299) -> Result<Vec<u8>> {
300    let mut session_record = session_store
301        .load_session(remote_address)
302        .await?
303        .ok_or_else(|| SignalProtocolError::SessionNotFound(remote_address.clone()))?;
304
305    let ptext = try_decrypt_from_record(
306        &mut session_record,
307        remote_address,
308        local_address,
309        ciphertext,
310        CiphertextMessageType::Whisper,
311        csprng,
312    )?;
313
314    // Why are we performing this check after decryption instead of before?
315    let their_identity_key = session_record
316        .session_state()
317        .expect("successfully decrypted; must have a current state")
318        .remote_identity_key()
319        .expect("successfully decrypted; must have a remote identity key")
320        .expect("successfully decrypted; must have a remote identity key");
321
322    if !identity_store
323        .is_trusted_identity(remote_address, &their_identity_key, Direction::Receiving)
324        .await?
325    {
326        log::warn!(
327            "Identity key {} is not trusted for remote address {}",
328            hex::encode(their_identity_key.public_key().public_key_bytes()),
329            remote_address,
330        );
331        return Err(SignalProtocolError::UntrustedIdentity(
332            remote_address.clone(),
333        ));
334    }
335
336    identity_store
337        .save_identity(remote_address, &their_identity_key)
338        .await?;
339
340    session_store
341        .store_session(remote_address, &session_record)
342        .await?;
343
344    Ok(ptext)
345}
346
347// ── Session management (Sesame) ──────────────────────────────────────────────
348
349/// Try to decrypt `ciphertext` against every session in `record`, in order.
350///
351/// Tries the current session first, then previous sessions. On success from
352/// a previous session, promotes that session to current (Sesame behavior).
353///
354/// `original_message_type` is `Whisper` for normal messages and `PreKey` for
355/// the inner `SignalMessage` of a pre-key message. When it is `PreKey`, we
356/// skip the fallback to previous sessions — a PreKey message establishes a
357/// fresh session and should always match the current one.
358pub(crate) fn try_decrypt_from_record<R: Rng + CryptoRng>(
359    record: &mut SessionRecord,
360    remote_address: &ProtocolAddress,
361    local_address: &ProtocolAddress,
362    ciphertext: &SignalMessage,
363    original_message_type: CiphertextMessageType,
364    csprng: &mut R,
365) -> Result<Vec<u8>> {
366    debug_assert!(matches!(
367        original_message_type,
368        CiphertextMessageType::Whisper | CiphertextMessageType::PreKey
369    ));
370    let ciphertext_version = ciphertext.message_version() as u32;
371
372    let log_failure = |label: &str, state: &SessionState, error: &SignalProtocolError| {
373        log::warn!(
374            "Failed to decrypt {:?} message with ratchet key: {} and counter: {}. \
375             Session loaded for {}. {} session has base key: {} and counter: {}. {}",
376            original_message_type,
377            hex::encode(ciphertext.sender_ratchet_key().public_key_bytes()),
378            ciphertext.counter(),
379            remote_address,
380            label,
381            state
382                .sender_ratchet_key_for_logging()
383                .unwrap_or_else(|e| format!("<error: {e}>")),
384            state.previous_counter(),
385            error
386        );
387    };
388
389    let mut errs = vec![];
390
391    // ── Try current session ──────────────────────────────────────────────────
392
393    if let Some(current_state) = record.session_state() {
394        let mut current_state = current_state.clone();
395
396        if current_state.session_version()? != ciphertext_version {
397            let e = SignalProtocolError::UnrecognizedMessageVersion(ciphertext_version);
398            log_failure("Current", &current_state, &e);
399            errs.push(e);
400        } else {
401            match try_decrypt_with_state(
402                &mut current_state,
403                remote_address,
404                local_address,
405                ciphertext,
406                original_message_type,
407                CurrentOrPrevious::Current,
408                csprng,
409            ) {
410                Ok(ptext) => {
411                    log::info!(
412                        "decrypted {:?} message from {} with current session state (base key {})",
413                        original_message_type,
414                        remote_address,
415                        current_state
416                            .sender_ratchet_key_for_logging()
417                            .expect("successful decrypt always has a valid base key"),
418                    );
419                    record.set_session_state(current_state);
420                    return Ok(ptext);
421                }
422                Err(e @ SignalProtocolError::DuplicatedMessage(_, _)) => return Err(e),
423                Err(e) => {
424                    log_failure("Current", &current_state, &e);
425                    errs.push(e);
426                    match original_message_type {
427                        CiphertextMessageType::PreKey => {
428                            // A PreKey message creates a session and then decrypts a Whisper message
429                            // using that session. No need to check older sessions.
430                            log::error!(
431                                "{}",
432                                format_decryption_failure_log(
433                                    remote_address,
434                                    &errs,
435                                    record,
436                                    ciphertext,
437                                )?
438                            );
439                            // Note that we don't propagate `e` here; we always return InvalidMessage,
440                            // as we would for a Whisper message that tried several sessions.
441                            return Err(SignalProtocolError::InvalidMessage(
442                                original_message_type,
443                                "decryption failed",
444                            ));
445                        }
446                        CiphertextMessageType::Whisper => {}
447                        CiphertextMessageType::SenderKey | CiphertextMessageType::Plaintext => {
448                            unreachable!("should not be using Double Ratchet for these")
449                        }
450                    }
451                }
452            }
453        }
454    }
455
456    // ── Try previous sessions (Whisper only) ─────────────────────────────────
457
458    let mut promoted = None;
459
460    for (idx, previous) in record.previous_session_states().enumerate() {
461        let mut previous = match previous {
462            Ok(previous) => previous,
463            Err(e) => {
464                let e: SignalProtocolError = e.into();
465                log::warn!(
466                    "Skipping corrupt previous session {} for {}: {}",
467                    idx,
468                    remote_address,
469                    e
470                );
471                errs.push(e);
472                continue;
473            }
474        };
475
476        if previous.session_version()? != ciphertext_version {
477            let e = SignalProtocolError::UnrecognizedMessageVersion(ciphertext_version);
478            log_failure("Previous", &previous, &e);
479            errs.push(e);
480            continue;
481        }
482
483        match try_decrypt_with_state(
484            &mut previous,
485            remote_address,
486            local_address,
487            ciphertext,
488            original_message_type,
489            CurrentOrPrevious::Previous,
490            csprng,
491        ) {
492            Ok(ptext) => {
493                log::info!(
494                    "decrypted {:?} message from {} with PREVIOUS session state (base key {})",
495                    original_message_type,
496                    remote_address,
497                    previous
498                        .sender_ratchet_key_for_logging()
499                        .expect("successful decrypt always has a valid base key"),
500                );
501                promoted = Some((ptext, idx, previous));
502                break;
503            }
504            Err(e @ SignalProtocolError::DuplicatedMessage(_, _)) => return Err(e),
505            Err(e) => {
506                log_failure("Previous", &previous, &e);
507                errs.push(e);
508            }
509        }
510    }
511
512    if let Some((ptext, idx, updated)) = promoted {
513        // Sesame: promote the successful previous session to current.
514        // The upcoming session management update will remove this promotion.
515        record.promote_old_session(idx, updated);
516        Ok(ptext)
517    } else {
518        let previous_state_count = || record.previous_session_states().len();
519        if let Some(current_state) = record.session_state() {
520            log::error!(
521                "No valid session for recipient: {}, current session base key {}, \
522                 number of previous states: {}",
523                remote_address,
524                current_state
525                    .sender_ratchet_key_for_logging()
526                    .unwrap_or_else(|e| format!("<error: {e}>")),
527                previous_state_count(),
528            );
529        } else {
530            log::error!(
531                "No valid session for recipient: {}, (no current session state), \
532                 number of previous states: {}",
533                remote_address,
534                previous_state_count(),
535            );
536        }
537        log::error!(
538            "{}",
539            format_decryption_failure_log(remote_address, &errs, record, ciphertext)?
540        );
541        Err(SignalProtocolError::InvalidMessage(
542            original_message_type,
543            "decryption failed",
544        ))
545    }
546}
547
548// ── Per-session decrypt ──────────────────────────────────────────────────────
549
550/// Attempt to decrypt `ciphertext` using the crypto state in `state`.
551///
552/// Caller must only pass version-compatible ciphertext/session pairs.
553///
554/// Constructs a [`TripleRatchet`], delegates the actual decryption, and writes
555/// updated state back on success. On failure, `state` is unchanged.
556pub(crate) fn try_decrypt_with_state<R: Rng + CryptoRng>(
557    state: &mut SessionState,
558    remote_address: &ProtocolAddress,
559    local_address: &ProtocolAddress,
560    ciphertext: &SignalMessage,
561    original_message_type: CiphertextMessageType,
562    curr_or_prev_for_logging: CurrentOrPrevious,
563    csprng: &mut R,
564) -> Result<Vec<u8>> {
565    debug_assert_eq!(
566        state.session_version()?,
567        ciphertext.message_version() as u32
568    );
569
570    let self_session = try_scoped::<bool, InvalidSessionError>(|| {
571        Ok(state.local_identity_key()?.is_same_account(
572            local_address,
573            &state
574                .remote_identity_key()?
575                .ok_or(InvalidSessionError("missing remote identity key"))?,
576            remote_address,
577        ))
578    })
579    .inspect_err(|e| log::warn!("Failed to determine self_session: {}", e))
580    .unwrap_or_default();
581    let mut session = TripleRatchet::from_session_state(state, self_session)?;
582
583    let ptext = session.decrypt(
584        remote_address,
585        local_address,
586        ciphertext,
587        original_message_type,
588        curr_or_prev_for_logging,
589        csprng,
590    )?;
591
592    session.apply_to_session_state(state);
593    state.clear_unacknowledged_pre_key_message();
594
595    Ok(ptext)
596}
597
598// ── Logging helpers ──────────────────────────────────────────────────────────
599
600pub(crate) fn format_decryption_failure_log(
601    remote_address_for_logging: &ProtocolAddress,
602    mut errs: &[SignalProtocolError],
603    record: &SessionRecord,
604    ciphertext: &SignalMessage,
605) -> Result<String> {
606    fn append_session_summary(
607        lines: &mut Vec<String>,
608        idx: usize,
609        state: std::result::Result<&SessionState, InvalidSessionError>,
610        err: Option<&SignalProtocolError>,
611    ) {
612        let chains = state.map(|state| state.all_receiver_chain_logging_info());
613        match (err, &chains) {
614            (Some(err), Ok(chains)) => {
615                lines.push(format!(
616                    "Candidate session {} failed with '{}', had {} receiver chains",
617                    idx,
618                    err,
619                    chains.len()
620                ));
621            }
622            (Some(err), Err(state_err)) => {
623                lines.push(format!(
624                    "Candidate session {idx} failed with '{err}'; \
625                     cannot get receiver chain info ({state_err})",
626                ));
627            }
628            (None, Ok(chains)) => {
629                lines.push(format!(
630                    "Candidate session {} had {} receiver chains",
631                    idx,
632                    chains.len()
633                ));
634            }
635            (None, Err(state_err)) => {
636                lines.push(format!(
637                    "Candidate session {idx}: cannot get receiver chain info ({state_err})",
638                ));
639            }
640        }
641
642        if let Ok(chains) = chains {
643            for chain in chains {
644                let chain_idx = match chain.1 {
645                    Some(i) => i.to_string(),
646                    None => "missing in protobuf".to_string(),
647                };
648                lines.push(format!(
649                    "Receiver chain with sender ratchet public key {} chain key index {}",
650                    hex::encode(chain.0),
651                    chain_idx
652                ));
653            }
654        }
655    }
656
657    let mut lines = vec![];
658    lines.push(format!(
659        "Message from {} failed to decrypt; sender ratchet public key {} message counter {}",
660        remote_address_for_logging,
661        hex::encode(ciphertext.sender_ratchet_key().public_key_bytes()),
662        ciphertext.counter()
663    ));
664
665    if let Some(current_session) = record.session_state() {
666        let err = errs.first();
667        if err.is_some() {
668            errs = &errs[1..];
669        }
670        append_session_summary(&mut lines, 0, Ok(current_session), err);
671    } else {
672        lines.push("No current session".to_string());
673    }
674
675    for (idx, (state, err)) in record
676        .previous_session_states()
677        .zip(errs.iter().map(Some).chain(std::iter::repeat(None)))
678        .enumerate()
679    {
680        let state = match state {
681            Ok(ref state) => Ok(state),
682            Err(err) => Err(err),
683        };
684        append_session_summary(&mut lines, idx + 1, state, err);
685    }
686
687    Ok(lines.join("\n"))
688}
689
690#[derive(Clone, Copy, Display)]
691pub(crate) enum CurrentOrPrevious {
692    /// current
693    Current,
694    /// previous
695    Previous,
696}
697
698// ── Comparison proptest ──────────────────────────────────────────────────────
699//
700// Verifies that the refactored encrypt/decrypt path produces identical results
701// to the legacy snapshot for any message sequence.
702#[cfg(test)]
703mod legacy_interop_tests {
704    // These tests live next to `session_management` rather than under
705    // `rust/protocol/tests/` because they compare the refactored code against
706    // the private `session_cipher_legacy` implementation and also assert
707    // byte-level equivalence of internal persisted state. That makes them
708    // implementation-regression tests for this refactor, not normal public API
709    // integration tests.
710    //
711    // This harness is temporary. Once we are confident in the refactor, remove
712    // `session_cipher_legacy` and the new-vs-legacy equivalence tests along
713    // with it.
714    use futures_util::FutureExt;
715    use libsignal_protocol_test_support::Event;
716    use proptest::prelude::*;
717    use prost::Message;
718    use rand::SeedableRng;
719    use rand_chacha::ChaCha8Rng;
720
721    use super::*;
722    use crate::proto::storage::RecordStructure;
723    use crate::ratchet::{
724        AliceSignalProtocolParameters, BobSignalProtocolParameters,
725        initialize_alice_session_record, initialize_bob_session_record,
726    };
727    use crate::{
728        DecryptionErrorMessage, DeviceId, GenericSignedPreKey, IdentityKeyPair,
729        InMemSignalProtocolStore, KeyPair, KyberPreKeyId, KyberPreKeyRecord, PlaintextContent,
730        PreKeyBundle, PreKeyId, PreKeyRecord, ProtocolAddress, SessionRecord,
731        SessionUsabilityRequirements, SignalProtocolError, SignedPreKeyId, SignedPreKeyRecord,
732        Timestamp, extract_decryption_error_message_from_serialized_content, process_prekey_bundle,
733        session_cipher_legacy as legacy,
734    };
735
736    #[derive(Clone, Copy, PartialEq, Eq, Debug)]
737    enum MessageStatus {
738        Sent,
739        Dropped,
740        Delivered,
741    }
742
743    #[derive(Clone)]
744    struct DualLocalState {
745        new_store: InMemSignalProtocolStore,
746        legacy_store: InMemSignalProtocolStore,
747        pre_key_count: u32,
748    }
749
750    struct DualParticipant {
751        address: ProtocolAddress,
752        message_queue: Vec<(CiphertextMessage, u64)>,
753        state: DualLocalState,
754        snapshots: Vec<DualLocalState>,
755        message_send_log: Vec<MessageStatus>,
756    }
757
758    /// Build a matched (alice, bob) session pair from a seeded RNG.
759    fn setup_stores(
760        rng: &mut ChaCha8Rng,
761    ) -> (
762        InMemSignalProtocolStore,
763        InMemSignalProtocolStore,
764        ProtocolAddress,
765        ProtocolAddress,
766    ) {
767        let alice_identity = IdentityKeyPair::generate(rng);
768        let bob_identity = IdentityKeyPair::generate(rng);
769
770        let alice_base_key = KeyPair::generate(rng);
771        let bob_signed_pre_key = KeyPair::generate(rng);
772        let bob_kyber_key = crate::kem::KeyPair::generate(crate::kem::KeyType::Kyber1024, rng);
773
774        let alice_params = AliceSignalProtocolParameters::new(
775            alice_identity,
776            alice_base_key,
777            *bob_identity.identity_key(),
778            bob_signed_pre_key.public_key,
779            bob_signed_pre_key.public_key,
780            bob_kyber_key.public_key.clone(),
781            false,
782        );
783
784        let alice_record =
785            initialize_alice_session_record(&alice_params, rng).expect("alice session init");
786        let kyber_ct: Box<[u8]> = alice_record
787            .get_kyber_ciphertext()
788            .expect("session valid")
789            .expect("has kyber ciphertext")
790            .clone()
791            .into_boxed_slice();
792
793        let bob_params = BobSignalProtocolParameters::new(
794            bob_identity,
795            bob_signed_pre_key,
796            None,
797            bob_kyber_key,
798            *alice_identity.identity_key(),
799            alice_base_key.public_key,
800            &kyber_ct,
801            false,
802        );
803
804        let bob_record = initialize_bob_session_record(&bob_params, &bob_signed_pre_key)
805            .expect("bob session init");
806
807        let alice_address = ProtocolAddress::new(
808            "57721566-4901-5328-6060-651209008240".to_owned(),
809            DeviceId::new(1).unwrap(),
810        );
811        let bob_address = ProtocolAddress::new(
812            "26149721-2847-6427-8375-542683860869".to_owned(),
813            DeviceId::new(1).unwrap(),
814        );
815
816        let mut alice_store = InMemSignalProtocolStore::new(alice_identity, 1).unwrap();
817        let mut bob_store = InMemSignalProtocolStore::new(bob_identity, 2).unwrap();
818
819        alice_store
820            .session_store
821            .store_session(&bob_address, &alice_record)
822            .now_or_never()
823            .unwrap()
824            .unwrap();
825        bob_store
826            .session_store
827            .store_session(&alice_address, &bob_record)
828            .now_or_never()
829            .unwrap()
830            .unwrap();
831
832        (alice_store, bob_store, alice_address, bob_address)
833    }
834
835    /// Create a pre-key bundle for Bob, storing new key material in his store.
836    ///
837    /// The `*_id` parameters must not collide with any IDs already in the
838    /// store. Using a monotonically increasing generation counter (1, 2, …)
839    /// is sufficient.
840    fn create_bob_bundle(
841        bob_store: &mut InMemSignalProtocolStore,
842        pre_key_id: u32,
843        signed_pre_key_id: u32,
844        kyber_pre_key_id: u32,
845        rng: &mut ChaCha8Rng,
846    ) -> PreKeyBundle {
847        let identity_key_pair = bob_store
848            .get_identity_key_pair()
849            .now_or_never()
850            .unwrap()
851            .unwrap();
852
853        let pre_key = KeyPair::generate(rng);
854        let signed_pre_key = KeyPair::generate(rng);
855        let kyber_key = crate::kem::KeyPair::generate(crate::kem::KeyType::Kyber1024, rng);
856
857        let pk_id = PreKeyId::from(pre_key_id);
858        let spk_id = SignedPreKeyId::from(signed_pre_key_id);
859        let kpk_id = KyberPreKeyId::from(kyber_pre_key_id);
860
861        let spk_sig = identity_key_pair
862            .private_key()
863            .calculate_signature(&signed_pre_key.public_key.serialize(), rng)
864            .unwrap();
865        let kpk_sig = identity_key_pair
866            .private_key()
867            .calculate_signature(&kyber_key.public_key.serialize(), rng)
868            .unwrap();
869
870        bob_store
871            .save_pre_key(pk_id, &PreKeyRecord::new(pk_id, &pre_key))
872            .now_or_never()
873            .unwrap()
874            .unwrap();
875        bob_store
876            .save_signed_pre_key(
877                spk_id,
878                &SignedPreKeyRecord::new(
879                    spk_id,
880                    Timestamp::from_epoch_millis(42),
881                    &signed_pre_key,
882                    &spk_sig,
883                ),
884            )
885            .now_or_never()
886            .unwrap()
887            .unwrap();
888        bob_store
889            .save_kyber_pre_key(
890                kpk_id,
891                &KyberPreKeyRecord::new(
892                    kpk_id,
893                    Timestamp::from_epoch_millis(43),
894                    &kyber_key,
895                    &kpk_sig,
896                ),
897            )
898            .now_or_never()
899            .unwrap()
900            .unwrap();
901
902        let reg_id = bob_store
903            .get_local_registration_id()
904            .now_or_never()
905            .unwrap()
906            .unwrap();
907
908        PreKeyBundle::new(
909            reg_id,
910            DeviceId::new(1).unwrap(),
911            Some((pk_id, pre_key.public_key)),
912            spk_id,
913            signed_pre_key.public_key,
914            spk_sig.to_vec(),
915            kpk_id,
916            kyber_key.public_key.clone(),
917            kpk_sig.to_vec(),
918            *identity_key_pair.identity_key(),
919        )
920        .unwrap()
921    }
922
923    #[test]
924    fn encrypt_preserves_corruption_error_instead_of_session_not_found() {
925        let mut rng = ChaCha8Rng::seed_from_u64(0xC0FFEE);
926        let (mut alice_store, _bob_store, alice_address, bob_address) = setup_stores(&mut rng);
927        let now = SystemTime::now();
928
929        let good_record = alice_store
930            .session_store
931            .load_session(&bob_address)
932            .now_or_never()
933            .expect("sync")
934            .expect("load succeeded")
935            .expect("session exists");
936
937        let serialized = good_record.serialize().expect("serialize");
938        let mut record_pb = RecordStructure::decode(serialized.as_slice()).expect("decode record");
939        record_pb
940            .current_session
941            .as_mut()
942            .expect("current session")
943            .remote_identity_public = vec![0xFF];
944
945        let corrupted_record = SessionRecord::deserialize(record_pb.encode_to_vec().as_slice())
946            .expect("deserialize corrupted record");
947
948        alice_store
949            .session_store
950            .store_session(&bob_address, &corrupted_record)
951            .now_or_never()
952            .expect("sync")
953            .expect("store succeeded");
954
955        let legacy_err = legacy::legacy_message_encrypt(
956            b"test",
957            &bob_address,
958            &alice_address,
959            &mut alice_store.session_store,
960            &mut alice_store.identity_store,
961            now,
962            &mut rng,
963        )
964        .now_or_never()
965        .expect("sync")
966        .expect_err("legacy encrypt should fail on corrupted state");
967        assert!(
968            matches!(
969                legacy_err,
970                SignalProtocolError::InvalidSessionStructure("invalid remote identity key")
971            ),
972            "unexpected legacy error: {legacy_err:?}"
973        );
974
975        alice_store
976            .session_store
977            .store_session(&bob_address, &corrupted_record)
978            .now_or_never()
979            .expect("sync")
980            .expect("store succeeded");
981
982        let new_err = message_encrypt(
983            b"test",
984            &bob_address,
985            &alice_address,
986            &mut alice_store.session_store,
987            &mut alice_store.identity_store,
988            now,
989            &mut rng,
990        )
991        .now_or_never()
992        .expect("sync")
993        .expect_err("new encrypt should fail on corrupted state");
994        assert!(
995            matches!(
996                new_err,
997                SignalProtocolError::InvalidSessionStructure("invalid remote identity key")
998            ),
999            "unexpected new error: {new_err:?}"
1000        );
1001    }
1002
1003    #[test]
1004    fn encrypt_ignores_corrupt_unused_receiver_chain() {
1005        let mut rng = ChaCha8Rng::seed_from_u64(0xACE55);
1006        let (mut alice_store, _bob_store, alice_address, bob_address) = setup_stores(&mut rng);
1007        let now = SystemTime::now();
1008
1009        let good_record = alice_store
1010            .session_store
1011            .load_session(&bob_address)
1012            .now_or_never()
1013            .expect("sync")
1014            .expect("load succeeded")
1015            .expect("session exists");
1016
1017        let serialized = good_record.serialize().expect("serialize");
1018        let mut record_pb = RecordStructure::decode(serialized.as_slice()).expect("decode record");
1019        let current_session = record_pb.current_session.as_mut().expect("current session");
1020        assert!(
1021            !current_session.receiver_chains.is_empty(),
1022            "expected at least one receiver chain"
1023        );
1024        current_session.receiver_chains[0].sender_ratchet_key = vec![0xFF];
1025
1026        let corrupted_record = SessionRecord::deserialize(record_pb.encode_to_vec().as_slice())
1027            .expect("deserialize corrupted record");
1028
1029        alice_store
1030            .session_store
1031            .store_session(&bob_address, &corrupted_record)
1032            .now_or_never()
1033            .expect("sync")
1034            .expect("store succeeded");
1035
1036        let mut legacy_rng = rng.clone();
1037        let legacy_ct = legacy::legacy_message_encrypt(
1038            b"test",
1039            &bob_address,
1040            &alice_address,
1041            &mut alice_store.session_store,
1042            &mut alice_store.identity_store,
1043            now,
1044            &mut legacy_rng,
1045        )
1046        .now_or_never()
1047        .expect("sync")
1048        .expect("legacy encrypt should ignore unused receiver-chain corruption");
1049
1050        alice_store
1051            .session_store
1052            .store_session(&bob_address, &corrupted_record)
1053            .now_or_never()
1054            .expect("sync")
1055            .expect("store succeeded");
1056
1057        let new_ct = message_encrypt(
1058            b"test",
1059            &bob_address,
1060            &alice_address,
1061            &mut alice_store.session_store,
1062            &mut alice_store.identity_store,
1063            now,
1064            &mut rng,
1065        )
1066        .now_or_never()
1067        .expect("sync")
1068        .expect("new encrypt should ignore unused receiver-chain corruption");
1069
1070        let legacy_msg = match legacy_ct {
1071            CiphertextMessage::SignalMessage(m) => m,
1072            other => panic!(
1073                "expected SignalMessage from legacy enc, got {:?}",
1074                other.message_type()
1075            ),
1076        };
1077        let new_msg = match new_ct {
1078            CiphertextMessage::SignalMessage(m) => m,
1079            other => panic!(
1080                "expected SignalMessage from new enc, got {:?}",
1081                other.message_type()
1082            ),
1083        };
1084
1085        assert_eq!(legacy_msg.serialized(), new_msg.serialized());
1086    }
1087
1088    #[test]
1089    fn decrypt_skips_corrupt_previous_session_and_uses_later_valid_previous() {
1090        let mut rng = ChaCha8Rng::seed_from_u64(0xBAD5EED);
1091        let (mut alice_store, mut bob_store, alice_address, bob_address) = setup_stores(&mut rng);
1092        let now = SystemTime::now();
1093
1094        let delayed_plaintext = b"delayed on session A".to_vec();
1095
1096        let delayed_ct = legacy::legacy_message_encrypt(
1097            &delayed_plaintext,
1098            &bob_address,
1099            &alice_address,
1100            &mut alice_store.session_store,
1101            &mut alice_store.identity_store,
1102            now,
1103            &mut rng,
1104        )
1105        .now_or_never()
1106        .expect("sync")
1107        .expect("delayed legacy enc");
1108
1109        let delayed_signal_msg = match delayed_ct {
1110            CiphertextMessage::SignalMessage(m) => m,
1111            other => panic!(
1112                "expected SignalMessage for delayed msg, got {:?}",
1113                other.message_type()
1114            ),
1115        };
1116
1117        let bundle = create_bob_bundle(&mut bob_store, 1, 1, 1, &mut rng);
1118        process_prekey_bundle(
1119            &bob_address,
1120            &alice_address,
1121            &mut alice_store.session_store,
1122            &mut alice_store.identity_store,
1123            &bundle,
1124            now,
1125            &mut rng,
1126        )
1127        .now_or_never()
1128        .expect("sync")
1129        .expect("process_prekey_bundle");
1130
1131        let session_b_init = message_encrypt(
1132            b"session B init",
1133            &bob_address,
1134            &alice_address,
1135            &mut alice_store.session_store,
1136            &mut alice_store.identity_store,
1137            now,
1138            &mut rng,
1139        )
1140        .now_or_never()
1141        .expect("sync")
1142        .expect("session B init enc");
1143
1144        message_decrypt(
1145            &session_b_init,
1146            &alice_address,
1147            &bob_address,
1148            &mut bob_store.session_store,
1149            &mut bob_store.identity_store,
1150            &mut bob_store.pre_key_store,
1151            &bob_store.signed_pre_key_store,
1152            &mut bob_store.kyber_pre_key_store,
1153            &mut rng,
1154        )
1155        .now_or_never()
1156        .expect("sync")
1157        .expect("session B init dec");
1158
1159        let session_b_ack = message_encrypt(
1160            b"session B ack",
1161            &alice_address,
1162            &bob_address,
1163            &mut bob_store.session_store,
1164            &mut bob_store.identity_store,
1165            now,
1166            &mut rng,
1167        )
1168        .now_or_never()
1169        .expect("sync")
1170        .expect("session B ack enc");
1171
1172        let session_b_ack_signal = match &session_b_ack {
1173            CiphertextMessage::SignalMessage(m) => m,
1174            other => panic!(
1175                "expected Whisper for session B ack, got {:?}",
1176                other.message_type()
1177            ),
1178        };
1179        message_decrypt_signal(
1180            session_b_ack_signal,
1181            &bob_address,
1182            &alice_address,
1183            &mut alice_store.session_store,
1184            &mut alice_store.identity_store,
1185            &mut rng,
1186        )
1187        .now_or_never()
1188        .expect("sync")
1189        .expect("session B ack dec");
1190
1191        let bob_record = bob_store
1192            .session_store
1193            .load_session(&alice_address)
1194            .now_or_never()
1195            .expect("sync")
1196            .expect("load succeeded")
1197            .expect("session exists");
1198        let serialized = bob_record.serialize().expect("serialize");
1199        let mut record_pb = RecordStructure::decode(serialized.as_slice()).expect("decode record");
1200        assert_eq!(
1201            record_pb.previous_sessions.len(),
1202            1,
1203            "expected one valid previous session"
1204        );
1205        record_pb.previous_sessions.insert(0, vec![0xFF]);
1206        let corrupted_record = SessionRecord::deserialize(record_pb.encode_to_vec().as_slice())
1207            .expect("deserialize mutated record");
1208        bob_store
1209            .session_store
1210            .store_session(&alice_address, &corrupted_record)
1211            .now_or_never()
1212            .expect("sync")
1213            .expect("store succeeded");
1214
1215        let ptext = message_decrypt_signal(
1216            &delayed_signal_msg,
1217            &alice_address,
1218            &bob_address,
1219            &mut bob_store.session_store,
1220            &mut bob_store.identity_store,
1221            &mut rng,
1222        )
1223        .now_or_never()
1224        .expect("sync")
1225        .expect("delayed msg dec via valid later previous session");
1226
1227        assert_eq!(ptext, delayed_plaintext);
1228    }
1229
1230    fn setup_two_alice_receiver_chains_on_bob(
1231        rng: &mut ChaCha8Rng,
1232    ) -> (
1233        InMemSignalProtocolStore,
1234        InMemSignalProtocolStore,
1235        ProtocolAddress,
1236        ProtocolAddress,
1237        SignalMessage,
1238    ) {
1239        let (mut alice_store, mut bob_store, alice_address, bob_address) = setup_stores(rng);
1240        let now = SystemTime::now();
1241
1242        let delayed_ct = message_encrypt(
1243            b"delayed old",
1244            &bob_address,
1245            &alice_address,
1246            &mut alice_store.session_store,
1247            &mut alice_store.identity_store,
1248            now,
1249            rng,
1250        )
1251        .now_or_never()
1252        .expect("sync")
1253        .expect("delayed old enc");
1254        let delayed_signal_msg = match delayed_ct {
1255            CiphertextMessage::SignalMessage(m) => m,
1256            other => panic!(
1257                "expected delayed SignalMessage, got {:?}",
1258                other.message_type()
1259            ),
1260        };
1261
1262        let trigger_ct = message_encrypt(
1263            b"trigger old chain advancement",
1264            &bob_address,
1265            &alice_address,
1266            &mut alice_store.session_store,
1267            &mut alice_store.identity_store,
1268            now,
1269            rng,
1270        )
1271        .now_or_never()
1272        .expect("sync")
1273        .expect("trigger enc");
1274        let trigger_signal_msg = match trigger_ct {
1275            CiphertextMessage::SignalMessage(m) => m,
1276            other => panic!(
1277                "expected trigger SignalMessage, got {:?}",
1278                other.message_type()
1279            ),
1280        };
1281        message_decrypt_signal(
1282            &trigger_signal_msg,
1283            &alice_address,
1284            &bob_address,
1285            &mut bob_store.session_store,
1286            &mut bob_store.identity_store,
1287            rng,
1288        )
1289        .now_or_never()
1290        .expect("sync")
1291        .expect("trigger dec");
1292
1293        let bob_reply_ct = message_encrypt(
1294            b"bob reply new ratchet",
1295            &alice_address,
1296            &bob_address,
1297            &mut bob_store.session_store,
1298            &mut bob_store.identity_store,
1299            now,
1300            rng,
1301        )
1302        .now_or_never()
1303        .expect("sync")
1304        .expect("bob reply enc");
1305        let bob_reply_signal_msg = match bob_reply_ct {
1306            CiphertextMessage::SignalMessage(m) => m,
1307            other => panic!(
1308                "expected bob reply SignalMessage, got {:?}",
1309                other.message_type()
1310            ),
1311        };
1312        message_decrypt_signal(
1313            &bob_reply_signal_msg,
1314            &bob_address,
1315            &alice_address,
1316            &mut alice_store.session_store,
1317            &mut alice_store.identity_store,
1318            rng,
1319        )
1320        .now_or_never()
1321        .expect("sync")
1322        .expect("bob reply dec");
1323
1324        let alice_new_chain_ct = message_encrypt(
1325            b"alice new chain",
1326            &bob_address,
1327            &alice_address,
1328            &mut alice_store.session_store,
1329            &mut alice_store.identity_store,
1330            now,
1331            rng,
1332        )
1333        .now_or_never()
1334        .expect("sync")
1335        .expect("alice new chain enc");
1336        let alice_new_chain_signal_msg = match alice_new_chain_ct {
1337            CiphertextMessage::SignalMessage(m) => m,
1338            other => panic!(
1339                "expected alice new chain SignalMessage, got {:?}",
1340                other.message_type()
1341            ),
1342        };
1343        message_decrypt_signal(
1344            &alice_new_chain_signal_msg,
1345            &alice_address,
1346            &bob_address,
1347            &mut bob_store.session_store,
1348            &mut bob_store.identity_store,
1349            rng,
1350        )
1351        .now_or_never()
1352        .expect("sync")
1353        .expect("alice new chain dec");
1354
1355        (
1356            alice_store,
1357            bob_store,
1358            alice_address,
1359            bob_address,
1360            delayed_signal_msg,
1361        )
1362    }
1363
1364    #[test]
1365    fn decrypt_ignores_corrupt_unmatched_receiver_chain() {
1366        let mut rng = ChaCha8Rng::seed_from_u64(0xD311A9);
1367        let (_alice_store, mut bob_store, alice_address, bob_address, delayed_signal_msg) =
1368            setup_two_alice_receiver_chains_on_bob(&mut rng);
1369
1370        let bob_record = bob_store
1371            .session_store
1372            .load_session(&alice_address)
1373            .now_or_never()
1374            .expect("sync")
1375            .expect("load succeeded")
1376            .expect("session exists");
1377        let serialized = bob_record.serialize().expect("serialize");
1378        let mut record_pb = RecordStructure::decode(serialized.as_slice()).expect("decode record");
1379        let current_session = record_pb.current_session.as_mut().expect("current session");
1380        assert!(
1381            current_session.receiver_chains.len() >= 2,
1382            "expected at least two receiver chains"
1383        );
1384
1385        let matched_key = delayed_signal_msg.sender_ratchet_key().serialize().to_vec();
1386        let matched_idx = current_session
1387            .receiver_chains
1388            .iter()
1389            .position(|chain| chain.sender_ratchet_key == matched_key)
1390            .expect("matching receiver chain present");
1391        let unmatched_idx = (0..current_session.receiver_chains.len())
1392            .find(|idx| *idx != matched_idx)
1393            .expect("unmatched receiver chain present");
1394
1395        current_session.receiver_chains[unmatched_idx].sender_ratchet_key = vec![0xFF];
1396
1397        let corrupted_record = SessionRecord::deserialize(record_pb.encode_to_vec().as_slice())
1398            .expect("deserialize mutated record");
1399        bob_store
1400            .session_store
1401            .store_session(&alice_address, &corrupted_record)
1402            .now_or_never()
1403            .expect("sync")
1404            .expect("store succeeded");
1405
1406        let ptext = message_decrypt_signal(
1407            &delayed_signal_msg,
1408            &alice_address,
1409            &bob_address,
1410            &mut bob_store.session_store,
1411            &mut bob_store.identity_store,
1412            &mut rng,
1413        )
1414        .now_or_never()
1415        .expect("sync")
1416        .expect("decrypt should ignore unmatched corrupt receiver chain");
1417
1418        assert_eq!(ptext, b"delayed old");
1419    }
1420
1421    #[test]
1422    fn decrypt_fails_on_corrupt_matched_receiver_chain() {
1423        let mut rng = ChaCha8Rng::seed_from_u64(0xD311AA);
1424        let (_alice_store, mut bob_store, alice_address, bob_address, delayed_signal_msg) =
1425            setup_two_alice_receiver_chains_on_bob(&mut rng);
1426
1427        let bob_record = bob_store
1428            .session_store
1429            .load_session(&alice_address)
1430            .now_or_never()
1431            .expect("sync")
1432            .expect("load succeeded")
1433            .expect("session exists");
1434        let serialized = bob_record.serialize().expect("serialize");
1435        let mut record_pb = RecordStructure::decode(serialized.as_slice()).expect("decode record");
1436        let current_session = record_pb.current_session.as_mut().expect("current session");
1437
1438        let matched_key = delayed_signal_msg.sender_ratchet_key().serialize().to_vec();
1439        let matched_idx = current_session
1440            .receiver_chains
1441            .iter()
1442            .position(|chain| chain.sender_ratchet_key == matched_key)
1443            .expect("matching receiver chain present");
1444        current_session.receiver_chains[matched_idx]
1445            .chain_key
1446            .as_mut()
1447            .expect("chain key present")
1448            .key = vec![0xFF];
1449
1450        let corrupted_record = SessionRecord::deserialize(record_pb.encode_to_vec().as_slice())
1451            .expect("deserialize mutated record");
1452        bob_store
1453            .session_store
1454            .store_session(&alice_address, &corrupted_record)
1455            .now_or_never()
1456            .expect("sync")
1457            .expect("store succeeded");
1458
1459        let err = message_decrypt_signal(
1460            &delayed_signal_msg,
1461            &alice_address,
1462            &bob_address,
1463            &mut bob_store.session_store,
1464            &mut bob_store.identity_store,
1465            &mut rng,
1466        )
1467        .now_or_never()
1468        .expect("sync")
1469        .expect_err("decrypt should fail on corrupt matched receiver chain");
1470
1471        assert!(
1472            matches!(
1473                err,
1474                SignalProtocolError::InvalidMessage(
1475                    CiphertextMessageType::Whisper,
1476                    "decryption failed"
1477                )
1478            ),
1479            "unexpected error: {err:?}"
1480        );
1481    }
1482
1483    // ── Dual-path simulation helpers ────────────────────────────────────
1484    //
1485    // Run every operation on both the refactored and legacy code paths,
1486    // asserting identical outputs (ciphertexts, plaintexts, or error
1487    // variants).  RNG sync follows the same clone-before-each-op pattern
1488    // as proptest_ciphertext_equality.
1489
1490    fn assert_store_state_equivalent(
1491        new_store: &InMemSignalProtocolStore,
1492        leg_store: &InMemSignalProtocolStore,
1493        peer_addr: &ProtocolAddress,
1494        context: &str,
1495    ) {
1496        let new_session = new_store
1497            .session_store
1498            .load_session(peer_addr)
1499            .now_or_never()
1500            .expect("sync")
1501            .expect("new load session");
1502        let leg_session = leg_store
1503            .session_store
1504            .load_session(peer_addr)
1505            .now_or_never()
1506            .expect("sync")
1507            .expect("legacy load session");
1508
1509        let new_session_bytes = new_session.map(|record| record.serialize().expect("serialize"));
1510        let leg_session_bytes = leg_session.map(|record| record.serialize().expect("serialize"));
1511        assert_eq!(
1512            new_session_bytes, leg_session_bytes,
1513            "{context}: session records diverged"
1514        );
1515
1516        let new_identity = new_store
1517            .identity_store
1518            .get_identity(peer_addr)
1519            .now_or_never()
1520            .expect("sync")
1521            .expect("new load identity")
1522            .map(|identity| identity.serialize());
1523        let leg_identity = leg_store
1524            .identity_store
1525            .get_identity(peer_addr)
1526            .now_or_never()
1527            .expect("sync")
1528            .expect("legacy load identity")
1529            .map(|identity| identity.serialize());
1530        assert_eq!(
1531            new_identity, leg_identity,
1532            "{context}: trusted identities diverged"
1533        );
1534    }
1535
1536    /// Encrypt on both paths with cloned RNG. Assert ciphertexts and sender
1537    /// state are byte-identical.
1538    fn dual_encrypt(
1539        plaintext: &[u8],
1540        recv_addr: &ProtocolAddress,
1541        send_addr: &ProtocolAddress,
1542        new_sender: &mut InMemSignalProtocolStore,
1543        leg_sender: &mut InMemSignalProtocolStore,
1544        now: SystemTime,
1545        rng: &mut ChaCha8Rng,
1546    ) -> SignalMessage {
1547        let mut leg_rng = rng.clone();
1548
1549        let new_ct = message_encrypt(
1550            plaintext,
1551            recv_addr,
1552            send_addr,
1553            &mut new_sender.session_store,
1554            &mut new_sender.identity_store,
1555            now,
1556            rng,
1557        )
1558        .now_or_never()
1559        .expect("sync")
1560        .expect("new encrypt");
1561
1562        let leg_ct = legacy::legacy_message_encrypt(
1563            plaintext,
1564            recv_addr,
1565            send_addr,
1566            &mut leg_sender.session_store,
1567            &mut leg_sender.identity_store,
1568            now,
1569            &mut leg_rng,
1570        )
1571        .now_or_never()
1572        .expect("sync")
1573        .expect("legacy encrypt");
1574
1575        assert_eq!(
1576            new_ct.serialize(),
1577            leg_ct.serialize(),
1578            "new and legacy produced different ciphertexts"
1579        );
1580        assert_eq!(
1581            new_ct.message_type(),
1582            leg_ct.message_type(),
1583            "new and legacy produced different ciphertext types"
1584        );
1585        assert_store_state_equivalent(new_sender, leg_sender, recv_addr, "encrypt");
1586
1587        match new_ct {
1588            CiphertextMessage::SignalMessage(m) => m,
1589            other => panic!(
1590                "expected SignalMessage from dual_encrypt, got {:?}",
1591                other.message_type()
1592            ),
1593        }
1594    }
1595
1596    /// Encrypt on both paths with cloned RNG. Assert full ciphertext
1597    /// equivalence, including `PreKeySignalMessage`.
1598    fn dual_encrypt_any(
1599        plaintext: &[u8],
1600        recv_addr: &ProtocolAddress,
1601        send_addr: &ProtocolAddress,
1602        new_sender: &mut InMemSignalProtocolStore,
1603        leg_sender: &mut InMemSignalProtocolStore,
1604        now: SystemTime,
1605        rng: &mut ChaCha8Rng,
1606    ) -> CiphertextMessage {
1607        let mut leg_rng = rng.clone();
1608
1609        let new_ct = message_encrypt(
1610            plaintext,
1611            recv_addr,
1612            send_addr,
1613            &mut new_sender.session_store,
1614            &mut new_sender.identity_store,
1615            now,
1616            rng,
1617        )
1618        .now_or_never()
1619        .expect("sync")
1620        .expect("new encrypt");
1621
1622        let leg_ct = legacy::legacy_message_encrypt(
1623            plaintext,
1624            recv_addr,
1625            send_addr,
1626            &mut leg_sender.session_store,
1627            &mut leg_sender.identity_store,
1628            now,
1629            &mut leg_rng,
1630        )
1631        .now_or_never()
1632        .expect("sync")
1633        .expect("legacy encrypt");
1634
1635        assert_eq!(
1636            new_ct.serialize(),
1637            leg_ct.serialize(),
1638            "new and legacy produced different ciphertexts"
1639        );
1640        assert_eq!(
1641            new_ct.message_type(),
1642            leg_ct.message_type(),
1643            "new and legacy produced different ciphertext types"
1644        );
1645        assert_store_state_equivalent(new_sender, leg_sender, recv_addr, "encrypt");
1646        new_ct
1647    }
1648
1649    /// Decrypt on both paths with cloned RNG. Assert plaintexts and receiver
1650    /// state match.
1651    fn dual_decrypt(
1652        msg: &SignalMessage,
1653        sender_addr: &ProtocolAddress,
1654        recv_addr: &ProtocolAddress,
1655        new_receiver: &mut InMemSignalProtocolStore,
1656        leg_receiver: &mut InMemSignalProtocolStore,
1657        rng: &mut ChaCha8Rng,
1658    ) -> Vec<u8> {
1659        let mut leg_rng = rng.clone();
1660
1661        let new_pt = message_decrypt_signal(
1662            msg,
1663            sender_addr,
1664            recv_addr,
1665            &mut new_receiver.session_store,
1666            &mut new_receiver.identity_store,
1667            rng,
1668        )
1669        .now_or_never()
1670        .expect("sync")
1671        .expect("new decrypt");
1672
1673        let leg_pt = legacy::legacy_message_decrypt_signal(
1674            msg,
1675            sender_addr,
1676            &mut leg_receiver.session_store,
1677            &mut leg_receiver.identity_store,
1678            &mut leg_rng,
1679        )
1680        .now_or_never()
1681        .expect("sync")
1682        .expect("legacy decrypt");
1683
1684        assert_eq!(
1685            new_pt, leg_pt,
1686            "new and legacy produced different plaintexts"
1687        );
1688        assert_store_state_equivalent(new_receiver, leg_receiver, sender_addr, "decrypt");
1689        new_pt
1690    }
1691
1692    /// Decrypt any ciphertext on both paths with cloned RNG. Assert
1693    /// plaintexts and receiver-side state match.
1694    fn dual_decrypt_any(
1695        msg: &CiphertextMessage,
1696        sender_addr: &ProtocolAddress,
1697        receiver_addr: &ProtocolAddress,
1698        new_receiver: &mut InMemSignalProtocolStore,
1699        leg_receiver: &mut InMemSignalProtocolStore,
1700        rng: &mut ChaCha8Rng,
1701    ) -> Vec<u8> {
1702        let mut leg_rng = rng.clone();
1703
1704        let new_pt = message_decrypt(
1705            msg,
1706            sender_addr,
1707            receiver_addr,
1708            &mut new_receiver.session_store,
1709            &mut new_receiver.identity_store,
1710            &mut new_receiver.pre_key_store,
1711            &new_receiver.signed_pre_key_store,
1712            &mut new_receiver.kyber_pre_key_store,
1713            rng,
1714        )
1715        .now_or_never()
1716        .expect("sync")
1717        .expect("new decrypt");
1718
1719        let leg_pt = legacy::legacy_message_decrypt(
1720            msg,
1721            sender_addr,
1722            receiver_addr,
1723            &mut leg_receiver.session_store,
1724            &mut leg_receiver.identity_store,
1725            &mut leg_receiver.pre_key_store,
1726            &leg_receiver.signed_pre_key_store,
1727            &mut leg_receiver.kyber_pre_key_store,
1728            &mut leg_rng,
1729        )
1730        .now_or_never()
1731        .expect("sync")
1732        .expect("legacy decrypt");
1733
1734        assert_eq!(
1735            new_pt, leg_pt,
1736            "new and legacy produced different plaintexts"
1737        );
1738        assert_store_state_equivalent(new_receiver, leg_receiver, sender_addr, "decrypt");
1739        new_pt
1740    }
1741
1742    fn dual_decrypt_any_result(
1743        msg: &CiphertextMessage,
1744        sender_addr: &ProtocolAddress,
1745        receiver_addr: &ProtocolAddress,
1746        new_receiver: &mut InMemSignalProtocolStore,
1747        leg_receiver: &mut InMemSignalProtocolStore,
1748        rng: &mut ChaCha8Rng,
1749    ) -> Result<Vec<u8>> {
1750        let mut leg_rng = rng.clone();
1751
1752        let new_result = message_decrypt(
1753            msg,
1754            sender_addr,
1755            receiver_addr,
1756            &mut new_receiver.session_store,
1757            &mut new_receiver.identity_store,
1758            &mut new_receiver.pre_key_store,
1759            &new_receiver.signed_pre_key_store,
1760            &mut new_receiver.kyber_pre_key_store,
1761            rng,
1762        )
1763        .now_or_never()
1764        .expect("sync");
1765
1766        let leg_result = legacy::legacy_message_decrypt(
1767            msg,
1768            sender_addr,
1769            receiver_addr,
1770            &mut leg_receiver.session_store,
1771            &mut leg_receiver.identity_store,
1772            &mut leg_receiver.pre_key_store,
1773            &leg_receiver.signed_pre_key_store,
1774            &mut leg_receiver.kyber_pre_key_store,
1775            &mut leg_rng,
1776        )
1777        .now_or_never()
1778        .expect("sync");
1779
1780        match (new_result, leg_result) {
1781            (Ok(new_pt), Ok(leg_pt)) => {
1782                assert_eq!(
1783                    new_pt, leg_pt,
1784                    "new and legacy produced different plaintexts"
1785                );
1786                assert_store_state_equivalent(new_receiver, leg_receiver, sender_addr, "decrypt");
1787                Ok(new_pt)
1788            }
1789            (Err(new_err), Err(leg_err)) => {
1790                assert_eq!(
1791                    std::mem::discriminant(&new_err),
1792                    std::mem::discriminant(&leg_err),
1793                    "error variants differ: new={new_err:?}, legacy={leg_err:?}"
1794                );
1795                assert_store_state_equivalent(new_receiver, leg_receiver, sender_addr, "decrypt");
1796                Err(new_err)
1797            }
1798            (new_result, leg_result) => panic!(
1799                "new and legacy disagreed on decrypt result: new={new_result:?}, legacy={leg_result:?}"
1800            ),
1801        }
1802    }
1803
1804    /// Decrypt on both paths, assert both fail with the same error variant.
1805    fn dual_decrypt_expect_err(
1806        msg: &SignalMessage,
1807        sender_addr: &ProtocolAddress,
1808        recv_addr: &ProtocolAddress,
1809        new_receiver: &mut InMemSignalProtocolStore,
1810        leg_receiver: &mut InMemSignalProtocolStore,
1811        rng: &mut ChaCha8Rng,
1812    ) -> SignalProtocolError {
1813        let mut leg_rng = rng.clone();
1814
1815        let new_err = message_decrypt_signal(
1816            msg,
1817            sender_addr,
1818            recv_addr,
1819            &mut new_receiver.session_store,
1820            &mut new_receiver.identity_store,
1821            rng,
1822        )
1823        .now_or_never()
1824        .expect("sync")
1825        .expect_err("expected new decrypt to fail");
1826
1827        let leg_err = legacy::legacy_message_decrypt_signal(
1828            msg,
1829            sender_addr,
1830            &mut leg_receiver.session_store,
1831            &mut leg_receiver.identity_store,
1832            &mut leg_rng,
1833        )
1834        .now_or_never()
1835        .expect("sync")
1836        .expect_err("expected legacy decrypt to fail");
1837
1838        assert_eq!(
1839            std::mem::discriminant(&new_err),
1840            std::mem::discriminant(&leg_err),
1841            "error variants differ: new={new_err:?}, legacy={leg_err:?}"
1842        );
1843        new_err
1844    }
1845
1846    // ── DualSession convenience wrapper ─────────────────────────────────
1847
1848    /// Paired new+legacy session state for readable scenario tests.
1849    struct DualSession {
1850        na: InMemSignalProtocolStore,
1851        nb: InMemSignalProtocolStore,
1852        la: InMemSignalProtocolStore,
1853        lb: InMemSignalProtocolStore,
1854        alice: ProtocolAddress,
1855        bob: ProtocolAddress,
1856        rng: ChaCha8Rng,
1857        now: SystemTime,
1858    }
1859
1860    impl DualSession {
1861        fn new(seed: u64) -> Self {
1862            let mut rng = ChaCha8Rng::seed_from_u64(seed);
1863            let (na, nb, alice, bob) = setup_stores(&mut rng);
1864            let (la, lb) = (na.clone(), nb.clone());
1865            Self {
1866                na,
1867                nb,
1868                la,
1869                lb,
1870                alice,
1871                bob,
1872                rng,
1873                now: SystemTime::now(),
1874            }
1875        }
1876
1877        fn alice_sends(&mut self, plaintext: &[u8]) -> SignalMessage {
1878            dual_encrypt(
1879                plaintext,
1880                &self.bob,
1881                &self.alice,
1882                &mut self.na,
1883                &mut self.la,
1884                self.now,
1885                &mut self.rng,
1886            )
1887        }
1888
1889        fn bob_sends(&mut self, plaintext: &[u8]) -> SignalMessage {
1890            dual_encrypt(
1891                plaintext,
1892                &self.alice,
1893                &self.bob,
1894                &mut self.nb,
1895                &mut self.lb,
1896                self.now,
1897                &mut self.rng,
1898            )
1899        }
1900
1901        fn bob_receives(&mut self, msg: &SignalMessage) -> Vec<u8> {
1902            dual_decrypt(
1903                msg,
1904                &self.alice,
1905                &self.bob,
1906                &mut self.nb,
1907                &mut self.lb,
1908                &mut self.rng,
1909            )
1910        }
1911
1912        fn alice_receives(&mut self, msg: &SignalMessage) -> Vec<u8> {
1913            dual_decrypt(
1914                msg,
1915                &self.bob,
1916                &self.alice,
1917                &mut self.na,
1918                &mut self.la,
1919                &mut self.rng,
1920            )
1921        }
1922
1923        fn bob_receives_err(&mut self, msg: &SignalMessage) -> SignalProtocolError {
1924            dual_decrypt_expect_err(
1925                msg,
1926                &self.alice,
1927                &self.bob,
1928                &mut self.nb,
1929                &mut self.lb,
1930                &mut self.rng,
1931            )
1932        }
1933
1934        #[allow(dead_code)]
1935        fn alice_receives_err(&mut self, msg: &SignalMessage) -> SignalProtocolError {
1936            dual_decrypt_expect_err(
1937                msg,
1938                &self.bob,
1939                &self.alice,
1940                &mut self.na,
1941                &mut self.la,
1942                &mut self.rng,
1943            )
1944        }
1945    }
1946
1947    impl DualParticipant {
1948        fn new(
1949            _name: &'static str,
1950            address: ProtocolAddress,
1951            rng: &mut (impl rand::Rng + rand::CryptoRng),
1952        ) -> Self {
1953            let identity = IdentityKeyPair::generate(rng);
1954            let store = InMemSignalProtocolStore::new(identity, rng.random()).unwrap();
1955            Self {
1956                address,
1957                message_queue: Vec::new(),
1958                state: DualLocalState {
1959                    new_store: store.clone(),
1960                    legacy_store: store,
1961                    pre_key_count: 0,
1962                },
1963                snapshots: Vec::new(),
1964                message_send_log: Vec::new(),
1965            }
1966        }
1967
1968        fn address(&self) -> &ProtocolAddress {
1969            &self.address
1970        }
1971
1972        fn has_pending_incoming_messages(&self) -> bool {
1973            !self.message_queue.is_empty()
1974        }
1975
1976        fn assert_equivalent_with(&self, them: &Self, context: &str) {
1977            assert_store_state_equivalent(
1978                &self.state.new_store,
1979                &self.state.legacy_store,
1980                &them.address,
1981                context,
1982            );
1983        }
1984
1985        async fn process_pre_key(
1986            &mut self,
1987            them: &mut Self,
1988            use_one_time_pre_key: bool,
1989            rng: &mut ChaCha8Rng,
1990        ) {
1991            let their_signed_pre_key_pair = KeyPair::generate(rng);
1992            let their_signed_pre_key_public = their_signed_pre_key_pair.public_key.serialize();
1993            let identity_key_pair = them.state.new_store.get_identity_key_pair().await.unwrap();
1994            let their_signed_pre_key_signature = identity_key_pair
1995                .private_key()
1996                .calculate_signature(&their_signed_pre_key_public, rng)
1997                .unwrap();
1998
1999            them.state.pre_key_count += 1;
2000            let signed_pre_key_id: SignedPreKeyId = them.state.pre_key_count.into();
2001            let signed_pre_key_record = SignedPreKeyRecord::new(
2002                signed_pre_key_id,
2003                Timestamp::from_epoch_millis(42),
2004                &their_signed_pre_key_pair,
2005                &their_signed_pre_key_signature,
2006            );
2007            them.state
2008                .new_store
2009                .save_signed_pre_key(signed_pre_key_id, &signed_pre_key_record)
2010                .await
2011                .unwrap();
2012            them.state
2013                .legacy_store
2014                .save_signed_pre_key(signed_pre_key_id, &signed_pre_key_record)
2015                .await
2016                .unwrap();
2017
2018            them.state.pre_key_count += 1;
2019            let pre_key_id: PreKeyId = them.state.pre_key_count.into();
2020            let pre_key_info = if use_one_time_pre_key {
2021                let one_time_pre_key = KeyPair::generate(rng);
2022                let pre_key_record = PreKeyRecord::new(pre_key_id, &one_time_pre_key);
2023                them.state
2024                    .new_store
2025                    .save_pre_key(pre_key_id, &pre_key_record)
2026                    .await
2027                    .unwrap();
2028                them.state
2029                    .legacy_store
2030                    .save_pre_key(pre_key_id, &pre_key_record)
2031                    .await
2032                    .unwrap();
2033                Some((pre_key_id, one_time_pre_key.public_key))
2034            } else {
2035                None
2036            };
2037
2038            let their_kyber_pre_key_pair =
2039                crate::kem::KeyPair::generate(crate::kem::KeyType::Kyber1024, rng);
2040            let their_kyber_pre_key_public = their_kyber_pre_key_pair.public_key.serialize();
2041            let their_kyber_pre_key_signature = identity_key_pair
2042                .private_key()
2043                .calculate_signature(&their_kyber_pre_key_public, rng)
2044                .unwrap();
2045
2046            them.state.pre_key_count += 1;
2047            let kyber_pre_key_id: KyberPreKeyId = them.state.pre_key_count.into();
2048            let kyber_pre_key_record = KyberPreKeyRecord::new(
2049                kyber_pre_key_id,
2050                Timestamp::from_epoch_millis(42),
2051                &their_kyber_pre_key_pair,
2052                &their_kyber_pre_key_signature,
2053            );
2054            them.state
2055                .new_store
2056                .save_kyber_pre_key(kyber_pre_key_id, &kyber_pre_key_record)
2057                .await
2058                .unwrap();
2059            them.state
2060                .legacy_store
2061                .save_kyber_pre_key(kyber_pre_key_id, &kyber_pre_key_record)
2062                .await
2063                .unwrap();
2064
2065            let their_pre_key_bundle = PreKeyBundle::new(
2066                them.state
2067                    .new_store
2068                    .get_local_registration_id()
2069                    .await
2070                    .unwrap(),
2071                DeviceId::new(1).unwrap(),
2072                pre_key_info,
2073                signed_pre_key_id,
2074                their_signed_pre_key_pair.public_key,
2075                their_signed_pre_key_signature.into_vec(),
2076                kyber_pre_key_id,
2077                their_kyber_pre_key_pair.public_key,
2078                their_kyber_pre_key_signature.into_vec(),
2079                *identity_key_pair.identity_key(),
2080            )
2081            .unwrap();
2082
2083            let mut legacy_rng = rng.clone();
2084            process_prekey_bundle(
2085                &them.address,
2086                &self.address,
2087                &mut self.state.new_store.session_store,
2088                &mut self.state.new_store.identity_store,
2089                &their_pre_key_bundle,
2090                SystemTime::UNIX_EPOCH,
2091                rng,
2092            )
2093            .await
2094            .unwrap();
2095            process_prekey_bundle(
2096                &them.address,
2097                &self.address,
2098                &mut self.state.legacy_store.session_store,
2099                &mut self.state.legacy_store.identity_store,
2100                &their_pre_key_bundle,
2101                SystemTime::UNIX_EPOCH,
2102                &mut legacy_rng,
2103            )
2104            .await
2105            .unwrap();
2106
2107            self.assert_equivalent_with(them, "process_pre_key/self");
2108            them.assert_equivalent_with(self, "process_pre_key/them");
2109            assert!(
2110                self.state
2111                    .new_store
2112                    .load_session(&them.address)
2113                    .await
2114                    .unwrap()
2115                    .expect("just created")
2116                    .has_usable_sender_chain(
2117                        SystemTime::UNIX_EPOCH,
2118                        SessionUsabilityRequirements::all(),
2119                    )
2120                    .unwrap()
2121            );
2122        }
2123
2124        async fn send_message(&mut self, them: &mut Self, rng: &mut ChaCha8Rng) {
2125            self.send_message_with_id(them, self.message_send_log.len().try_into().unwrap(), rng)
2126                .await;
2127            self.message_send_log.push(MessageStatus::Sent);
2128        }
2129
2130        async fn send_message_with_id(&mut self, them: &mut Self, id: u64, rng: &mut ChaCha8Rng) {
2131            let has_usable_sender_chain = self
2132                .state
2133                .new_store
2134                .load_session(&them.address)
2135                .await
2136                .unwrap()
2137                .and_then(|session| {
2138                    session
2139                        .has_usable_sender_chain(
2140                            SystemTime::UNIX_EPOCH,
2141                            SessionUsabilityRequirements::all(),
2142                        )
2143                        .ok()
2144                })
2145                .unwrap_or(false);
2146
2147            if !has_usable_sender_chain {
2148                self.process_pre_key(them, rng.random_bool(0.75), rng).await;
2149            }
2150
2151            let buffer = id.to_le_bytes();
2152            let outgoing_message = dual_encrypt_any(
2153                &buffer,
2154                &them.address,
2155                &self.address,
2156                &mut self.state.new_store,
2157                &mut self.state.legacy_store,
2158                SystemTime::UNIX_EPOCH,
2159                rng,
2160            );
2161
2162            let incoming_message = match outgoing_message.message_type() {
2163                CiphertextMessageType::PreKey => CiphertextMessage::PreKeySignalMessage(
2164                    PreKeySignalMessage::try_from(outgoing_message.serialize()).unwrap(),
2165                ),
2166                CiphertextMessageType::Whisper => CiphertextMessage::SignalMessage(
2167                    SignalMessage::try_from(outgoing_message.serialize()).unwrap(),
2168                ),
2169                other_type => panic!("unexpected type {other_type:?}"),
2170            };
2171
2172            them.message_queue.push((incoming_message, id));
2173            self.assert_equivalent_with(them, "send");
2174        }
2175
2176        async fn receive_messages(&mut self, them: &mut Self, rng: &mut ChaCha8Rng) {
2177            for (incoming_message, expected) in self.message_queue.split_off(0) {
2178                match incoming_message {
2179                    CiphertextMessage::SignalMessage(_)
2180                    | CiphertextMessage::PreKeySignalMessage(_) => {
2181                        match dual_decrypt_any_result(
2182                            &incoming_message,
2183                            &them.address,
2184                            &self.address,
2185                            &mut self.state.new_store,
2186                            &mut self.state.legacy_store,
2187                            rng,
2188                        ) {
2189                            Ok(decrypted) => {
2190                                assert_eq!(expected.to_le_bytes(), &decrypted[..]);
2191                                them.ack(expected);
2192                            }
2193                            Err(_) => {
2194                                let error_msg = DecryptionErrorMessage::for_original(
2195                                    incoming_message.serialize(),
2196                                    incoming_message.message_type(),
2197                                    Timestamp::from_epoch_millis(expected),
2198                                    1,
2199                                )
2200                                .expect("can encode DEM");
2201                                them.message_queue.push((
2202                                    CiphertextMessage::PlaintextContent(error_msg.into()),
2203                                    u64::MAX,
2204                                ));
2205                            }
2206                        }
2207                    }
2208                    CiphertextMessage::SenderKeyMessage(_) => unreachable!(),
2209                    CiphertextMessage::PlaintextContent(content) => {
2210                        self.handle_decryption_error(them, content, rng).await;
2211                    }
2212                }
2213            }
2214            self.assert_equivalent_with(them, "receive");
2215            them.assert_equivalent_with(self, "receive/peer");
2216        }
2217
2218        fn drop_message(&mut self, them: &mut Self) {
2219            match self.message_queue.pop() {
2220                None | Some((CiphertextMessage::PlaintextContent(_), _)) => {}
2221                Some((_, id)) => them.nack(id),
2222            }
2223        }
2224
2225        fn shuffle_messages(&mut self, rng: &mut impl rand::Rng) {
2226            use rand::seq::SliceRandom as _;
2227            self.message_queue.shuffle(rng);
2228        }
2229
2230        async fn handle_decryption_error(
2231            &mut self,
2232            them: &mut Self,
2233            content: PlaintextContent,
2234            rng: &mut ChaCha8Rng,
2235        ) {
2236            let dem = extract_decryption_error_message_from_serialized_content(content.body())
2237                .expect("all PlaintextContent is DEM");
2238            assert_eq!(dem.device_id(), 1);
2239
2240            let id = dem.timestamp().epoch_millis();
2241            let Some(status) = self.message_send_log.get(usize::try_from(id).unwrap()) else {
2242                panic!(
2243                    "failed to decrypt an unsent message {id} ({} total sent)",
2244                    self.message_send_log.len()
2245                )
2246            };
2247            match status {
2248                MessageStatus::Sent => {}
2249                MessageStatus::Dropped => {
2250                    panic!("got a decryption error for dropped message {id}");
2251                }
2252                MessageStatus::Delivered => {
2253                    panic!("got a decryption error for successfully delivered message {id}");
2254                }
2255            }
2256
2257            let ratchet_key = dem
2258                .ratchet_key()
2259                .expect("all DEMs for 1:1 messages have ratchet keys");
2260            if self
2261                .state
2262                .new_store
2263                .load_session(&them.address)
2264                .await
2265                .unwrap()
2266                .is_some_and(|session| {
2267                    session
2268                        .current_ratchet_key_matches(ratchet_key)
2269                        .expect("structurally valid session")
2270                })
2271            {
2272                self.archive_session(&them.address).await;
2273            }
2274
2275            self.send_message_with_id(them, id, rng).await;
2276        }
2277
2278        async fn archive_session(&mut self, their_address: &ProtocolAddress) {
2279            for store in [&mut self.state.new_store, &mut self.state.legacy_store] {
2280                if let Some(mut session) = store.load_session(their_address).await.unwrap() {
2281                    session.archive_current_state().unwrap();
2282                    store.store_session(their_address, &session).await.unwrap();
2283                }
2284            }
2285        }
2286
2287        fn snapshot_state(&mut self) {
2288            self.snapshots.push(self.state.clone());
2289        }
2290
2291        fn restore_from_snapshot_if_exists(&mut self, i: u8) {
2292            let i = usize::from(i);
2293            if i < self.snapshots.len() {
2294                self.state = self.snapshots.remove(i);
2295            }
2296        }
2297
2298        fn ack(&mut self, id: u64) {
2299            self.update_status(id, MessageStatus::Delivered);
2300        }
2301
2302        fn nack(&mut self, id: u64) {
2303            self.update_status(id, MessageStatus::Dropped);
2304        }
2305
2306        fn update_status(&mut self, id: u64, updated_status: MessageStatus) {
2307            let Some(status) = self.message_send_log.get_mut(usize::try_from(id).unwrap()) else {
2308                panic!(
2309                    "tried to update unsent message {id} ({} total sent)",
2310                    self.message_send_log.len()
2311                )
2312            };
2313            match status {
2314                MessageStatus::Sent => *status = updated_status,
2315                MessageStatus::Dropped => panic!("updated dropped message {id}"),
2316                MessageStatus::Delivered => panic!("updated delivered message {id}"),
2317            }
2318        }
2319
2320        async fn run_event(&mut self, them: &mut Self, event: Event, rng: &mut ChaCha8Rng) {
2321            match event {
2322                Event::Archive => self.archive_session(them.address()).await,
2323                Event::Snapshot => self.snapshot_state(),
2324                Event::Restore { index } => self.restore_from_snapshot_if_exists(index),
2325                Event::Receive => self.receive_messages(them, rng).await,
2326                Event::Drop => self.drop_message(them),
2327                Event::Shuffle => self.shuffle_messages(rng),
2328                Event::Send { count_times_eight } => {
2329                    for _ in 0..(count_times_eight / 8) {
2330                        self.send_message(them, rng).await;
2331                    }
2332                }
2333            }
2334            self.assert_equivalent_with(them, "event");
2335            them.assert_equivalent_with(self, "event/peer");
2336        }
2337    }
2338
2339    // ── Scenario tests ──────────────────────────────────────────────────
2340
2341    /// Ordinary skipped-key handling remains interoperable when messages are
2342    /// delivered with gaps, later arrive out of order, and both directions
2343    /// continue sending before the session fully catches up.
2344    #[test]
2345    fn scenario_interleaved_delivery_with_gaps_and_recovery() {
2346        let mut s = DualSession::new(0xBEEF_0001);
2347
2348        // Alice sends a burst to Bob. Bob receives only the first and third,
2349        // leaving a gap that must be recovered later from stored skipped keys.
2350        let a_msgs: Vec<_> = (0u8..4)
2351            .map(|i| (s.alice_sends(&[b'A', i]), vec![b'A', i]))
2352            .collect();
2353        assert_eq!(s.bob_receives(&a_msgs[0].0), a_msgs[0].1, "alice msg 0");
2354        assert_eq!(s.bob_receives(&a_msgs[2].0), a_msgs[2].1, "alice msg 2");
2355
2356        // Before Alice's burst is fully drained, Bob sends his own burst.
2357        // Alice receives only the later message first, exercising the same
2358        // skipped-key path in the opposite direction.
2359        let b_msgs: Vec<_> = (0u8..3)
2360            .map(|i| (s.bob_sends(&[b'B', i]), vec![b'B', i]))
2361            .collect();
2362        assert_eq!(s.alice_receives(&b_msgs[2].0), b_msgs[2].1, "bob msg 2");
2363
2364        // The missing earlier messages now arrive and must still decrypt.
2365        assert_eq!(s.bob_receives(&a_msgs[1].0), a_msgs[1].1, "alice msg 1");
2366        assert_eq!(s.bob_receives(&a_msgs[3].0), a_msgs[3].1, "alice msg 3");
2367        assert_eq!(s.alice_receives(&b_msgs[0].0), b_msgs[0].1, "bob msg 0");
2368        assert_eq!(s.alice_receives(&b_msgs[1].0), b_msgs[1].1, "bob msg 1");
2369
2370        // After recovering the gaps, both directions should continue in steady
2371        // state without any special handling.
2372        let alice_followup = s.alice_sends(b"alice steady");
2373        assert_eq!(s.bob_receives(&alice_followup), b"alice steady");
2374        let bob_followup = s.bob_sends(b"bob steady");
2375        assert_eq!(s.alice_receives(&bob_followup), b"bob steady");
2376    }
2377
2378    /// Skip past MAX_FORWARD_JUMPS — both paths must reject with the same
2379    /// error.  Encrypts on just the new path for performance (25k+ messages);
2380    /// both receivers start from identical untouched state.
2381    #[test]
2382    fn scenario_chain_jump_over_limit() {
2383        let mut rng = ChaCha8Rng::seed_from_u64(0xBEEF_0005);
2384        let (mut na, mut nb, alice, bob) = setup_stores(&mut rng);
2385        let mut lb = nb.clone();
2386        let now = SystemTime::now();
2387
2388        let count = crate::consts::MAX_FORWARD_JUMPS + 2;
2389        let mut last = None;
2390        for _ in 0..count {
2391            let ct = message_encrypt(
2392                b"x",
2393                &bob,
2394                &alice,
2395                &mut na.session_store,
2396                &mut na.identity_store,
2397                now,
2398                &mut rng,
2399            )
2400            .now_or_never()
2401            .expect("sync")
2402            .expect("encrypt");
2403            last = Some(match ct {
2404                CiphertextMessage::SignalMessage(m) => m,
2405                _ => panic!("not SignalMessage"),
2406            });
2407        }
2408        let msg = last.unwrap();
2409
2410        // Bob has received nothing — new and legacy stores are identical.
2411        let mut leg_rng = rng.clone();
2412        let new_err = message_decrypt_signal(
2413            &msg,
2414            &alice,
2415            &bob,
2416            &mut nb.session_store,
2417            &mut nb.identity_store,
2418            &mut rng,
2419        )
2420        .now_or_never()
2421        .expect("sync")
2422        .expect_err("should exceed jump limit");
2423
2424        let leg_err = legacy::legacy_message_decrypt_signal(
2425            &msg,
2426            &alice,
2427            &mut lb.session_store,
2428            &mut lb.identity_store,
2429            &mut leg_rng,
2430        )
2431        .now_or_never()
2432        .expect("sync")
2433        .expect_err("should exceed jump limit");
2434
2435        assert_eq!(
2436            std::mem::discriminant(&new_err),
2437            std::mem::discriminant(&leg_err),
2438            "error variants differ: new={new_err:?}, legacy={leg_err:?}"
2439        );
2440        assert!(
2441            matches!(new_err, SignalProtocolError::InvalidMessage(..)),
2442            "expected InvalidMessage, got {new_err:?}"
2443        );
2444    }
2445
2446    /// The initial unacknowledged send must be bit-identical as a
2447    /// `PreKeySignalMessage`, and both sides must end up with identical
2448    /// session state after the ack round-trip.
2449    #[test]
2450    fn scenario_prekey_session_establishment_equivalence() {
2451        let mut rng = ChaCha8Rng::seed_from_u64(0xBEEF_0008);
2452        let alice_identity = IdentityKeyPair::generate(&mut rng);
2453        let bob_identity = IdentityKeyPair::generate(&mut rng);
2454        let alice = ProtocolAddress::new(
2455            "9d0652a3-dcc3-4d11-975f-74d61598733f".to_owned(),
2456            DeviceId::new(1).unwrap(),
2457        );
2458        let bob = ProtocolAddress::new(
2459            "796abedb-ca4e-4f18-8803-1fde5b921f9f".to_owned(),
2460            DeviceId::new(1).unwrap(),
2461        );
2462        let now = SystemTime::now();
2463
2464        let alice_base = InMemSignalProtocolStore::new(alice_identity, 1).expect("alice store");
2465        let mut bob_base = InMemSignalProtocolStore::new(bob_identity, 2).expect("bob store");
2466        let bundle = create_bob_bundle(&mut bob_base, 1, 1, 1, &mut rng);
2467
2468        let (mut alice_new, mut alice_legacy) = (alice_base.clone(), alice_base.clone());
2469        let (mut bob_new, mut bob_legacy) = (bob_base.clone(), bob_base.clone());
2470
2471        let mut legacy_rng = rng.clone();
2472        process_prekey_bundle(
2473            &bob,
2474            &alice,
2475            &mut alice_new.session_store,
2476            &mut alice_new.identity_store,
2477            &bundle,
2478            now,
2479            &mut rng,
2480        )
2481        .now_or_never()
2482        .expect("sync")
2483        .expect("new process_prekey_bundle");
2484        process_prekey_bundle(
2485            &bob,
2486            &alice,
2487            &mut alice_legacy.session_store,
2488            &mut alice_legacy.identity_store,
2489            &bundle,
2490            now,
2491            &mut legacy_rng,
2492        )
2493        .now_or_never()
2494        .expect("sync")
2495        .expect("legacy process_prekey_bundle");
2496        assert_store_state_equivalent(&alice_new, &alice_legacy, &bob, "post-bundle");
2497
2498        let init = dual_encrypt_any(
2499            b"session init",
2500            &bob,
2501            &alice,
2502            &mut alice_new,
2503            &mut alice_legacy,
2504            now,
2505            &mut rng,
2506        );
2507        assert!(
2508            matches!(init, CiphertextMessage::PreKeySignalMessage(_)),
2509            "expected first message after bundle processing to be PreKey"
2510        );
2511
2512        assert_eq!(
2513            dual_decrypt_any(&init, &alice, &bob, &mut bob_new, &mut bob_legacy, &mut rng),
2514            b"session init"
2515        );
2516
2517        let ack = dual_encrypt_any(
2518            b"session ack",
2519            &alice,
2520            &bob,
2521            &mut bob_new,
2522            &mut bob_legacy,
2523            now,
2524            &mut rng,
2525        );
2526        assert!(
2527            matches!(ack, CiphertextMessage::SignalMessage(_)),
2528            "expected ack to be a SignalMessage"
2529        );
2530        assert_eq!(
2531            dual_decrypt_any(
2532                &ack,
2533                &bob,
2534                &alice,
2535                &mut alice_new,
2536                &mut alice_legacy,
2537                &mut rng
2538            ),
2539            b"session ack"
2540        );
2541
2542        let followup = dual_encrypt_any(
2543            b"steady state",
2544            &bob,
2545            &alice,
2546            &mut alice_new,
2547            &mut alice_legacy,
2548            now,
2549            &mut rng,
2550        );
2551        assert!(
2552            matches!(followup, CiphertextMessage::SignalMessage(_)),
2553            "expected acknowledged session to emit SignalMessage"
2554        );
2555        assert_eq!(
2556            dual_decrypt_any(
2557                &followup,
2558                &alice,
2559                &bob,
2560                &mut bob_new,
2561                &mut bob_legacy,
2562                &mut rng
2563            ),
2564            b"steady state"
2565        );
2566    }
2567
2568    /// Flip a byte in the MAC — both paths must reject identically, and
2569    /// the original message must still decrypt afterward.
2570    #[test]
2571    fn scenario_corrupted_ciphertext() {
2572        let mut s = DualSession::new(0xBEEF_0006);
2573
2574        // Warm up the session with a round-trip
2575        let msg = s.alice_sends(b"hello");
2576        assert_eq!(s.bob_receives(&msg), b"hello");
2577        let msg = s.bob_sends(b"hi");
2578        assert_eq!(s.alice_receives(&msg), b"hi");
2579
2580        // Alice sends a message; corrupt the last byte (in the MAC)
2581        let msg = s.alice_sends(b"secret");
2582        let mut corrupted_bytes = msg.serialized().to_vec();
2583        let len = corrupted_bytes.len();
2584        corrupted_bytes[len - 1] ^= 0xFF;
2585        let corrupted =
2586            SignalMessage::try_from(corrupted_bytes.as_slice()).expect("parse corrupted message");
2587
2588        let err = s.bob_receives_err(&corrupted);
2589        assert!(
2590            matches!(
2591                err,
2592                SignalProtocolError::InvalidMessage(CiphertextMessageType::Whisper, _)
2593            ),
2594            "expected InvalidMessage(Whisper, _), got {err:?}"
2595        );
2596
2597        // The original (uncorrupted) message still decrypts — failed MAC
2598        // check does not persist state changes.
2599        assert_eq!(s.bob_receives(&msg), b"secret");
2600    }
2601
2602    /// Replay an already-decrypted message — both paths must detect the
2603    /// duplicate.
2604    #[test]
2605    fn scenario_replay_message() {
2606        let mut s = DualSession::new(0xBEEF_0007);
2607
2608        let msg = s.alice_sends(b"once");
2609        assert_eq!(s.bob_receives(&msg), b"once");
2610
2611        // Same ciphertext again — should be detected as duplicate
2612        let err = s.bob_receives_err(&msg);
2613        assert!(
2614            matches!(err, SignalProtocolError::DuplicatedMessage(..)),
2615            "expected DuplicatedMessage, got {err:?}"
2616        );
2617    }
2618
2619    proptest! {
2620        /// Reuse the existing session-reset event model from `test-support`,
2621        /// but execute every encrypt/decrypt on both new and legacy codepaths.
2622        #[test]
2623        fn proptest_event_model_matches_legacy(
2624            actions in prop::collection::vec(
2625                (prop::bool::ANY, proptest_arbitrary_interop::arb::<Event>()),
2626                0..40,
2627            ),
2628        ) {
2629            let mut rng = ChaCha8Rng::seed_from_u64(0);
2630            let mut alice = DualParticipant::new(
2631                "alice",
2632                ProtocolAddress::new("9d0652a3-dcc3-4d11-975f-74d61598733f".to_owned(), DeviceId::new(1).unwrap()),
2633                &mut rng,
2634            );
2635            let mut bob = DualParticipant::new(
2636                "bob",
2637                ProtocolAddress::new("796abedb-ca4e-4f18-8803-1fde5b921f9f".to_owned(), DeviceId::new(1).unwrap()),
2638                &mut rng,
2639            );
2640
2641            for (who, event) in actions {
2642                let (me, them) = if who {
2643                    (&mut alice, &mut bob)
2644                } else {
2645                    (&mut bob, &mut alice)
2646                };
2647                me.run_event(them, event, &mut rng)
2648                    .now_or_never()
2649                    .expect("sync");
2650            }
2651
2652            while alice.has_pending_incoming_messages() || bob.has_pending_incoming_messages() {
2653                alice
2654                    .receive_messages(&mut bob, &mut rng)
2655                    .now_or_never()
2656                    .expect("sync");
2657                bob.receive_messages(&mut alice, &mut rng)
2658                    .now_or_never()
2659                    .expect("sync");
2660            }
2661
2662            for _ in 0..8 {
2663                alice
2664                    .send_message(&mut bob, &mut rng)
2665                    .now_or_never()
2666                    .expect("sync");
2667                bob.receive_messages(&mut alice, &mut rng)
2668                    .now_or_never()
2669                    .expect("sync");
2670                bob.send_message(&mut alice, &mut rng)
2671                    .now_or_never()
2672                    .expect("sync");
2673                alice
2674                    .receive_messages(&mut bob, &mut rng)
2675                    .now_or_never()
2676                    .expect("sync");
2677            }
2678
2679            alice.assert_equivalent_with(&bob, "final/alice");
2680            bob.assert_equivalent_with(&alice, "final/bob");
2681        }
2682    }
2683
2684    proptest! {
2685        /// New code can take over a session whose state was last written by
2686        /// legacy code.
2687        ///
2688        /// Runs `legacy_actions` using legacy enc+dec on both sides, then
2689        /// switches both sides to new enc+dec for `new_actions`. The session
2690        /// state — chain keys, ratchet state, SPQR state — was written by the
2691        /// legacy decrypt path; new code must read and advance it correctly.
2692        #[test]
2693        fn proptest_legacy_handover_to_new(
2694            seed in 0u64..u64::MAX,
2695            legacy_actions in prop::collection::vec(
2696                (prop::bool::ANY, prop::collection::vec(any::<u8>(), 0..=64)),
2697                1..=10,
2698            ),
2699            new_actions in prop::collection::vec(
2700                (prop::bool::ANY, prop::collection::vec(any::<u8>(), 0..=64)),
2701                1..=10,
2702            ),
2703        ) {
2704            let mut rng = ChaCha8Rng::seed_from_u64(seed);
2705            let (mut alice_store, mut bob_store, alice_address, bob_address) =
2706                setup_stores(&mut rng);
2707            let now = std::time::SystemTime::now();
2708
2709            // Phase 1: legacy enc + legacy dec advance the session state.
2710            for (alice_sends, plaintext) in &legacy_actions {
2711                let (sender, receiver, recv_addr, send_addr) = if *alice_sends {
2712                    (&mut alice_store, &mut bob_store, &bob_address, &alice_address)
2713                } else {
2714                    (&mut bob_store, &mut alice_store, &alice_address, &bob_address)
2715                };
2716
2717                let ct = legacy::legacy_message_encrypt(
2718                    plaintext,
2719                    recv_addr,
2720                    send_addr,
2721                    &mut sender.session_store,
2722                    &mut sender.identity_store,
2723                    now,
2724                    &mut rng,
2725                )
2726                .now_or_never()
2727                .expect("sync")
2728                .expect("legacy enc");
2729
2730                let signal_msg = match &ct {
2731                    CiphertextMessage::SignalMessage(m) => m,
2732                    other => panic!(
2733                        "expected SignalMessage in legacy phase, got {:?}",
2734                        other.message_type()
2735                    ),
2736                };
2737
2738                let ptext = legacy::legacy_message_decrypt_signal(
2739                    signal_msg,
2740                    send_addr,
2741                    &mut receiver.session_store,
2742                    &mut receiver.identity_store,
2743                    &mut rng,
2744                )
2745                .now_or_never()
2746                .expect("sync")
2747                .expect("legacy dec");
2748
2749                prop_assert_eq!(ptext, plaintext.clone(), "legacy phase: wrong plaintext");
2750            }
2751
2752            // Phase 2: new code takes over the session state left by legacy.
2753            for (alice_sends, plaintext) in &new_actions {
2754                let (sender, receiver, recv_addr, send_addr) = if *alice_sends {
2755                    (&mut alice_store, &mut bob_store, &bob_address, &alice_address)
2756                } else {
2757                    (&mut bob_store, &mut alice_store, &alice_address, &bob_address)
2758                };
2759
2760                let ct = message_encrypt(
2761                    plaintext,
2762                    recv_addr,
2763                    send_addr,
2764                    &mut sender.session_store,
2765                    &mut sender.identity_store,
2766                    now,
2767                    &mut rng,
2768                )
2769                .now_or_never()
2770                .expect("sync")
2771                .expect("new enc");
2772
2773                let signal_msg = match &ct {
2774                    CiphertextMessage::SignalMessage(m) => m,
2775                    other => panic!(
2776                        "expected SignalMessage in new phase, got {:?}",
2777                        other.message_type()
2778                    ),
2779                };
2780
2781                let ptext = message_decrypt_signal(
2782                    signal_msg,
2783                    send_addr,
2784                    recv_addr,
2785                    &mut receiver.session_store,
2786                    &mut receiver.identity_store,
2787                    &mut rng,
2788                )
2789                .now_or_never()
2790                .expect("sync")
2791                .expect("new dec");
2792
2793                prop_assert_eq!(ptext, plaintext.clone(), "new phase: wrong plaintext");
2794            }
2795        }
2796
2797        /// A message encrypted by legacy code on a previous session is correctly
2798        /// decrypted by new code after a session transition.
2799        ///
2800        /// Scenario:
2801        ///   1. `pre_actions` exchanges on session A using legacy enc+dec.
2802        ///   2. Alice encrypts a delayed Whisper on session A using legacy enc
2803        ///      (not yet delivered to Bob).
2804        ///   3. Alice processes a new pre-key bundle from Bob → session B
2805        ///      (session A is archived in Alice's previous_sessions).
2806        ///   4. Alice and Bob establish session B on both sides and exchange
2807        ///      `post_actions` using new enc+dec. Bob's session A' is archived to
2808        ///      his previous_sessions when he receives Alice's first session-B
2809        ///      PreKeySignalMessage.
2810        ///   5. The delayed message from step 2 is delivered to Bob via new
2811        ///      message_decrypt_signal. try_decrypt_from_record must fail on
2812        ///      the current session B' and succeed on the previous session A',
2813        ///      exercising promote_old_session.
2814        #[test]
2815        fn proptest_delayed_message_via_previous_session(
2816            seed in 0u64..u64::MAX,
2817            pre_actions in prop::collection::vec(
2818                (prop::bool::ANY, prop::collection::vec(any::<u8>(), 0..=6)),
2819                0..=6,
2820            ),
2821            post_actions in prop::collection::vec(
2822                (prop::bool::ANY, prop::collection::vec(any::<u8>(), 0..=64)),
2823                0..=6,
2824            ),
2825            delayed_plaintext in prop::collection::vec(any::<u8>(), 1..=64),
2826        ) {
2827            let mut rng = ChaCha8Rng::seed_from_u64(seed);
2828            let (mut alice_store, mut bob_store, alice_address, bob_address) =
2829                setup_stores(&mut rng);
2830            let now = std::time::SystemTime::now();
2831
2832            // ── Phase 1: legacy enc+dec on session A ─────────────────────────
2833
2834            for (alice_sends, plaintext) in &pre_actions {
2835                let (sender, receiver, recv_addr, send_addr) = if *alice_sends {
2836                    (&mut alice_store, &mut bob_store, &bob_address, &alice_address)
2837                } else {
2838                    (&mut bob_store, &mut alice_store, &alice_address, &bob_address)
2839                };
2840
2841                let ct = legacy::legacy_message_encrypt(
2842                    plaintext,
2843                    recv_addr,
2844                    send_addr,
2845                    &mut sender.session_store,
2846                    &mut sender.identity_store,
2847                    now,
2848                    &mut rng,
2849                )
2850                .now_or_never()
2851                .expect("sync")
2852                .expect("pre legacy enc");
2853
2854                let signal_msg = match &ct {
2855                    CiphertextMessage::SignalMessage(m) => m,
2856                    other => panic!(
2857                        "expected SignalMessage in pre phase, got {:?}",
2858                        other.message_type()
2859                    ),
2860                };
2861
2862                let ptext = legacy::legacy_message_decrypt_signal(
2863                    signal_msg,
2864                    send_addr,
2865                    &mut receiver.session_store,
2866                    &mut receiver.identity_store,
2867                    &mut rng,
2868                )
2869                .now_or_never()
2870                .expect("sync")
2871                .expect("pre legacy dec");
2872
2873                prop_assert_eq!(ptext, plaintext.clone(), "pre phase: wrong plaintext");
2874            }
2875
2876            // ── Encrypt delayed message (not yet delivered) ──────────────────
2877
2878            // This Whisper is encrypted by Alice on session A's current chain.
2879            // Bob's session A' is at the same chain index, so it can decrypt
2880            // it later.
2881            let delayed_ct = legacy::legacy_message_encrypt(
2882                &delayed_plaintext,
2883                &bob_address,
2884                &alice_address,
2885                &mut alice_store.session_store,
2886                &mut alice_store.identity_store,
2887                now,
2888                &mut rng,
2889            )
2890            .now_or_never()
2891            .expect("sync")
2892            .expect("delayed legacy enc");
2893
2894            let delayed_signal_msg = match delayed_ct {
2895                CiphertextMessage::SignalMessage(m) => m,
2896                other => panic!(
2897                    "expected SignalMessage for delayed msg, got {:?}",
2898                    other.message_type()
2899                ),
2900            };
2901
2902            // ── Session transition: A → B ─────────────────────────────────────
2903
2904            // Alice processes a new pre-key bundle from Bob. This calls
2905            // promote_state, archiving session A to Alice's previous_sessions.
2906            let bundle = create_bob_bundle(&mut bob_store, 1, 1, 1, &mut rng);
2907            process_prekey_bundle(
2908                &bob_address,
2909                &alice_address,
2910                &mut alice_store.session_store,
2911                &mut alice_store.identity_store,
2912                &bundle,
2913                now,
2914                &mut rng,
2915            )
2916            .now_or_never()
2917            .expect("sync")
2918            .expect("process_prekey_bundle");
2919
2920            // Alice sends her first message on session B (a PreKeySignalMessage
2921            // since B is unacknowledged). When Bob decrypts it, process_prekey
2922            // fires and archives his session A' to previous_sessions.
2923            let session_b_init = message_encrypt(
2924                b"session B init",
2925                &bob_address,
2926                &alice_address,
2927                &mut alice_store.session_store,
2928                &mut alice_store.identity_store,
2929                now,
2930                &mut rng,
2931            )
2932            .now_or_never()
2933            .expect("sync")
2934            .expect("session B init enc");
2935
2936            message_decrypt(
2937                &session_b_init,
2938                &alice_address,
2939                &bob_address,
2940                &mut bob_store.session_store,
2941                &mut bob_store.identity_store,
2942                &mut bob_store.pre_key_store,
2943                &bob_store.signed_pre_key_store,
2944                &mut bob_store.kyber_pre_key_store,
2945                &mut rng,
2946            )
2947            .now_or_never()
2948            .expect("sync")
2949            .expect("session B init dec");
2950            // Bob now has: current = session B', previous_sessions = [session A']
2951
2952            // Bob acknowledges session B on Alice's side. Without this, Alice
2953            // would keep wrapping messages as PreKeySignalMessage, and each
2954            // would trigger another process_prekey on Bob's side, nesting
2955            // sessions further. After this round-trip both sides send Whispers.
2956            let session_b_ack = message_encrypt(
2957                b"session B ack",
2958                &alice_address,
2959                &bob_address,
2960                &mut bob_store.session_store,
2961                &mut bob_store.identity_store,
2962                now,
2963                &mut rng,
2964            )
2965            .now_or_never()
2966            .expect("sync")
2967            .expect("session B ack enc");
2968
2969            let session_b_ack_signal = match &session_b_ack {
2970                CiphertextMessage::SignalMessage(m) => m,
2971                other => panic!(
2972                    "expected Whisper for session B ack, got {:?}",
2973                    other.message_type()
2974                ),
2975            };
2976            message_decrypt_signal(
2977                session_b_ack_signal,
2978                &bob_address,
2979                &alice_address,
2980                &mut alice_store.session_store,
2981                &mut alice_store.identity_store,
2982                &mut rng,
2983            )
2984            .now_or_never()
2985            .expect("sync")
2986            .expect("session B ack dec");
2987            // Alice's session B is now acknowledged; all her sends are Whispers.
2988
2989            // ── Phase 2: new enc+dec on session B ────────────────────────────
2990
2991            for (alice_sends, plaintext) in &post_actions {
2992                let (sender, receiver, recv_addr, send_addr) = if *alice_sends {
2993                    (&mut alice_store, &mut bob_store, &bob_address, &alice_address)
2994                } else {
2995                    (&mut bob_store, &mut alice_store, &alice_address, &bob_address)
2996                };
2997
2998                let ct = message_encrypt(
2999                    plaintext,
3000                    recv_addr,
3001                    send_addr,
3002                    &mut sender.session_store,
3003                    &mut sender.identity_store,
3004                    now,
3005                    &mut rng,
3006                )
3007                .now_or_never()
3008                .expect("sync")
3009                .expect("post new enc");
3010
3011                let signal_msg = match &ct {
3012                    CiphertextMessage::SignalMessage(m) => m,
3013                    other => panic!(
3014                        "expected SignalMessage in post phase, got {:?}",
3015                        other.message_type()
3016                    ),
3017                };
3018
3019                let ptext = message_decrypt_signal(
3020                    signal_msg,
3021                    send_addr,
3022                    recv_addr,
3023                    &mut receiver.session_store,
3024                    &mut receiver.identity_store,
3025                    &mut rng,
3026                )
3027                .now_or_never()
3028                .expect("sync")
3029                .expect("post new dec");
3030
3031                prop_assert_eq!(ptext, plaintext.clone(), "post phase: wrong plaintext");
3032            }
3033
3034            // ── Deliver delayed message ───────────────────────────────────────
3035
3036            // Bob's current session is B'. The delayed message was encrypted
3037            // under session A. try_decrypt_from_record must:
3038            //   1. Try session B' → fail (wrong ratchet key / counter).
3039            //   2. Try session A' from previous_sessions → succeed.
3040            //   3. Call promote_old_session, making A' the current session.
3041            let ptext = message_decrypt_signal(
3042                &delayed_signal_msg,
3043                &alice_address,
3044                &bob_address,
3045                &mut bob_store.session_store,
3046                &mut bob_store.identity_store,
3047                &mut rng,
3048            )
3049            .now_or_never()
3050            .expect("sync")
3051            .expect("delayed msg dec via previous session");
3052
3053            prop_assert_eq!(
3054                ptext,
3055                delayed_plaintext.clone(),
3056                "delayed message: wrong plaintext"
3057            );
3058        }
3059
3060        /// New encrypt and legacy encrypt produce byte-identical ciphertexts
3061        /// when given the same RNG state.
3062        ///
3063        /// Runs two parallel session pairs from the same initial state.
3064        /// Before each encrypt, the RNG is cloned so that both the new and
3065        /// legacy paths start from the same randomness.  If the RNG
3066        /// consumption is identical (one `spqr::send` call per encrypt,
3067        /// one `KeyPair::generate` per DH ratchet step on decrypt), the
3068        /// ciphertexts must be equal.  The receiver sessions are advanced
3069        /// with the same split so that subsequent iterations stay in sync.
3070        #[test]
3071        fn proptest_ciphertext_equality(
3072            seed in 0u64..u64::MAX,
3073            actions in prop::collection::vec(
3074                (prop::bool::ANY, prop::collection::vec(any::<u8>(), 0..=64)),
3075                1..=20,
3076            ),
3077        ) {
3078            let mut rng = ChaCha8Rng::seed_from_u64(seed);
3079            let (mut alice_new, mut bob_new, alice_address, bob_address) =
3080                setup_stores(&mut rng);
3081            // Clone the freshly-initialized stores so both paths start from
3082            // identical state.
3083            let (mut alice_legacy, mut bob_legacy) = (alice_new.clone(), bob_new.clone());
3084            let now = SystemTime::now();
3085
3086            for (alice_sends, plaintext) in &actions {
3087                // Borrow the right stores for sender/receiver on each path.
3088                let (
3089                    (sender_new, receiver_new),
3090                    (sender_legacy, receiver_legacy),
3091                    sender_addr,
3092                    receiver_addr,
3093                ) = if *alice_sends {
3094                    (
3095                        (&mut alice_new, &mut bob_new),
3096                        (&mut alice_legacy, &mut bob_legacy),
3097                        &alice_address,
3098                        &bob_address,
3099                    )
3100                } else {
3101                    (
3102                        (&mut bob_new, &mut alice_new),
3103                        (&mut bob_legacy, &mut alice_legacy),
3104                        &bob_address,
3105                        &alice_address,
3106                    )
3107                };
3108
3109                // Both encrypt calls start from the same RNG position.
3110                let mut enc_rng = rng.clone();
3111
3112                let new_ct = message_encrypt(
3113                    plaintext,
3114                    receiver_addr,
3115                    sender_addr,
3116                    &mut sender_new.session_store,
3117                    &mut sender_new.identity_store,
3118                    now,
3119                    &mut rng,
3120                )
3121                .now_or_never()
3122                .expect("sync")
3123                .expect("new encrypt succeeded");
3124
3125                let legacy_ct = legacy::legacy_message_encrypt(
3126                    plaintext,
3127                    receiver_addr,
3128                    sender_addr,
3129                    &mut sender_legacy.session_store,
3130                    &mut sender_legacy.identity_store,
3131                    now,
3132                    &mut enc_rng,
3133                )
3134                .now_or_never()
3135                .expect("sync")
3136                .expect("legacy encrypt succeeded");
3137
3138                let new_msg = match &new_ct {
3139                    CiphertextMessage::SignalMessage(m) => m,
3140                    other => panic!(
3141                        "expected SignalMessage from new enc, got {:?}",
3142                        other.message_type()
3143                    ),
3144                };
3145                let legacy_msg = match &legacy_ct {
3146                    CiphertextMessage::SignalMessage(m) => m,
3147                    other => panic!(
3148                        "expected SignalMessage from legacy enc, got {:?}",
3149                        other.message_type()
3150                    ),
3151                };
3152
3153                prop_assert_eq!(
3154                    new_msg.serialized(),
3155                    legacy_msg.serialized(),
3156                    "new and legacy produced different ciphertexts from the same RNG state"
3157                );
3158
3159                // Advance both receiver sessions with the same RNG split so
3160                // their states stay in sync for the next iteration.
3161                let mut dec_rng = rng.clone();
3162
3163                let _ = message_decrypt_signal(
3164                    new_msg,
3165                    sender_addr,
3166                    receiver_addr,
3167                    &mut receiver_new.session_store,
3168                    &mut receiver_new.identity_store,
3169                    &mut rng,
3170                )
3171                .now_or_never()
3172                .expect("sync")
3173                .expect("new decrypt succeeded");
3174
3175                let _ = legacy::legacy_message_decrypt_signal(
3176                    legacy_msg,
3177                    sender_addr,
3178                    &mut receiver_legacy.session_store,
3179                    &mut receiver_legacy.identity_store,
3180                    &mut dec_rng,
3181                )
3182                .now_or_never()
3183                .expect("sync")
3184                .expect("legacy decrypt succeeded");
3185            }
3186        }
3187    }
3188}