openmls 0.9.2

//! ### Don't Panic!
//!
//! Functions in this module should never panic. However, if there is a bug in
//! the implementation, a function will return an unrecoverable `LibraryError`.
//! This means that some functions that are not expected to fail and throw an
//! error, will still return a `Result` since they may throw a `LibraryError`.

use std::collections::VecDeque;

use openmls_traits::crypto::OpenMlsCrypto;

use openmls_traits::types::Ciphersuite;

use crate::ciphersuite::AeadNonce;
#[cfg(feature = "virtual-clients-draft")]
use crate::tree::dual_use_ratchet::DualUseRatchet;
use crate::tree::secret_tree::*;

use super::*;

/// The generation of a given [`SenderRatchet`].
pub(crate) type Generation = u32;
/// Stores the configuration parameters for `DecryptionRatchet`s.
///
/// **Parameters**
///
/// - out_of_order_tolerance:
///   This parameter defines a window for which decryption secrets are kept.
///   This is useful in case the DS cannot guarantee that all application messages have total order within an epoch.
///   Use this carefully, since keeping decryption secrets affects forward secrecy within an epoch.
///   The default value is 5.
/// - maximum_forward_distance:
///   This parameter defines how many incoming messages can be skipped. This is useful if the DS
///   drops application messages. The default value is 1000.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SenderRatchetConfiguration {
    out_of_order_tolerance: Generation,
    maximum_forward_distance: Generation,
}

impl SenderRatchetConfiguration {
    /// Create a new configuration
    pub fn new(out_of_order_tolerance: Generation, maximum_forward_distance: Generation) -> Self {
        Self {
            out_of_order_tolerance,
            maximum_forward_distance,
        }
    }
    /// Get a reference to the sender ratchet configuration's out of order tolerance.
    pub fn out_of_order_tolerance(&self) -> Generation {
        self.out_of_order_tolerance
    }

    /// Get a reference to the sender ratchet configuration's maximum forward distance.
    pub fn maximum_forward_distance(&self) -> Generation {
        self.maximum_forward_distance
    }
}

impl Default for SenderRatchetConfiguration {
    fn default() -> Self {
        Self::new(5, 1000)
    }
}

/// The key material derived from a [`RatchetSecret`] meant for use with a
/// nonce-based symmetric encryption scheme.
pub(crate) type RatchetKeyMaterial = (AeadKey, AeadNonce);

/// A ratchet that can output key material either for encryption
/// ([`EncryptionRatchet`](SenderRatchet)) or decryption
/// ([`DecryptionRatchet`]). A [`DecryptionRatchet`] can be configured with an
/// `out_of_order_tolerance` and a `maximum_forward_distance` (see
/// [`SenderRatchetConfiguration`]) while an Encryption Ratchet never keeps past
/// secrets around. With the `virtual-clients-draft` feature, own sender
/// ratchets are
/// [`DualUseRatchet`](super::dual_use_ratchet::DualUseRatchet)s, which can
/// output key material for both encryption and decryption.
#[derive(Serialize, Deserialize)]
#[serde(from = "SerdeSenderRatchet")]
#[cfg_attr(any(feature = "test-utils", test), derive(PartialEq, Clone))]
#[cfg_attr(any(feature = "crypto-debug", test), derive(Debug))]
pub(crate) enum SenderRatchet {
    #[cfg_attr(feature = "virtual-clients-draft", allow(unused))]
    EncryptionRatchet(RatchetSecret),
    DecryptionRatchet(DecryptionRatchet),
    #[cfg(feature = "virtual-clients-draft")]
    DualUse(DualUseRatchet),
}

/// Wire representation of [`SenderRatchet`] used to deserialize it via
/// `#[serde(from = ...)]`. Unlike [`SenderRatchet`], the `DualUse` variant
/// here is always present (its payload just changes shape depending on the
/// `virtual-clients-draft` feature), so a value persisted by a build with the
/// feature enabled can still be read by a build with the feature disabled,
/// and vice versa. This avoids two failure modes of deserializing directly
/// into [`SenderRatchet`]: an own ratchet persisted as `EncryptionRatchet`
/// before the feature was turned on would otherwise never become a
/// `DualUse` ratchet, and an own ratchet persisted as `DualUse` would
/// otherwise fail to deserialize at all once the feature is turned off.
#[derive(Deserialize)]
enum SerdeSenderRatchet {
    EncryptionRatchet(RatchetSecret),
    DecryptionRatchet(DecryptionRatchet),
    DualUse(SerdeDualUseRatchet),
}

