Refactor 1:1 messaging code

rust/protocol/src/double_ratchet.rs

@@ -0,0 +1,381 @@
+//
+// Copyright 2026 Signal Messenger, LLC.
+// SPDX-License-Identifier: AGPL-3.0-only
+//
+
+//! Double Ratchet implementation.
+//!
+//! The root key, sender chain, and counters are deserialized into typed
+//! fields. Receiver chains are kept in protobuf form and deserialized
+//! lazily on demand.
+//!
+//! References:
+//! - [Double Ratchet spec](https://signal.org/docs/specifications/doubleratchet/)
+
+use rand::{CryptoRng, Rng};
+
+use crate::proto::storage::{SessionStructure, session_structure};
+use crate::ratchet::{ChainKey, MessageKeyGenerator, RootKey};
+use crate::state::InvalidSessionError;
+use crate::{
+    CiphertextMessageType, KeyPair, PrivateKey, PublicKey, Result, SignalProtocolError, consts,
+};
+
+// ── State ────────────────────────────────────────────────────────────
+
+/// The state of a Double Ratchet session.
+///
+/// Contains the root key, sending and receiving chains, and cached
+/// message keys for out-of-order messages. Receiver chains are stored
+/// in protobuf form and deserialized lazily on demand.
+#[derive(Clone)]
+pub(crate) struct RatchetState {
+    pub root_key: RootKey,
+    pub sender_chain: Option<SenderChain>,
+    pub receiver_chains: Vec<session_structure::Chain>,
+    pub previous_counter: u32,
+    /// Maximum number of messages we'll skip ahead in a single chain.
+    /// Set to `consts::MAX_FORWARD_JUMPS` for normal sessions,
+    /// `usize::MAX` for self-sessions (note-to-self), following the
+    /// same pattern as SPQR's `max_jump` chain parameter.
+    pub max_forward_jumps: usize,
+}
+
+/// The sending side of a ratchet: our current ephemeral key pair and
+/// the chain key for encrypting outgoing messages.
+#[derive(Clone)]
+pub(crate) struct SenderChain {
+    pub ratchet_key: KeyPair,
+    pub chain_key: ChainKey,
+}
+
+// Receiver chains are stored as raw `session_structure::Chain` protobuf in
+// `RatchetState::receiver_chains`. Only the ratchet key and chain key are
+// deserialized on demand; skipped message keys stay in protobuf form and are
+// only deserialized individually when a matching counter is found.
+
+// ── Serialization bridge ──────────────────────────────────────────────
+
+impl RatchetState {
+    /// Deserialize the ratchet-relevant fields from a `SessionStructure`.
+    ///
+    /// Only reads root key, counters, and sender chain from `session`.
+    /// Receiver chains are passed in separately (moved, not cloned) by
+    /// the caller.
+    ///
+    /// Identity keys, registration IDs, pending pre-key state, and SPQR
+    /// state are owned by higher layers and are not touched here.
+    ///
+    /// `self_session` must be computed by the caller (requires identity
+    /// key comparison, which is not ratchet-layer knowledge).
+    pub(crate) fn from_pb(
+        session: &SessionStructure,
+        self_session: bool,
+        receiver_chains: Vec<session_structure::Chain>,
+    ) -> std::result::Result<Self, InvalidSessionError> {
+        let root_key_bytes: [u8; 32] = session
+            .root_key
+            .as_slice()
+            .try_into()
+            .map_err(|_| InvalidSessionError("invalid root key"))?;
+
+        let sender_chain = session
+            .sender_chain
+            .as_ref()
+            .map(SenderChain::from_pb)
+            .transpose()?;
+
+        Ok(Self {
+            root_key: RootKey::new(root_key_bytes),
+            sender_chain,
+            receiver_chains,
+            previous_counter: session.previous_counter,
+            max_forward_jumps: if self_session {
+                usize::MAX
+            } else {
+                consts::MAX_FORWARD_JUMPS
+            },
+        })
+    }
+
+    /// Write the ratchet state back into a `SessionStructure`.
+    ///
+    /// Only updates the ratchet-relevant fields; all other fields in
+    /// `session` are left unchanged.
+    pub(crate) fn apply_to_pb(self, session: &mut SessionStructure) {
+        let Self {
+            root_key,
+            sender_chain,
+            receiver_chains,
+            previous_counter,
+            max_forward_jumps: _, // not serialized; derived from session context
+        } = self;
+        session.root_key = root_key.key().to_vec();
+        session.previous_counter = previous_counter;
+        session.sender_chain = sender_chain.as_ref().map(SenderChain::to_pb);
+        session.receiver_chains = receiver_chains;
+    }
+}
+
+impl SenderChain {
+    fn from_pb(chain: &session_structure::Chain) -> std::result::Result<Self, InvalidSessionError> {
+        let public_key = PublicKey::deserialize(&chain.sender_ratchet_key)
+            .map_err(|_| InvalidSessionError("invalid sender ratchet public key"))?;
+        let private_key = PrivateKey::deserialize(&chain.sender_ratchet_key_private)
+            .map_err(|_| InvalidSessionError("invalid sender ratchet private key"))?;
+        let chain_key = ChainKey::from_pb(
+            chain
+                .chain_key
+                .as_ref()
+                .ok_or(InvalidSessionError("missing sender chain key"))?,
+        )?;
+        Ok(Self {
+            ratchet_key: KeyPair {
+                public_key,
+                private_key,
+            },
+            chain_key,
+        })
+    }
+
+    fn to_pb(&self) -> session_structure::Chain {
+        session_structure::Chain {
+            sender_ratchet_key: self.ratchet_key.public_key.serialize().to_vec(),
+            sender_ratchet_key_private: self.ratchet_key.private_key.serialize().to_vec(),
+            chain_key: Some(self.chain_key.to_pb()),
+            message_keys: vec![],
+        }
+    }
+}
+
+impl ChainKey {
+    fn from_pb(
+        pb: &session_structure::chain::ChainKey,
+    ) -> std::result::Result<Self, InvalidSessionError> {
+        let key: [u8; 32] = pb
+            .key
+            .as_slice()
+            .try_into()
+            .map_err(|_| InvalidSessionError("invalid chain key"))?;
+        Ok(Self::new(key, pb.index))
+    }
+
+    fn to_pb(&self) -> session_structure::chain::ChainKey {
+        session_structure::chain::ChainKey {
+            index: self.index(),
+            key: self.key().to_vec(),
+        }
+    }
+}
+
+// ── Operations ───────────────────────────────────────────────────────
+
+impl RatchetState {
+    fn take_root_key(&mut self) -> RootKey {
+        std::mem::replace(&mut self.root_key, RootKey::new([0; 32]))
+    }
+
+    /// Ensure a receiver chain exists for a remote ephemeral key, returning
+    /// its chain key.
+    ///
+    /// If we already have a receiver chain for `their_ephemeral`, return
+    /// its chain key. Otherwise, perform a DH ratchet step: derive a new
+    /// receiver chain from the current root key and the remote ephemeral,
+    /// then generate a fresh sender chain.
+    pub fn ensure_receiver_chain<R: Rng + CryptoRng>(
+        &mut self,
+        their_ephemeral: &PublicKey,
+        csprng: &mut R,
+    ) -> Result<ChainKey> {
+        if let Some(chain_key) = self.find_receiver_chain_key(their_ephemeral)? {
+            return Ok(chain_key);
+        }
+
+        self.dh_ratchet_step(their_ephemeral, csprng)
+    }
+
+    /// Consume the message key for a specific counter value.
+    ///
+    /// If the counter is behind the current chain index, take a cached
+    /// (skipped) key. If it's ahead, advance the chain — caching the
+    /// intermediate keys for out-of-order delivery — up to `max_forward_jumps`.
+    /// Either way, the key is consumed and cannot be retrieved again.
+    pub fn consume_message_key(
+        &mut self,
+        their_ephemeral: &PublicKey,
+        mut chain_key: ChainKey,
+        counter: u32,
+        // The original message type, for error reporting only.
+        original_message_type: CiphertextMessageType,
+        remote_address_for_logging: &str,
+    ) -> Result<MessageKeyGenerator> {
+        let chain_index = chain_key.index();
+
+        if chain_index > counter {
+            // Counter is in the past — look up a cached key.
+            return self
+                .take_skipped_key(their_ephemeral, counter)?
+                .ok_or_else(|| {
+                    log::info!(
+                        "{remote_address_for_logging} Duplicate message for counter: {counter}"
+                    );
+                    SignalProtocolError::DuplicatedMessage(chain_index, counter)
+                });
+        }
+
+        let jump = (counter - chain_index) as usize;
+        if jump > self.max_forward_jumps {
+            log::error!(
+                "{remote_address_for_logging} Exceeded future message limit: {}, index: {chain_index}, counter: {counter}",
+                self.max_forward_jumps,
+            );
+            return Err(SignalProtocolError::InvalidMessage(
+                original_message_type,
+                "message from too far into the future",
+            ));
+        } else if jump > consts::MAX_FORWARD_JUMPS {
+            // This only happens if it is a session with self
+            log::info!(
+                "{remote_address_for_logging} Jumping ahead {jump} messages (index: {chain_index}, counter: {counter})"
+            );
+        }
+
+        // Advance the chain to the target counter, caching skipped keys.
+        while chain_key.index() < counter {
+            self.store_skipped_key(their_ephemeral, chain_key.message_keys());
+            chain_key = chain_key.next_chain_key();
+        }
+
+        // Update the receiver chain to the next key past the one we're returning.
+        self.set_receiver_chain_key(their_ephemeral, chain_key.next_chain_key());
+
+        Ok(chain_key.message_keys())
+    }
+
+    // ── Internals ────────────────────────────────────────────────────
+
+    fn find_receiver_chain_key(
+        &self,
+        their_ephemeral: &PublicKey,
+    ) -> std::result::Result<Option<ChainKey>, InvalidSessionError> {
+        let Some(idx) = self.find_receiver_chain_index(their_ephemeral) else {
+            return Ok(None);
+        };
+        let chain_key_pb = self.receiver_chains[idx]
+            .chain_key
+            .as_ref()
+            .ok_or(InvalidSessionError("missing receiver chain key"))?;
+        Ok(Some(ChainKey::from_pb(chain_key_pb)?))
+    }
+
+    fn dh_ratchet_step<R: Rng + CryptoRng>(
+        &mut self,
+        their_ephemeral: &PublicKey,
+        csprng: &mut R,
+    ) -> Result<ChainKey> {
+        let sender_private_key = self
+            .sender_chain
+            .as_ref()
+            .ok_or(InvalidSessionError("missing sender chain"))?
+            .ratchet_key
+            .private_key;
+
+        // Receiving half-step: root_key + DH(our_sender, their_ephemeral)
+        let current_root_key = self.take_root_key();
+        let (new_root_key, receiver_chain_key) =
+            current_root_key.create_chain(their_ephemeral, &sender_private_key)?;
+
+        // Sending half-step: new_root_key + DH(new_ephemeral, their_ephemeral)
+        let new_sender_key = KeyPair::generate(csprng);
+        let (final_root_key, sender_chain_key) =
+            new_root_key.create_chain(their_ephemeral, &new_sender_key.private_key)?;
+
+        // Record the previous sender chain counter before we replace it.
+        let current_index = self
+            .sender_chain
+            .as_ref()
+            .expect("checked above")
+            .chain_key
+            .index();
+        self.previous_counter = current_index.saturating_sub(1);
+
+        self.root_key = final_root_key;
+
+        self.receiver_chains.push(session_structure::Chain {
+            sender_ratchet_key: their_ephemeral.serialize().to_vec(),
+            sender_ratchet_key_private: vec![],
+            chain_key: Some(receiver_chain_key.to_pb()),
+            message_keys: vec![],
+        });
+        while self.receiver_chains.len() > consts::MAX_RECEIVER_CHAINS {
+            self.receiver_chains.remove(0);
+        }
+
+        self.sender_chain = Some(SenderChain {
+            ratchet_key: new_sender_key,
+            chain_key: sender_chain_key,
+        });
+
+        Ok(receiver_chain_key)
+    }
+
+    fn take_skipped_key(
+        &mut self,
+        their_ephemeral: &PublicKey,
+        counter: u32,
+    ) -> std::result::Result<Option<MessageKeyGenerator>, InvalidSessionError> {
+        let Some(chain_idx) = self.find_receiver_chain_index(their_ephemeral) else {
+            return Ok(None);
+        };
+        let keys = &mut self.receiver_chains[chain_idx].message_keys;
+
+        // Scan by protobuf index field — no deserialization needed for non-matches.
+        let Some(pos) = keys.iter().position(|mk| mk.index == counter) else {
+            return Ok(None);
+        };
+
+        // Remove before deserializing. If from_pb fails the key was corrupt
+        // and unrecoverable — slightly different from main which preserves it,
+        // but the outcome (error) is the same.
+        let key_pb = keys.remove(pos);
+        MessageKeyGenerator::from_pb(key_pb)
+            .map(Some)
+            .map_err(InvalidSessionError)
+    }
+
+    fn store_skipped_key(&mut self, their_ephemeral: &PublicKey, key: MessageKeyGenerator) {
+        let chain_idx = self
+            .find_receiver_chain_index(their_ephemeral)
+            .expect("store_skipped_key called for non-existent chain");
+        let keys = &mut self.receiver_chains[chain_idx].message_keys;
+
+        // TODO: This insert(0) is O(n), making the skip-ahead loop O(n²).
+        // We could switch to push() (appending newest at end) with rposition()
+        // for search and remove(0) for trimming. That makes the common path
+        // O(1) per key. Deferred because it changes serialized key order,
+        // breaking bit-for-bit compatibility with the legacy implementation.
+        keys.insert(0, key.into_pb());
+        if keys.len() > consts::MAX_MESSAGE_KEYS {
+            keys.pop();
+        }
+    }
+
+    fn set_receiver_chain_key(&mut self, their_ephemeral: &PublicKey, chain_key: ChainKey) {
+        let chain_idx = self
+            .find_receiver_chain_index(their_ephemeral)
+            .expect("set_receiver_chain_key called for non-existent chain");
+        self.receiver_chains[chain_idx].chain_key = Some(chain_key.to_pb());
+    }
+
+    fn find_receiver_chain_index(&self, their_ephemeral: &PublicKey) -> Option<usize> {
+        self.receiver_chains.iter().position(|chain| {
+            match PublicKey::deserialize(&chain.sender_ratchet_key) {
+                Ok(key) => &key == their_ephemeral,
+                Err(_) => {
+                    log::warn!("skipping corrupt receiver chain with invalid ratchet key");
+                    false
+                }
+            }
+        })
+    }
+}