#[cfg(feature = "virtual-clients-draft")]
type SerdeDualUseRatchet = DualUseRatchet;

#[cfg(not(feature = "virtual-clients-draft"))]
#[derive(Deserialize)]
struct SerdeDualUseRatchet {
    ratchet_head: RatchetSecret,
}

impl From<SerdeSenderRatchet> for SenderRatchet {
    fn from(repr: SerdeSenderRatchet) -> Self {
        match repr {
            SerdeSenderRatchet::EncryptionRatchet(ratchet_secret) => {
                #[cfg(feature = "virtual-clients-draft")]
                {
                    SenderRatchet::DualUse(DualUseRatchet::from(ratchet_secret))
                }
                #[cfg(not(feature = "virtual-clients-draft"))]
                {
                    SenderRatchet::EncryptionRatchet(ratchet_secret)
                }
            }
            SerdeSenderRatchet::DecryptionRatchet(dec_ratchet) => {
                SenderRatchet::DecryptionRatchet(dec_ratchet)
            }
            SerdeSenderRatchet::DualUse(dual_ratchet) => {
                #[cfg(feature = "virtual-clients-draft")]
                {
                    SenderRatchet::DualUse(dual_ratchet)
                }
                #[cfg(not(feature = "virtual-clients-draft"))]
                {
                    SenderRatchet::EncryptionRatchet(dual_ratchet.ratchet_head)
                }
            }
        }
    }
}

impl SenderRatchet {
    #[cfg(test)]
    pub(crate) fn generation(&self) -> Generation {
        match self {
            SenderRatchet::EncryptionRatchet(enc_ratchet) => enc_ratchet.generation(),
            SenderRatchet::DecryptionRatchet(dec_ratchet) => dec_ratchet.generation(),
            #[cfg(feature = "virtual-clients-draft")]
            SenderRatchet::DualUse(dual_ratchet) => dual_ratchet.generation(),
        }
    }
}

/// The core of both types of [`SenderRatchet`]. It contains the current head of
/// the ratchet chain, as well as its current [`Generation`]. It can be
/// initialized with a given secret and then ratcheted forward, outputting
/// [`RatchetKeyMaterial`] and increasing its [`Generation`] each time.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[cfg_attr(any(feature = "test-utils", test), derive(PartialEq))]
pub(crate) struct RatchetSecret {
    secret: Secret,
    generation: Generation,
}

impl RatchetSecret {
    /// Create an initial [`RatchetSecret`] with `generation = 0` from the given
    /// [`Secret`].
    pub(crate) fn initial_ratchet_secret(secret: Secret) -> Self {
        Self {
            secret,
            generation: 0,
        }
    }

    /// Return the generation of this [`RatchetSecret`].
    pub(crate) fn generation(&self) -> Generation {
        self.generation
    }

    /// Consume this [`RatchetSecret`] to derive a pair of [`RatchetSecrets`],
    /// as well as the [`RatchetSecret`] of the next generation and return both.
    pub(crate) fn ratchet_forward(
        &mut self,
        crypto: &impl OpenMlsCrypto,
        ciphersuite: Ciphersuite,
    ) -> Result<(Generation, RatchetKeyMaterial), SecretTreeError> {
        log::trace!("Ratcheting forward in generation {}.", self.generation);
        log_crypto!(trace, "    with secret {:x?}", self.secret);

        // Check if the generation is getting too large.
        if self.generation == u32::MAX {
            return Err(SecretTreeError::RatchetTooLong);
        }
        let nonce = derive_tree_secret(
            ciphersuite,
            &self.secret,
            "nonce",
            self.generation,
            ciphersuite.aead_nonce_length(),
            crypto,
        )?;
        let key = derive_tree_secret(
            ciphersuite,
            &self.secret,
            "key",
            self.generation,
            ciphersuite.aead_key_length(),
            crypto,
        )?;
        self.secret = derive_tree_secret(
            ciphersuite,
            &self.secret,
            "secret",
            self.generation,
            ciphersuite.hash_length(),
            crypto,
        )?;
        let generation = self.generation;
        self.generation += 1;
        Ok((
            generation,
            (
                AeadKey::from_secret(key, ciphersuite),
                AeadNonce::from_secret(nonce),
            ),
        ))
    }

    #[cfg(test)]
    pub(crate) fn set_generation(&mut self, generation: Generation) {
        self.generation = generation
    }
}

/// [`SenderRatchet`] used to derive key material for decryption. It keeps the
/// [`RatchetKeyMaterial`] of epochs around until they are retrieved. This
/// behaviour can be configured via the `out_of_order_tolerance` and
/// `maximum_forward_distance` of the given [`SenderRatchetConfiguration`].
#[derive(Serialize, Deserialize)]
#[cfg_attr(any(feature = "test-utils", test), derive(PartialEq, Clone))]
#[cfg_attr(any(feature = "crypto-debug", test), derive(Debug))]
pub struct DecryptionRatchet {
    past_secrets: VecDeque<Option<RatchetKeyMaterial>>,
    ratchet_head: RatchetSecret,
}

impl DecryptionRatchet {
    /// Creates e new SenderRatchet
    pub(crate) fn new(secret: Secret) -> Self {
        Self {
            past_secrets: VecDeque::new(),
            ratchet_head: RatchetSecret::initial_ratchet_secret(secret),
        }
    }

    /// Get the generation of the ratchet head.
    #[cfg(test)]
    pub(crate) fn generation(&self) -> Generation {
        self.ratchet_head.generation()
    }

    #[cfg(test)]
    pub(crate) fn ratchet_secret_mut(&mut self) -> &mut RatchetSecret {
        &mut self.ratchet_head
    }

    /// Returns the secret for `generation` without changing the ratchet,
    /// together with the change that [`Self::consume`] applies to mark the
    /// generation as used. Returns an error if the generation is out of bound
    /// or was used before.
    pub(crate) fn peek(
        &self,
        ciphersuite: Ciphersuite,
        crypto: &impl OpenMlsCrypto,
        generation: Generation,
        configuration: &SenderRatchetConfiguration,
    ) -> Result<(RatchetKeyMaterial, PeekedGeneration), SecretTreeError> {
        log::debug!("peek decryption secret");
        let head_generation = self.ratchet_head.generation();
        // If generation is too distant in the future
        if head_generation < u32::MAX - configuration.maximum_forward_distance()
            && generation > head_generation + configuration.maximum_forward_distance()
        {
            return Err(SecretTreeError::TooDistantInTheFuture);
        }
        // If generation is too distant in the past
        if generation < head_generation
            && (head_generation - generation) > configuration.out_of_order_tolerance()
        {
            log::error!("  Generation is too far in the past (broke out of order tolerance ({}) {generation} < {head_generation}).", configuration.out_of_order_tolerance());
            return Err(SecretTreeError::TooDistantInThePast);
        }
        // If generation is the one the ratchet is currently at or in the future
        if generation >= head_generation {
            let mut ratchet_head = self.ratchet_head.clone();
            // Newest first. Only the newest `out_of_order_tolerance` skipped
            // secrets can remain in the window after consuming.
            let mut skipped_secrets = VecDeque::new();
            for _ in 0..(generation - head_generation) {
                let ratchet_secrets = ratchet_head
                    .ratchet_forward(crypto, ciphersuite)
                    .map(|(_, key_material)| key_material)?;
                skipped_secrets.push_front(ratchet_secrets);
                skipped_secrets.truncate(configuration.out_of_order_tolerance() as usize);
            }
            let ratchet_secrets = ratchet_head
                .ratchet_forward(crypto, ciphersuite)
                .map(|(_, key_material)| key_material)?;
            Ok((
                ratchet_secrets,
                PeekedGeneration::Forward {
                    from_head: head_generation,
                    ratchet_head,
                    skipped_secrets,
                },
            ))
        } else {
            let index = self.past_secret_index(generation)?;
            let ratchet_secrets = self
                .past_secrets
                .get(index)
                .ok_or(SecretTreeError::IndexOutOfBounds)?
                .clone()
                // If the requested generation was used to decrypt a message
                // earlier, throw an error.
                .ok_or(SecretTreeError::SecretReuseError)?;
            Ok((ratchet_secrets, PeekedGeneration::Past { generation }))
        }
    }