rust/protocol/src/handshake.rs

@@ -0,0 +1,56 @@
+//
+// Copyright 2026 Signal Messenger, LLC.
+// SPDX-License-Identifier: AGPL-3.0-only
+//
+
+//! Handshake trait for key agreement protocols.
+//!
+//! Abstracts over different key agreement protocols (PQXDH, and hypothetical
+//! future variants). The trait separates key agreement from ratchet
+//! initialization and session management.
+//!
+//! See [`crate::pqxdh`] for the current production implementation.
+
+use rand::{CryptoRng, Rng};
+
+use crate::Result;
+
+/// A key agreement protocol used to establish a shared secret during
+/// session initialization.
+///
+/// Implementors handle the cryptographic key agreement (DH computations,
+/// KEM encapsulation/decapsulation, KDF). The resulting session secret is
+/// consumed by the ratchet layer to initialize session state. The initiator
+/// also produces a message that must be sent to the recipient for them to
+/// complete the agreement.
+///
+/// The `initiate` method returns `(InitiatorMessage, SessionSecret)` to
+/// enforce a clean boundary: the message is data for the wire, the secret
+/// is data for the ratchet. These are currently protocol-specific types
+/// but should eventually become opaque byte arrays.
+pub(crate) trait Handshake {
+    /// Parameters for the initiator (constructed from a pre-key bundle).
+    type InitiatorParams;
+
+    /// Parameters for the recipient (constructed from an incoming message).
+    type RecipientParams<'a>;
+
+    /// Data the initiator must send to the recipient for them to complete
+    /// the key agreement (e.g., a KEM ciphertext for PQXDH).
+    type InitiatorMessage;
+
+    /// The shared secret derived from the handshake, consumed by the
+    /// ratchet layer to initialize session state.
+    type SessionSecret;
+
+    /// Perform the initiator side of the key agreement.
+    ///
+    /// Returns the message to send and the session secret to keep.
+    fn initiate<R: Rng + CryptoRng>(
+        params: &Self::InitiatorParams,
+        rng: &mut R,
+    ) -> Result<(Self::InitiatorMessage, Self::SessionSecret)>;
+
+    /// Perform the recipient side of the key agreement.
+    fn accept(params: &Self::RecipientParams<'_>) -> Result<Self::SessionSecret>;
+}

rust/protocol/src/lib.rs

@@ -24,22 +24,28 @@
 
 mod consts;
 mod crypto;
+mod double_ratchet;
 pub mod error;
 mod fingerprint;
 mod group_cipher;
+mod handshake;
 mod identity_key;
 pub mod incremental_mac;
 pub mod kem;
+pub mod pqxdh;
 mod proto;
 mod protocol;
 mod ratchet;
 mod sealed_sender;
 mod sender_keys;
 mod session;
-mod session_cipher;
+#[cfg(test)]
+mod session_cipher_legacy;
+mod session_management;
 mod state;
 mod storage;
 mod timestamp;
+mod triple_ratchet;
 
 use error::Result;
 pub use error::SignalProtocolError;
@@ -72,7 +78,7 @@ pub use sealed_sender::{
 };
 pub use sender_keys::SenderKeyRecord;
 pub use session::{process_prekey, process_prekey_bundle};
-pub use session_cipher::{
+pub use session_management::{
     message_decrypt, message_decrypt_prekey, message_decrypt_signal, message_encrypt,
 };
 pub use state::{

rust/protocol/src/pqxdh.rs

@@ -0,0 +1,358 @@
+//
+// Copyright 2026 Signal Messenger, LLC.
+// SPDX-License-Identifier: AGPL-3.0-only
+//
+
+//! PQXDH key agreement protocol.
+//!
+//! This module implements the PQXDH (Post-Quantum Extended Diffie-Hellman) key
+//! agreement, extracting the pure key agreement computation from ratchet
+//! initialization. The output includes derived keys ready for ratchet setup;
+//! the actual ratchet initialization is handled separately in the internal
+//! ratchet module.
+//!
+//! ## Future direction
+//!
+//! The KDF output shape (`RootKey`, `ChainKey`, `[u8; 32]`) is currently
+//! coupled to the Double Ratchet's initialization requirements. Ideally, the
+//! handshake would output a single 32-byte secret and the ratchet layer would
+//! derive whatever it needs from that. This requires a protocol version bump
+//! and should be done alongside a future handshake protocol revision.
+
+use libsignal_core::derive_arrays;
+use rand::{CryptoRng, Rng};
+
+use crate::handshake::Handshake;
+use crate::ratchet::{ChainKey, RootKey};
+use crate::{
+    CiphertextMessageType, IdentityKey, IdentityKeyPair, KeyPair, PublicKey, Result,
+    SignalProtocolError, kem,
+};
+
+/// The PQXDH key agreement protocol.
+///
+/// Implements [`Handshake`] for the Post-Quantum Extended Diffie-Hellman
+/// protocol (4 EC DH + 1 ML-KEM encapsulation/decapsulation).
+pub(crate) struct Pqxdh;
+
+impl Handshake for Pqxdh {
+    type InitiatorParams = InitiatorParameters;
+    type RecipientParams<'a> = RecipientParameters<'a>;
+    type InitiatorMessage = kem::SerializedCiphertext;
+    type SessionSecret = HandshakeKeys;
+
+    fn initiate<R: Rng + CryptoRng>(
+        params: &Self::InitiatorParams,
+        rng: &mut R,
+    ) -> Result<(Self::InitiatorMessage, Self::SessionSecret)> {
+        let result = pqxdh_initiate(params, rng)?;
+        Ok((result.kyber_ciphertext, result.keys))
+    }
+
+    fn accept(params: &Self::RecipientParams<'_>) -> Result<Self::SessionSecret> {
+        pqxdh_accept(params)
+    }
+}
+
+/// The initial PQR (post-quantum ratchet) key derived from the handshake.
+pub(crate) type InitialPQRKey = [u8; 32];
+
+/// Keys derived from a PQXDH handshake, ready for ratchet initialization.
+///
+/// This bundles the KDF output in the shape the ratchet layer expects.
+/// See module-level docs for why this is coupled and the plan to decouple.
+pub(crate) struct HandshakeKeys {
+    pub root_key: RootKey,
+    pub chain_key: ChainKey,
+    pub pqr_key: InitialPQRKey,
+}
+
+impl HandshakeKeys {
+    /// Derive ratchet initialization keys from raw PQXDH shared secret material.
+    fn derive(secret_input: &[u8]) -> Self {
+        Self::derive_with_label(
+            b"WhisperText_X25519_SHA-256_CRYSTALS-KYBER-1024",
+            secret_input,
+        )
+    }
+
+    fn derive_with_label(label: &[u8], secret_input: &[u8]) -> Self {
+        let (root_key_bytes, chain_key_bytes, pqr_bytes) = derive_arrays(|bytes| {
+            hkdf::Hkdf::<sha2::Sha256>::new(None, secret_input)
+                .expand(label, bytes)
+                .expect("valid length")
+        });
+
+        Self {
+            root_key: RootKey::new(root_key_bytes),
+            chain_key: ChainKey::new(chain_key_bytes, 0),
+            pqr_key: pqr_bytes,
+        }
+    }
+}
+
+// ── Initiator ────────────────────────────────────────────────────────
+
+/// The output of a PQXDH key agreement from the initiator's side.
+///
+/// Contains the derived handshake keys and the KEM ciphertext that the
+/// recipient needs to complete the agreement.
+pub(crate) struct InitiatorAgreement {
+    keys: HandshakeKeys,
+    kyber_ciphertext: kem::SerializedCiphertext,
+}
+
+/// Parameters for the initiator side of a PQXDH key agreement.
+///
+/// The initiator fetches the recipient's pre-key bundle from the server
+/// and uses it together with their own identity and ephemeral keys to
+/// compute a shared secret.
+pub struct InitiatorParameters {
+    our_identity_key_pair: IdentityKeyPair,
+    our_ephemeral_key_pair: KeyPair,
+
+    their_identity_key: IdentityKey,
+    their_signed_pre_key: PublicKey,
+    their_one_time_pre_key: Option<PublicKey>,
+    their_ratchet_key: PublicKey,
+    their_kyber_pre_key: kem::PublicKey,
+}
+
+impl InitiatorParameters {
+    pub fn new(
+        our_identity_key_pair: IdentityKeyPair,
+        our_ephemeral_key_pair: KeyPair,
+        their_identity_key: IdentityKey,
+        their_signed_pre_key: PublicKey,
+        their_ratchet_key: PublicKey,
+        their_kyber_pre_key: kem::PublicKey,
+    ) -> Self {
+        Self {
+            our_identity_key_pair,
+            our_ephemeral_key_pair,
+            their_identity_key,
+            their_one_time_pre_key: None,
+            their_signed_pre_key,
+            their_ratchet_key,
+            their_kyber_pre_key,
+        }
+    }
+
+    pub fn set_their_one_time_pre_key(&mut self, ec_public: PublicKey) {
+        self.their_one_time_pre_key = Some(ec_public);
+    }
+
+    #[inline]
+    pub fn our_identity_key_pair(&self) -> &IdentityKeyPair {
+        &self.our_identity_key_pair
+    }
+
+    #[inline]
+    pub fn our_ephemeral_key_pair(&self) -> &KeyPair {
+        &self.our_ephemeral_key_pair
+    }
+
+    #[inline]
+    pub fn their_identity_key(&self) -> &IdentityKey {
+        &self.their_identity_key
+    }
+
+    #[inline]
+    pub fn their_signed_pre_key(&self) -> &PublicKey {
+        &self.their_signed_pre_key
+    }
+
+    #[inline]
+    pub fn their_one_time_pre_key(&self) -> Option<&PublicKey> {
+        self.their_one_time_pre_key.as_ref()
+    }
+
+    #[inline]
+    pub fn their_kyber_pre_key(&self) -> &kem::PublicKey {
+        &self.their_kyber_pre_key
+    }
+
+    #[inline]
+    pub fn their_ratchet_key(&self) -> &PublicKey {
+        &self.their_ratchet_key
+    }
+}
+
+/// Perform the initiator side of the PQXDH key agreement.
+///
+/// Computes DH shared secrets and KEM encapsulation, then applies the KDF
+/// to produce keys ready for ratchet initialization.
+pub(crate) fn pqxdh_initiate<R: Rng + CryptoRng>(
+    parameters: &InitiatorParameters,
+    mut csprng: &mut R,
+) -> Result<InitiatorAgreement> {
+    let mut secrets = Vec::with_capacity(32 * 6);
+
+    secrets.extend_from_slice(&[0xFFu8; 32]); // discontinuity bytes
+
+    secrets.extend_from_slice(
+        &parameters
+            .our_identity_key_pair
+            .private_key()
+            .calculate_agreement(&parameters.their_signed_pre_key)?,
+    );
+
+    let our_ephemeral_private_key = parameters.our_ephemeral_key_pair.private_key;
+
+    secrets.extend_from_slice(
+        &our_ephemeral_private_key
+            .calculate_agreement(parameters.their_identity_key.public_key())?,
+    );
+
+    secrets.extend_from_slice(
+        &our_ephemeral_private_key.calculate_agreement(&parameters.their_signed_pre_key)?,
+    );
+
+    if let Some(their_one_time_prekey) = &parameters.their_one_time_pre_key {
+        secrets.extend_from_slice(
+            &our_ephemeral_private_key.calculate_agreement(their_one_time_prekey)?,
+        );
+    }
+
+    let kyber_ciphertext = {
+        let (ss, ct) = parameters.their_kyber_pre_key.encapsulate(&mut csprng)?;
+        secrets.extend_from_slice(ss.as_ref());
+        ct
+    };
+
+    Ok(InitiatorAgreement {
+        keys: HandshakeKeys::derive(&secrets),
+        kyber_ciphertext,
+    })
+}
+
+// ── Recipient ────────────────────────────────────────────────────────
+
+/// Parameters for the recipient side of a PQXDH key agreement.
+///
+/// The recipient uses their own pre-keys together with the initiator's
+/// identity and base keys (received in the pre-key message) to compute
+/// the same shared secret.
+pub struct RecipientParameters<'a> {
+    our_identity_key_pair: IdentityKeyPair,
+    our_signed_pre_key_pair: KeyPair,
+    our_one_time_pre_key_pair: Option<KeyPair>,
+    our_kyber_pre_key_pair: kem::KeyPair,
+
+    their_identity_key: IdentityKey,
+    their_ephemeral_key: PublicKey,
+    their_kyber_ciphertext: &'a kem::SerializedCiphertext,
+}
+
+impl<'a> RecipientParameters<'a> {
+    pub fn new(
+        our_identity_key_pair: IdentityKeyPair,
+        our_signed_pre_key_pair: KeyPair,
+        our_one_time_pre_key_pair: Option<KeyPair>,
+        our_kyber_pre_key_pair: kem::KeyPair,
+        their_identity_key: IdentityKey,
+        their_ephemeral_key: PublicKey,
+        their_kyber_ciphertext: &'a kem::SerializedCiphertext,
+    ) -> Self {
+        Self {
+            our_identity_key_pair,
+            our_signed_pre_key_pair,
+            our_one_time_pre_key_pair,
+            our_kyber_pre_key_pair,
+            their_identity_key,
+            their_ephemeral_key,
+            their_kyber_ciphertext,
+        }
+    }
+
+    #[inline]
+    pub fn our_identity_key_pair(&self) -> &IdentityKeyPair {
+        &self.our_identity_key_pair
+    }
+
+    #[inline]
+    pub fn our_signed_pre_key_pair(&self) -> &KeyPair {
+        &self.our_signed_pre_key_pair
+    }
+
+    #[inline]
+    pub fn our_one_time_pre_key_pair(&self) -> Option<&KeyPair> {
+        self.our_one_time_pre_key_pair.as_ref()
+    }
+
+    #[inline]
+    pub fn our_kyber_pre_key_pair(&self) -> &kem::KeyPair {
+        &self.our_kyber_pre_key_pair
+    }
+
+    #[inline]
+    pub fn their_identity_key(&self) -> &IdentityKey {
+        &self.their_identity_key
+    }
+
+    #[inline]
+    pub fn their_ephemeral_key(&self) -> &PublicKey {
+        &self.their_ephemeral_key
+    }
+
+    #[inline]
+    pub fn their_kyber_ciphertext(&self) -> &kem::SerializedCiphertext {
+        self.their_kyber_ciphertext
+    }
+}
+
+/// Perform the recipient side of the PQXDH key agreement.
+///
+/// Computes DH shared secrets and KEM decapsulation, then applies the KDF
+/// to produce keys ready for ratchet initialization.
+pub(crate) fn pqxdh_accept(parameters: &RecipientParameters) -> Result<HandshakeKeys> {
+    // Validate the initiator's base key before doing any computation.
+    if !parameters.their_ephemeral_key.is_canonical() {
+        return Err(SignalProtocolError::InvalidMessage(
+            CiphertextMessageType::PreKey,
+            "incoming base key is invalid",
+        ));
+    }
+
+    let mut secrets = Vec::with_capacity(32 * 6);
+
+    secrets.extend_from_slice(&[0xFFu8; 32]); // discontinuity bytes
+
+    secrets.extend_from_slice(
+        &parameters
+            .our_signed_pre_key_pair
+            .private_key
+            .calculate_agreement(parameters.their_identity_key.public_key())?,
+    );
+
+    secrets.extend_from_slice(
+        &parameters
+            .our_identity_key_pair
+            .private_key()
+            .calculate_agreement(&parameters.their_ephemeral_key)?,
+    );
+
+    secrets.extend_from_slice(
+        &parameters
+            .our_signed_pre_key_pair
+            .private_key
+            .calculate_agreement(&parameters.their_ephemeral_key)?,
+    );
+
+    if let Some(our_one_time_pre_key_pair) = &parameters.our_one_time_pre_key_pair {
+        secrets.extend_from_slice(
+            &our_one_time_pre_key_pair
+                .private_key
+                .calculate_agreement(&parameters.their_ephemeral_key)?,
+        );
+    }
+
+    secrets.extend_from_slice(
+        &parameters
+            .our_kyber_pre_key_pair
+            .secret_key
+            .decapsulate(parameters.their_kyber_ciphertext)?,
+    );
+
+    Ok(HandshakeKeys::derive(&secrets))
+}

rust/protocol/src/protocol.rs

@@ -203,7 +203,7 @@ impl SignalMessage {
 
         let Some(expected) = Self::serialize_addresses(sender_address, recipient_address) else {
             log::warn!(
-                "Local addresses not valid Service IDs: sender={}, recipient={}",
+                "Locally supplied addresses not valid Service IDs: sender={}, recipient={}",
                 sender_address,
                 recipient_address,
             );

rust/protocol/src/ratchet.rs

@@ -4,39 +4,25 @@
 //
 
 mod keys;
-mod params;
 
-use libsignal_core::derive_arrays;
 use rand::{CryptoRng, Rng};
 
 pub(crate) use self::keys::{ChainKey, MessageKeyGenerator, RootKey};
-pub use self::params::{AliceSignalProtocolParameters, BobSignalProtocolParameters};
+use crate::handshake::Handshake;
+use crate::pqxdh::{HandshakeKeys, Pqxdh};
+// Re-export the parameter types for backward compatibility.
+// Callers (session.rs, tests) use these via `ratchet::`.
+pub use crate::pqxdh::{InitiatorParameters, RecipientParameters};
 use crate::protocol::CIPHERTEXT_MESSAGE_CURRENT_VERSION;
 use crate::state::SessionState;
 use crate::{KeyPair, Result, SessionRecord, SignalProtocolError, consts};
 
-type InitialPQRKey = [u8; 32];
-
-fn derive_keys(secret_input: &[u8]) -> (RootKey, ChainKey, InitialPQRKey) {
-    derive_keys_with_label(
-        b"WhisperText_X25519_SHA-256_CRYSTALS-KYBER-1024",
-        secret_input,
-    )
-}
-
-fn derive_keys_with_label(label: &[u8], secret_input: &[u8]) -> (RootKey, ChainKey, InitialPQRKey) {
-    let (root_key_bytes, chain_key_bytes, pqr_bytes) = derive_arrays(|bytes| {
-        hkdf::Hkdf::<sha2::Sha256>::new(None, secret_input)
-            .expand(label, bytes)
-            .expect("valid length")
-    });
-
-    let root_key = RootKey::new(root_key_bytes);
-    let chain_key = ChainKey::new(chain_key_bytes, 0);
-    let pqr_key: InitialPQRKey = pqr_bytes;
-
-    (root_key, chain_key, pqr_key)
-}
+// Backward-compatible aliases for the old names. These keep existing
+// external callers (tests, bridge code) compiling during the transition.
+#[doc(hidden)]
+pub type AliceSignalProtocolParameters = InitiatorParameters;
+#[doc(hidden)]
+pub type BobSignalProtocolParameters<'a> = RecipientParameters<'a>;
 
 fn spqr_chain_params(self_connection: bool) -> spqr::ChainParams {
     #[allow(clippy::needless_update)]
@@ -51,47 +37,44 @@ fn spqr_chain_params(self_connection: bool) -> spqr::ChainParams {
     }
 }
 
+/// Initialize a session from the initiator's side.
+///
+/// Performs the PQXDH key agreement and then sets up the Double Ratchet
+/// and SPQR state.
 pub(crate) fn initialize_alice_session<R: Rng + CryptoRng>(
-    parameters: &AliceSignalProtocolParameters,
-    mut csprng: &mut R,
+    parameters: &InitiatorParameters,
+    csprng: &mut R,
+) -> Result<SessionState> {
+    let (
+        kyber_ciphertext,
+        HandshakeKeys {
+            root_key,
+            chain_key,
+            pqr_key,
+        },
+    ) = Pqxdh::initiate(parameters, csprng)?;
+
+    initialize_initiator_session(
+        parameters,
+        root_key,
+        chain_key,
+        pqr_key,
+        kyber_ciphertext,
+        csprng,
+    )
+}
+
+fn initialize_initiator_session<R: Rng + CryptoRng>(
+    parameters: &InitiatorParameters,
+    root_key: RootKey,
+    chain_key: ChainKey,
+    pqr_key: [u8; 32],
+    kyber_ciphertext: crate::kem::SerializedCiphertext,
+    csprng: &mut R,
 ) -> Result<SessionState> {
     let local_identity = parameters.our_identity_key_pair().identity_key();
 
-    let mut secrets = Vec::with_capacity(32 * 6);
-
-    secrets.extend_from_slice(&[0xFFu8; 32]); // "discontinuity bytes"
-
-    let our_base_private_key = parameters.our_base_key_pair().private_key;
-
-    secrets.extend_from_slice(
-        &parameters
-            .our_identity_key_pair()
-            .private_key()
-            .calculate_agreement(parameters.their_signed_pre_key())?,
-    );
-
-    secrets.extend_from_slice(
-        &our_base_private_key.calculate_agreement(parameters.their_identity_key().public_key())?,
-    );
-
-    secrets.extend_from_slice(
-        &our_base_private_key.calculate_agreement(parameters.their_signed_pre_key())?,
-    );
-
-    if let Some(their_one_time_prekey) = parameters.their_one_time_pre_key() {
-        secrets
-            .extend_from_slice(&our_base_private_key.calculate_agreement(their_one_time_prekey)?);
-    }
-
-    let kyber_ciphertext = {
-        let (ss, ct) = parameters.their_kyber_pre_key().encapsulate(&mut csprng)?;
-        secrets.extend_from_slice(ss.as_ref());
-        ct
-    };
-
-    let (root_key, chain_key, pqr_key) = derive_keys(&secrets);
-
-    let sending_ratchet_key = KeyPair::generate(&mut csprng);
+    let sending_ratchet_key = KeyPair::generate(csprng);
     let (sending_chain_root_key, sending_chain_chain_key) = root_key.create_chain(
         parameters.their_ratchet_key(),
         &sending_ratchet_key.private_key,
@@ -118,7 +101,7 @@ pub(crate) fn initialize_alice_session<R: Rng + CryptoRng>(
         local_identity,
         parameters.their_identity_key(),
         &sending_chain_root_key,
-        &parameters.our_base_key_pair().public_key,
+        &parameters.our_ephemeral_key_pair().public_key,
         pqr_state,
     )
     .with_receiver_chain(parameters.their_ratchet_key(), &chain_key)
@@ -129,61 +112,38 @@ pub(crate) fn initialize_alice_session<R: Rng + CryptoRng>(
     Ok(session)
 }
 
+/// Initialize a session from the recipient's side.
+///
+/// Performs the PQXDH key agreement and then sets up the Double Ratchet
+/// and SPQR state.
 pub(crate) fn initialize_bob_session(
-    parameters: &BobSignalProtocolParameters,
+    parameters: &RecipientParameters,
+    our_ratchet_key_pair: &KeyPair,
 ) -> Result<SessionState> {
-    // validate their base key
-    if !parameters.their_base_key().is_canonical() {
-        return Err(SignalProtocolError::InvalidMessage(
-            crate::CiphertextMessageType::PreKey,
-            "incoming base key is invalid",
-        ));
-    }
+    let HandshakeKeys {
+        root_key,
+        chain_key,
+        pqr_key,
+    } = Pqxdh::accept(parameters)?;
 
+    initialize_recipient_session(
+        parameters,
+        our_ratchet_key_pair,
+        root_key,
+        chain_key,
+        pqr_key,
+    )
+}
+
+fn initialize_recipient_session(
+    parameters: &RecipientParameters,
+    our_ratchet_key_pair: &KeyPair,
+    root_key: RootKey,
+    chain_key: ChainKey,
+    pqr_key: [u8; 32],
+) -> Result<SessionState> {
     let local_identity = parameters.our_identity_key_pair().identity_key();
 
-    let mut secrets = Vec::with_capacity(32 * 6);
-
-    secrets.extend_from_slice(&[0xFFu8; 32]); // "discontinuity bytes"
-
-    secrets.extend_from_slice(
-        &parameters
-            .our_signed_pre_key_pair()
-            .private_key
-            .calculate_agreement(parameters.their_identity_key().public_key())?,
-    );
-
-    secrets.extend_from_slice(
-        &parameters
-            .our_identity_key_pair()
-            .private_key()
-            .calculate_agreement(parameters.their_base_key())?,
-    );
-
-    secrets.extend_from_slice(
-        &parameters
-            .our_signed_pre_key_pair()
-            .private_key
-            .calculate_agreement(parameters.their_base_key())?,
-    );
-
-    if let Some(our_one_time_pre_key_pair) = parameters.our_one_time_pre_key_pair() {
-        secrets.extend_from_slice(
-            &our_one_time_pre_key_pair
-                .private_key
-                .calculate_agreement(parameters.their_base_key())?,
-        );
-    }
-
-    secrets.extend_from_slice(
-        &parameters
-            .our_kyber_pre_key_pair()
-            .secret_key
-            .decapsulate(parameters.their_kyber_ciphertext())?,
-    );
-
-    let (root_key, chain_key, pqr_key) = derive_keys(&secrets);
-
     let self_session = local_identity == parameters.their_identity_key();
     let pqr_state = spqr::initial_state(spqr::Params {
         auth_key: &pqr_key,
@@ -199,21 +159,22 @@ pub(crate) fn initialize_bob_session(
             "post-quantum ratchet: error creating initial B2A state: {e}"
         ))
     })?;
+
     let session = SessionState::new(
         CIPHERTEXT_MESSAGE_CURRENT_VERSION,
         local_identity,
         parameters.their_identity_key(),
         &root_key,
-        parameters.their_base_key(),
+        parameters.their_ephemeral_key(),
         pqr_state,
     )
-    .with_sender_chain(parameters.our_ratchet_key_pair(), &chain_key);
+    .with_sender_chain(our_ratchet_key_pair, &chain_key);
 
     Ok(session)
 }
 
 pub fn initialize_alice_session_record<R: Rng + CryptoRng>(
-    parameters: &AliceSignalProtocolParameters,
+    parameters: &InitiatorParameters,
     csprng: &mut R,
 ) -> Result<SessionRecord> {
     Ok(SessionRecord::new(initialize_alice_session(
@@ -222,7 +183,11 @@ pub fn initialize_alice_session_record<R: Rng + CryptoRng>(
 }
 
 pub fn initialize_bob_session_record(
-    parameters: &BobSignalProtocolParameters,
+    parameters: &RecipientParameters,
+    our_ratchet_key_pair: &KeyPair,
 ) -> Result<SessionRecord> {
-    Ok(SessionRecord::new(initialize_bob_session(parameters)?))
+    Ok(SessionRecord::new(initialize_bob_session(
+        parameters,
+        our_ratchet_key_pair,
+    )?))
 }

rust/protocol/src/ratchet/keys.rs

@@ -10,11 +10,21 @@ use libsignal_core::derive_arrays;
 use crate::proto::storage::session_structure;
 use crate::{PrivateKey, PublicKey, Result, crypto};
 
+#[derive(Clone)]
 pub(crate) enum MessageKeyGenerator {
     Keys(MessageKeys),
     Seed((Vec<u8>, u32)),
 }
 
+impl fmt::Debug for MessageKeyGenerator {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Keys(k) => f.debug_tuple("Keys").field(&k.counter()).finish(),
+            Self::Seed((_, idx)) => f.debug_tuple("Seed").field(idx).finish(),
+        }
+    }
+}
+
 impl MessageKeyGenerator {
     pub(crate) fn new_from_seed(seed: &[u8], counter: u32) -> Self {
         Self::Seed((seed.to_vec(), counter))

rust/protocol/src/ratchet/params.rs

@@ -1,159 +0,0 @@
-//
-// Copyright 2020 Signal Messenger, LLC.
-// SPDX-License-Identifier: AGPL-3.0-only
-//
-
-use crate::{IdentityKey, IdentityKeyPair, KeyPair, PublicKey, kem};
-
-pub struct AliceSignalProtocolParameters {
-    our_identity_key_pair: IdentityKeyPair,
-    our_base_key_pair: KeyPair,
-
-    their_identity_key: IdentityKey,
-    their_signed_pre_key: PublicKey,
-    their_one_time_pre_key: Option<PublicKey>,
-    their_ratchet_key: PublicKey,
-    their_kyber_pre_key: kem::PublicKey,
-}
-
-impl AliceSignalProtocolParameters {
-    pub fn new(
-        our_identity_key_pair: IdentityKeyPair,
-        our_base_key_pair: KeyPair,
-        their_identity_key: IdentityKey,
-        their_signed_pre_key: PublicKey,
-        their_ratchet_key: PublicKey,
-        their_kyber_pre_key: kem::PublicKey,
-    ) -> Self {
-        Self {
-            our_identity_key_pair,
-            our_base_key_pair,
-            their_identity_key,
-            their_signed_pre_key,
-            their_one_time_pre_key: None,
-            their_ratchet_key,
-            their_kyber_pre_key,
-        }
-    }
-
-    pub fn set_their_one_time_pre_key(&mut self, ec_public: PublicKey) {
-        self.their_one_time_pre_key = Some(ec_public);
-    }
-
-    pub fn with_their_one_time_pre_key(mut self, ec_public: PublicKey) -> Self {
-        self.set_their_one_time_pre_key(ec_public);
-        self
-    }
-
-    #[inline]
-    pub fn our_identity_key_pair(&self) -> &IdentityKeyPair {
-        &self.our_identity_key_pair
-    }
-
-    #[inline]
-    pub fn our_base_key_pair(&self) -> &KeyPair {
-        &self.our_base_key_pair
-    }
-
-    #[inline]
-    pub fn their_identity_key(&self) -> &IdentityKey {
-        &self.their_identity_key
-    }
-
-    #[inline]
-    pub fn their_signed_pre_key(&self) -> &PublicKey {
-        &self.their_signed_pre_key
-    }
-
-    #[inline]
-    pub fn their_one_time_pre_key(&self) -> Option<&PublicKey> {
-        self.their_one_time_pre_key.as_ref()
-    }
-
-    #[inline]
-    pub fn their_kyber_pre_key(&self) -> &kem::PublicKey {
-        &self.their_kyber_pre_key
-    }
-
-    #[inline]
-    pub fn their_ratchet_key(&self) -> &PublicKey {
-        &self.their_ratchet_key
-    }
-}
-
-pub struct BobSignalProtocolParameters<'a> {
-    our_identity_key_pair: IdentityKeyPair,
-    our_signed_pre_key_pair: KeyPair,
-    our_one_time_pre_key_pair: Option<KeyPair>,
-    our_ratchet_key_pair: KeyPair,
-    our_kyber_pre_key_pair: kem::KeyPair,
-
-    their_identity_key: IdentityKey,
-    their_base_key: PublicKey,
-    their_kyber_ciphertext: &'a kem::SerializedCiphertext,
-}
-
-impl<'a> BobSignalProtocolParameters<'a> {
-    #[allow(clippy::too_many_arguments)]
-    pub fn new(
-        our_identity_key_pair: IdentityKeyPair,
-        our_signed_pre_key_pair: KeyPair,
-        our_one_time_pre_key_pair: Option<KeyPair>,
-        our_ratchet_key_pair: KeyPair,
-        our_kyber_pre_key_pair: kem::KeyPair,
-        their_identity_key: IdentityKey,
-        their_base_key: PublicKey,
-        their_kyber_ciphertext: &'a kem::SerializedCiphertext,
-    ) -> Self {
-        Self {
-            our_identity_key_pair,
-            our_signed_pre_key_pair,
-            our_one_time_pre_key_pair,
-            our_ratchet_key_pair,
-            our_kyber_pre_key_pair,
-            their_identity_key,
-            their_base_key,
-            their_kyber_ciphertext,
-        }
-    }
-
-    #[inline]
-    pub fn our_identity_key_pair(&self) -> &IdentityKeyPair {
-        &self.our_identity_key_pair
-    }
-
-    #[inline]
-    pub fn our_signed_pre_key_pair(&self) -> &KeyPair {
-        &self.our_signed_pre_key_pair
-    }
-
-    #[inline]
-    pub fn our_one_time_pre_key_pair(&self) -> Option<&KeyPair> {
-        self.our_one_time_pre_key_pair.as_ref()
-    }
-
-    #[inline]
-    pub fn our_ratchet_key_pair(&self) -> &KeyPair {
-        &self.our_ratchet_key_pair
-    }
-
-    #[inline]
-    pub fn our_kyber_pre_key_pair(&self) -> &kem::KeyPair {
-        &self.our_kyber_pre_key_pair
-    }
-
-    #[inline]
-    pub fn their_identity_key(&self) -> &IdentityKey {
-        &self.their_identity_key
-    }
-
-    #[inline]
-    pub fn their_base_key(&self) -> &PublicKey {
-        &self.their_base_key
-    }
-
-    #[inline]
-    pub fn their_kyber_ciphertext(&self) -> &kem::SerializedCiphertext {
-        self.their_kyber_ciphertext
-    }
-}

rust/protocol/src/sealed_sender.rs

@@ -23,7 +23,7 @@ use crate::{
     IdentityKeyStore, KeyPair, KyberPreKeyStore, PreKeySignalMessage, PreKeyStore, PrivateKey,
     ProtocolAddress, PublicKey, Result, ServiceId, ServiceIdFixedWidthBinaryBytes, SessionRecord,
     SessionStore, SignalMessage, SignalProtocolError, SignedPreKeyStore, Timestamp, crypto,
-    message_encrypt, proto, session_cipher,
+    message_encrypt, proto, session_management,
 };
 
 #[derive(Debug, Clone)]
@@ -2044,7 +2044,7 @@ pub async fn sealed_sender_decrypt(
     let message = match usmc.msg_type()? {
         CiphertextMessageType::Whisper => {
             let ctext = SignalMessage::try_from(usmc.contents()?)?;
-            session_cipher::message_decrypt_signal(
+            session_management::message_decrypt_signal(
                 &ctext,
                 &remote_address,
                 session_store,
@@ -2055,7 +2055,7 @@ pub async fn sealed_sender_decrypt(
         }
         CiphertextMessageType::PreKey => {
             let ctext = PreKeySignalMessage::try_from(usmc.contents()?)?;
-            session_cipher::message_decrypt_prekey(
+            session_management::message_decrypt_prekey(
                 &ctext,
                 &remote_address,
                 &local_address,

rust/protocol/src/session.rs

@@ -145,16 +145,16 @@ async fn process_prekey_impl(
 
     let parameters = BobSignalProtocolParameters::new(
         identity_store.get_identity_key_pair().await?,
-        our_signed_pre_key_pair, // signed pre key
+        our_signed_pre_key_pair,
         our_one_time_pre_key_pair,
-        our_signed_pre_key_pair, // ratchet key
         our_kyber_pre_key_pair,
         *message.identity_key(),
         *message.base_key(),
         kyber_ciphertext,
     );
 
-    let mut new_session = ratchet::initialize_bob_session(&parameters)?;
+    // The recipient's initial ratchet key is the signed pre-key.
+    let mut new_session = ratchet::initialize_bob_session(&parameters, &our_signed_pre_key_pair)?;
 
     new_session.set_local_registration_id(identity_store.get_local_registration_id().await?);
     new_session.set_remote_registration_id(message.registration_id());

rust/protocol/src/session_cipher_legacy.rs

@@ -1,7 +1,16 @@
 //
-// Copyright 2020-2022 Signal Messenger, LLC.
+// Copyright 2020-2026 Signal Messenger, LLC.
 // SPDX-License-Identifier: AGPL-3.0-only
 //
+// Frozen snapshot of session_cipher.rs taken before the protocol refactoring.
+// Used exclusively for interoperability tests: the legacy encrypt/decrypt paths
+// must be able to exchange messages with the new implementation.
+//
+// DO NOT modify the crypto logic in this file. The point is to have an
+// immutable reference implementation to test against.
+
+#![cfg(test)]
+#![allow(dead_code)]
 
 use std::time::SystemTime;
 
@@ -16,7 +25,7 @@ use crate::{
     SessionRecord, SessionStore, SignalMessage, SignalProtocolError, SignedPreKeyStore, session,
 };
 
-pub async fn message_encrypt<R: Rng + CryptoRng>(
+pub async fn legacy_message_encrypt<R: Rng + CryptoRng>(
     ptext: &[u8],
     remote_address: &ProtocolAddress,
     local_address: &ProtocolAddress,
@@ -162,7 +171,7 @@ pub async fn message_encrypt<R: Rng + CryptoRng>(
 }
 
 #[allow(clippy::too_many_arguments)]
-pub async fn message_decrypt<R: Rng + CryptoRng>(
+pub async fn legacy_message_decrypt<R: Rng + CryptoRng>(
     ciphertext: &CiphertextMessage,
     remote_address: &ProtocolAddress,
     local_address: &ProtocolAddress,
@@ -175,11 +184,11 @@ pub async fn message_decrypt<R: Rng + CryptoRng>(
 ) -> Result<Vec<u8>> {
     match ciphertext {
         CiphertextMessage::SignalMessage(m) => {
-            let _ = local_address;
-            message_decrypt_signal(m, remote_address, session_store, identity_store, csprng).await
+            legacy_message_decrypt_signal(m, remote_address, session_store, identity_store, csprng)
+                .await
         }
         CiphertextMessage::PreKeySignalMessage(m) => {
-            message_decrypt_prekey(
+            legacy_message_decrypt_prekey(
                 m,
                 remote_address,
                 local_address,
@@ -200,7 +209,7 @@ pub async fn message_decrypt<R: Rng + CryptoRng>(
 }
 
 #[allow(clippy::too_many_arguments)]
-pub async fn message_decrypt_prekey<R: Rng + CryptoRng>(
+pub async fn legacy_message_decrypt_prekey<R: Rng + CryptoRng>(
     ciphertext: &PreKeySignalMessage,
     remote_address: &ProtocolAddress,
     local_address: &ProtocolAddress,
@@ -285,7 +294,7 @@ pub async fn message_decrypt_prekey<R: Rng + CryptoRng>(
     Ok(ptext)
 }
 
-pub async fn message_decrypt_signal<R: Rng + CryptoRng>(
+pub async fn legacy_message_decrypt_signal<R: Rng + CryptoRng>(
     ciphertext: &SignalMessage,
     remote_address: &ProtocolAddress,
     session_store: &mut dyn SessionStore,

rust/protocol/src/state/session.rs

@@ -8,18 +8,21 @@ use std::time::{Duration, SystemTime};
 
 use bitflags::bitflags;
 use prost::Message;
+#[cfg(test)]
 use rand::{CryptoRng, Rng};
 use subtle::ConstantTimeEq;
 
 use crate::proto::storage::{RecordStructure, SessionStructure, session_structure};
 use crate::protocol::CIPHERTEXT_MESSAGE_PRE_KYBER_VERSION;
-use crate::ratchet::{ChainKey, MessageKeyGenerator, RootKey};
+#[cfg(test)]
+use crate::ratchet::MessageKeyGenerator;
+use crate::ratchet::{ChainKey, RootKey};
 use crate::state::{KyberPreKeyId, PreKeyId, SignedPreKeyId};
 use crate::{IdentityKey, KeyPair, PrivateKey, PublicKey, SignalProtocolError, consts, kem};
 
 /// A distinct error type to keep from accidentally propagating deserialization errors.
 #[derive(Debug)]
-pub(crate) struct InvalidSessionError(&'static str);
+pub(crate) struct InvalidSessionError(pub(crate) &'static str);
 
 impl std::fmt::Display for InvalidSessionError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -214,10 +217,12 @@ impl SessionState {
         self.session.previous_counter
     }
 
+    #[cfg(test)]
     pub(crate) fn set_previous_counter(&mut self, ctr: u32) {
         self.session.previous_counter = ctr;
     }
 
+    #[cfg(test)]
     pub(crate) fn root_key(&self) -> Result<RootKey, InvalidSessionError> {
         let root_key_bytes = self.session.root_key[..]
             .try_into()
@@ -225,6 +230,7 @@ impl SessionState {
         Ok(RootKey::new(root_key_bytes))
     }
 
+    #[cfg(test)]
     pub(crate) fn set_root_key(&mut self, root_key: &RootKey) {
         self.session.root_key = root_key.key().to_vec();
     }
@@ -428,6 +434,7 @@ impl SessionState {
         self.session.sender_chain = Some(new_chain);
     }
 
+    #[cfg(test)]
     pub(crate) fn get_message_keys(
         &mut self,
         sender: &PublicKey,
@@ -454,6 +461,7 @@ impl SessionState {
         Ok(None)
     }
 
+    #[cfg(test)]
     pub(crate) fn set_message_keys(
         &mut self,
         sender: &PublicKey,
@@ -474,6 +482,7 @@ impl SessionState {
         Ok(())
     }
 
+    #[cfg(test)]
     pub(crate) fn set_receiver_chain_key(
         &mut self,
         sender: &PublicKey,
@@ -598,6 +607,39 @@ impl SessionState {
             .map(|pending| &pending.ciphertext)
     }
 
+    /// Extract the Double Ratchet state from this session.
+    ///
+    /// This **moves** the receiver chains out of the session protobuf,
+    /// leaving them empty. The caller must either apply the ratchet state
+    /// back via [`apply_ratchet_state`](Self::apply_ratchet_state) or
+    /// discard this `SessionState` entirely. Reading receiver chains from
+    /// a taken session is invalid and will produce incorrect results.
+    ///
+    /// The caller must supply `self_session` (whether this is a note-to-self
+    /// session) since it requires identity key comparison, which is above the
+    /// ratchet layer.
+    pub(crate) fn take_ratchet_state(
+        &mut self,
+        self_session: bool,
+    ) -> crate::error::Result<crate::double_ratchet::RatchetState> {
+        let receiver_chains = std::mem::take(&mut self.session.receiver_chains);
+        Ok(crate::double_ratchet::RatchetState::from_pb(
+            &self.session,
+            self_session,
+            receiver_chains,
+        )?)
+    }
+
+    /// Write modified Double Ratchet state back into this session.
+    ///
+    /// Only the ratchet-relevant fields are updated; all other session
+    /// fields (identity keys, pending pre-key, SPQR state, etc.) are
+    /// left unchanged.
+    pub(crate) fn apply_ratchet_state(&mut self, ratchet: crate::double_ratchet::RatchetState) {
+        ratchet.apply_to_pb(&mut self.session);
+    }
+
+    #[cfg(test)]
     pub(crate) fn pq_ratchet_recv(
         &mut self,
         msg: &spqr::SerializedMessage,
@@ -608,6 +650,7 @@ impl SessionState {
         Ok(key)
     }
 
+    #[cfg(test)]
     pub(crate) fn pq_ratchet_send<R: Rng + CryptoRng>(
         &mut self,
         csprng: &mut R,
@@ -621,6 +664,19 @@ impl SessionState {
     pub(crate) fn pq_ratchet_state(&self) -> &spqr::SerializedState {
         &self.session.pq_ratchet_state
     }
+
+    /// Move the PQ ratchet state out, leaving an empty vec in its place.
+    ///
+    /// The caller must either write a new PQ ratchet state back via
+    /// [`set_pq_ratchet_state`](Self::set_pq_ratchet_state) or discard
+    /// this `SessionState` entirely.
+    pub(crate) fn take_pq_ratchet_state(&mut self) -> spqr::SerializedState {
+        std::mem::take(&mut self.session.pq_ratchet_state)
+    }
+
+    pub(crate) fn set_pq_ratchet_state(&mut self, state: spqr::SerializedState) {
+        self.session.pq_ratchet_state = state;
+    }
 }
 
 impl From<SessionStructure> for SessionState {

rust/protocol/src/triple_ratchet.rs

@@ -0,0 +1,311 @@
+//
+// Copyright 2026 Signal Messenger, LLC.
+// SPDX-License-Identifier: AGPL-3.0-only
+//
+
+//! The Triple Ratchet: Double Ratchet + SPQR combined into a single
+//! encrypt/decrypt interface.
+//!
+//! Handles MAC computation/verification and AES-CBC encryption/decryption.
+//!
+//! The session management layer treats this as an opaque box: give it
+//! plaintext, get ciphertext; give it ciphertext, get plaintext. Identity
+//! checking, session selection, pre-key handling, and storage are NOT
+//! this layer's concern.
+
+use rand::{CryptoRng, Rng};
+
+use crate::double_ratchet::RatchetState;
+use crate::ratchet::ChainKey;
+use crate::session_management::CurrentOrPrevious;
+use crate::state::SessionState;
+use crate::{
+    CiphertextMessageType, IdentityKey, KeyPair, ProtocolAddress, Result, SignalMessage,
+    SignalProtocolError,
+};
+
+/// Sender-side Triple Ratchet session.
+///
+/// This is intentionally narrower than [`TripleRatchet`]: encrypt only
+/// depends on the current sender chain, SPQR state, identities, and metadata.
+/// It does not deserialize receiver chains, so corrupt cold receiver-chain
+/// state does not block sending.
+///
+/// # Ownership contract
+///
+/// [`from_session_state`](Self::from_session_state) **moves** the PQ ratchet
+/// state out of the session (via [`SessionState::take_pq_ratchet_state`]),
+/// leaving it temporarily invalid. The caller must either call
+/// [`apply_to_session_state`](Self::apply_to_session_state) on success, or
+/// discard the `SessionState` entirely.
+pub(crate) struct OutgoingTripleRatchet {
+    sender_ratchet_key: KeyPair,
+    sender_chain_key: ChainKey,
+    previous_counter: u32,
+    pqr_state: spqr::SerializedState,
+    session_version: u8,
+    local_identity_key: IdentityKey,
+    remote_identity_key: IdentityKey,
+}
+
+impl OutgoingTripleRatchet {
+    pub(crate) fn from_session_state(state: &mut SessionState) -> Result<Self> {
+        let sender_ratchet_key = KeyPair {
+            public_key: state.sender_ratchet_key()?,
+            private_key: state.sender_ratchet_private_key()?,
+        };
+        let sender_chain_key = state.get_sender_chain_key()?;
+        let pqr_state = state.take_pq_ratchet_state();
+        let session_version: u8 = state.session_version()?.try_into().map_err(|_| {
+            SignalProtocolError::InvalidSessionStructure("version does not fit in u8")
+        })?;
+        let local_identity_key = state.local_identity_key()?;
+        let remote_identity_key =
+            state
+                .remote_identity_key()?
+                .ok_or(SignalProtocolError::InvalidSessionStructure(
+                    "missing remote identity key",
+                ))?;
+
+        Ok(Self {
+            sender_ratchet_key,
+            sender_chain_key,
+            previous_counter: state.previous_counter(),
+            pqr_state,
+            session_version,
+            local_identity_key,
+            remote_identity_key,
+        })
+    }
+
+    pub(crate) fn apply_to_session_state(self, state: &mut SessionState) {
+        state.set_sender_chain_key(&self.sender_chain_key);
+        state.set_pq_ratchet_state(self.pqr_state);
+    }
+
+    pub(crate) fn encrypt<R: Rng + CryptoRng>(
+        &mut self,
+        plaintext: &[u8],
+        local_address: Option<&ProtocolAddress>,
+        remote_address: &ProtocolAddress,
+        csprng: &mut R,
+    ) -> Result<SignalMessage> {
+        let spqr::Send {
+            state: new_pqr_state,
+            key: pqr_key,
+            msg: pqr_msg,
+        } = spqr::send(&self.pqr_state, csprng).map_err(|e| {
+            SignalProtocolError::InvalidState(
+                "encrypt",
+                format!("post-quantum ratchet send error: {e}"),
+            )
+        })?;
+
+        let message_keys = self.sender_chain_key.message_keys().generate_keys(pqr_key);
+
+        let ctext = signal_crypto::aes_256_cbc_encrypt(
+            plaintext,
+            message_keys.cipher_key(),
+            message_keys.iv(),
+        )
+        .map_err(|_| {
+            log::error!("session state corrupt for {remote_address}");
+            SignalProtocolError::InvalidSessionStructure("invalid sender chain message keys")
+        })?;
+
+        let addresses = local_address.map(|addr| (addr, remote_address));
+
+        let message = SignalMessage::new(
+            self.session_version,
+            message_keys.mac_key(),
+            addresses,
+            self.sender_ratchet_key.public_key,
+            self.sender_chain_key.index(),
+            self.previous_counter,
+            &ctext,
+            &self.local_identity_key,
+            &self.remote_identity_key,
+            &pqr_msg,
+        )?;
+
+        self.sender_chain_key = self.sender_chain_key.next_chain_key();
+        self.pqr_state = new_pqr_state;
+
+        Ok(message)
+    }
+
+    pub(crate) fn session_version(&self) -> u8 {
+        self.session_version
+    }
+
+    pub(crate) fn local_identity_key(&self) -> &IdentityKey {
+        &self.local_identity_key
+    }
+}
+
+/// A Triple Ratchet session combining Double Ratchet and SPQR.
+///
+/// Constructed from a [`SessionState`], this extracts the cryptographic
+/// state needed for decrypt into typed fields. After a successful
+/// operation, call [`apply_to_session_state`](Self::apply_to_session_state)
+/// to write the updated state back.
+///
+/// # Ownership contract
+///
+/// [`from_session_state`](Self::from_session_state) **moves** the receiver
+/// chains and PQ ratchet state out of the session, leaving it temporarily invalid.
+/// The caller must either:
+/// - Call [`apply_to_session_state`](Self::apply_to_session_state) on success, or
+/// - Discard the `SessionState` (e.g., it was a clone for trial decrypt).
+pub(crate) struct TripleRatchet {
+    ratchet: RatchetState,
+    pqr_state: spqr::SerializedState,
+    local_identity_key: IdentityKey,
+    remote_identity_key: IdentityKey,
+}
+
+impl TripleRatchet {
+    /// Construct from a [`SessionState`] by extracting crypto state.
+    ///
+    /// This moves receiver chains and PQ ratchet state out of `state`.
+    /// See the [ownership contract](Self#ownership-contract) for details.
+    ///
+    /// Fails if the session is missing required fields (root key, identity
+    /// keys, etc.). The caller should map the error appropriately for the
+    /// context (e.g., "no session available to decrypt").
+    pub(crate) fn from_session_state(state: &mut SessionState) -> Result<Self> {
+        let self_session = state.session_with_self()?;
+        let ratchet = state.take_ratchet_state(self_session)?;
+        let pqr_state = state.take_pq_ratchet_state();
+        let local_identity_key = state.local_identity_key()?;
+        let remote_identity_key =
+            state
+                .remote_identity_key()?
+                .ok_or(SignalProtocolError::InvalidSessionStructure(
+                    "missing remote identity key",
+                ))?;
+
+        Ok(Self {
+            ratchet,
+            pqr_state,
+            local_identity_key,
+            remote_identity_key,
+        })
+    }
+
+    /// Write the updated crypto state back to a [`SessionState`].
+    ///
+    /// Only call this after a successful decrypt — this is how we ensure
+    /// no state pollution on failure.
+    pub(crate) fn apply_to_session_state(self, state: &mut SessionState) {
+        state.apply_ratchet_state(self.ratchet);
+        state.set_pq_ratchet_state(self.pqr_state);
+    }
+
+    // -- Decrypt -------------------------------------------------------
+
+    /// Decrypt a [`SignalMessage`] to plaintext.
+    ///
+    /// Performs DR chain key derivation, SPQR key derivation, MAC
+    /// verification, and AES-CBC decryption. Ratchet and SPQR state are
+    /// only committed on success — a failed MAC or decryption leaves this
+    /// session unchanged.
+    ///
+    /// `original_message_type` is used for error classification only
+    /// (PreKey vs Whisper).
+    pub(crate) fn decrypt<R: Rng + CryptoRng>(
+        &mut self,
+        sender_address: &ProtocolAddress,
+        recipient_address: Option<&ProtocolAddress>,
+        ciphertext: &SignalMessage,
+        original_message_type: CiphertextMessageType,
+        current_or_previous_for_logging: CurrentOrPrevious,
+        csprng: &mut R,
+    ) -> Result<Vec<u8>> {
+        // DR: ensure we have a receiver chain, then consume the message key
+        let their_ephemeral = ciphertext.sender_ratchet_key();
+        let counter = ciphertext.counter();
+        let chain_key = self
+            .ratchet
+            .ensure_receiver_chain(their_ephemeral, csprng)?;
+        let message_key_gen = self.ratchet.consume_message_key(
+            their_ephemeral,
+            chain_key,
+            counter,
+            original_message_type,
+            &sender_address.to_string(),
+        )?;
+
+        // SPQR recv — compute key but don't commit state yet
+        let spqr::Recv {
+            state: new_pqr_state,
+            key: pqr_key,
+        } = spqr::recv(&self.pqr_state, ciphertext.pq_ratchet()).map_err(|e| match e {
+            spqr::Error::StateDecode => SignalProtocolError::InvalidState(
+                "decrypt",
+                format!("post-quantum ratchet error: {e}"),
+            ),
+            _ => {
+                log::info!("post-quantum ratchet error in decrypt: {e}");
+                SignalProtocolError::InvalidMessage(
+                    original_message_type,
+                    "post-quantum ratchet error",
+                )
+            }
+        })?;
+
+        // Derive final message keys by mixing DR chain key with SPQR key
+        let message_keys = message_key_gen.generate_keys(pqr_key);
+
+        // MAC verification
+        let mac_valid = match recipient_address {
+            Some(recipient_address) => ciphertext.verify_mac_with_addresses(
+                sender_address,
+                recipient_address,
+                &self.remote_identity_key,
+                &self.local_identity_key,
+                message_keys.mac_key(),
+            )?,
+            None => ciphertext.verify_mac(
+                &self.remote_identity_key,
+                &self.local_identity_key,
+                message_keys.mac_key(),
+            )?,
+        };
+        if !mac_valid {
+            return Err(SignalProtocolError::InvalidMessage(
+                original_message_type,
+                "MAC verification failed",
+            ));
+        }
+
+        // AES-CBC decrypt
+        let ptext = match signal_crypto::aes_256_cbc_decrypt(
+            ciphertext.body(),
+            message_keys.cipher_key(),
+            message_keys.iv(),
+        ) {
+            Ok(ptext) => ptext,
+            Err(signal_crypto::DecryptionError::BadKeyOrIv) => {
+                log::warn!(
+                    "{current_or_previous_for_logging} session state corrupt for {sender_address}",
+                );
+                return Err(SignalProtocolError::InvalidSessionStructure(
+                    "invalid receiver chain message keys",
+                ));
+            }
+            Err(signal_crypto::DecryptionError::BadCiphertext(msg)) => {
+                log::warn!("failed to decrypt 1:1 message: {msg}");
+                return Err(SignalProtocolError::InvalidMessage(
+                    original_message_type,
+                    "failed to decrypt",
+                ));
+            }
+        };
+
+        // Commit SPQR state only after all verification passed
+        self.pqr_state = new_pqr_state;
+
+        Ok(ptext)
+    }
+}

rust/protocol/tests/ratchet.rs

@@ -51,13 +51,12 @@ fn test_alice_and_bob_agree_on_chain_keys_with_kyber() -> Result<(), SignalProto
         bob_identity_key_pair,
         bob_signed_pre_key_pair,
         None,
-        bob_ephemeral_key_pair,
         bob_kyber_pre_key_pair,
         *alice_identity_key_pair.identity_key(),
         alice_base_key_pair.public_key,
         &kyber_ciphertext,
     );
-    let bob_record = initialize_bob_session_record(&bob_parameters)?;
+    let bob_record = initialize_bob_session_record(&bob_parameters, &bob_ephemeral_key_pair)?;
 
     assert_eq!(
         KYBER_AWARE_MESSAGE_VERSION,
@@ -132,14 +131,13 @@ fn test_bob_rejects_torsioned_basekey() -> Result<(), SignalProtocolError> {
         bob_identity_key_pair,
         bob_signed_pre_key_pair,
         None,
-        bob_ephemeral_key_pair,
         bob_kyber_pre_key_pair,
         *alice_identity_key_pair.identity_key(),
         tweaked_alice_base_key,
         &kyber_ciphertext,
     );
 
-    assert!(initialize_bob_session_record(&bob_parameters).is_err());
+    assert!(initialize_bob_session_record(&bob_parameters, &bob_ephemeral_key_pair).is_err());
 
     Ok(())
 }
@@ -194,14 +192,13 @@ fn test_bob_rejects_highbit_basekey() -> Result<(), SignalProtocolError> {
         bob_identity_key_pair,
         bob_signed_pre_key_pair,
         None,
-        bob_ephemeral_key_pair,
         bob_kyber_pre_key_pair,
         *alice_identity_key_pair.identity_key(),
         tweaked_alice_base_key,
         &kyber_ciphertext,
     );
 
-    assert!(initialize_bob_session_record(&bob_parameters).is_err());
+    assert!(initialize_bob_session_record(&bob_parameters, &bob_ephemeral_key_pair).is_err());
 
     Ok(())
 }

rust/protocol/tests/support/mod.rs

@@ -177,14 +177,13 @@ pub fn initialize_sessions_v4() -> Result<(SessionRecord, SessionRecord), Signal
         bob_identity,
         bob_base_key,
         None,
-        bob_ephemeral_key,
         bob_kyber_key,
         *alice_identity.identity_key(),
         alice_base_key.public_key,
         &kyber_ciphertext,
     );
 
-    let bob_session = initialize_bob_session_record(&bob_params)?;
+    let bob_session = initialize_bob_session_record(&bob_params, &bob_ephemeral_key)?;
 
     Ok((alice_session, bob_session))
 }