    /// Marks a generation returned by [`Self::peek`] as used. Afterwards, the
    /// generation's secret is deleted and the ratchet has moved forward as
    /// far as needed. Returns an error if the ratchet changed in a way that
    /// makes `peeked` stale.
    pub(crate) fn consume(
        &mut self,
        peeked: PeekedGeneration,
        configuration: &SenderRatchetConfiguration,
    ) -> Result<(), SecretTreeError> {
        match peeked {
            PeekedGeneration::Forward {
                from_head,
                ratchet_head,
                skipped_secrets,
            } => {
                if from_head != self.ratchet_head.generation() {
                    log::error!("The ratchet moved since the generation was peeked.");
                    return Err(SecretTreeError::LibraryError);
                }
                self.ratchet_head = ratchet_head;
                for ratchet_secrets in skipped_secrets.into_iter().rev() {
                    self.past_secrets.push_front(Some(ratchet_secrets));
                }
                // Add an entry to the past secrets queue to keep indexing consistent.
                self.past_secrets.push_front(None);
                self.past_secrets
                    .truncate(configuration.out_of_order_tolerance() as usize);
                Ok(())
            }
            PeekedGeneration::Past { generation } => {
                let index = self.past_secret_index(generation)?;
                self.past_secrets
                    .get_mut(index)
                    .ok_or(SecretTreeError::IndexOutOfBounds)?
                    // We use take here to replace the entry in the
                    // `past_secrets` with `None`, thus achieving FS for that
                    // secret.
                    .take()
                    .map(|_| ())
                    .ok_or(SecretTreeError::SecretReuseError)
            }
        }
    }

    /// Gets a secret from the SenderRatchet and marks it as used. Returns an
    /// error if the generation is out of bound.
    #[cfg(test)]
    pub(crate) fn secret_for_decryption(
        &mut self,
        ciphersuite: Ciphersuite,
        crypto: &impl OpenMlsCrypto,
        generation: Generation,
        configuration: &SenderRatchetConfiguration,
    ) -> Result<RatchetKeyMaterial, SecretTreeError> {
        let (ratchet_secrets, peeked) =
            self.peek(ciphersuite, crypto, generation, configuration)?;
        self.consume(peeked, configuration)?;
        Ok(ratchet_secrets)
    }

    /// Returns the index of a generation below the head in the past secrets
    /// queue.
    fn past_secret_index(&self, generation: Generation) -> Result<usize, SecretTreeError> {
        // If the requested generation is within the window of past secrets,
        // we should get a positive index.
        let window_index = i64::from(self.ratchet_head.generation()) - i64::from(generation) - 1;
        // We might not have the key material (e.g. we might have discarded
        // it when generating an encryption secret).
        usize::try_from(window_index).map_err(|_| {
            log::error!("  Generation is too far in the past (not in the window).");
            SecretTreeError::TooDistantInThePast
        })
    }
}

/// The change to a [`DecryptionRatchet`] that marks one generation as used,
/// computed by [`DecryptionRatchet::peek`].
pub(crate) enum PeekedGeneration {
    /// The generation is at or after the head. The ratchet moves its head to
    /// `ratchet_head` and keeps the skipped secrets that fit into its window.
    Forward {
        from_head: Generation,
        ratchet_head: RatchetSecret,
        skipped_secrets: VecDeque<RatchetKeyMaterial>,
    },
    /// The generation is in the window of past secrets.
    Past { generation: Generation },
}

#[cfg(test)]
mod persistence_tests {
    use super::*;

    // Reading an `EncryptionRatchet` as `DualUse` when enabling the `virtual-clients-draft` feature.
    #[cfg(feature = "virtual-clients-draft")]
    #[test]
    fn pre_feature_encryption_ratchet_is_promoted_to_dual_use() {
        let ratchet_head = RatchetSecret::initial_ratchet_secret(Secret::from_slice(&[0; 32]));
        let restored: SenderRatchet =
            SerdeSenderRatchet::EncryptionRatchet(ratchet_head.clone()).into();
        match restored {
            SenderRatchet::DualUse(dual_use) => {
                assert_eq!(dual_use.generation(), ratchet_head.generation())
            }
            _ => panic!("expected promotion to DualUse"),
        }
    }

    // A `DualUse` ratchet persisted when using the `virtual-clients-draft` that should
    // be downgraded to a `EncryptionRatchet` when not using the feature.
    #[cfg(not(feature = "virtual-clients-draft"))]
    #[test]
    fn persisted_dual_use_is_readable_without_the_feature() {
        let ratchet_head = RatchetSecret::initial_ratchet_secret(Secret::from_slice(&[0; 32]));
        let restored: SenderRatchet = SerdeSenderRatchet::DualUse(SerdeDualUseRatchet {
            ratchet_head: ratchet_head.clone(),
        })
        .into();
        match restored {
            SenderRatchet::EncryptionRatchet(rs) => {
                assert_eq!(rs.generation(), ratchet_head.generation())
            }
            _ => panic!("expected fallback to EncryptionRatchet"),
        }
    }
}