Skip to main content

openmls/treesync/
mod.rs

1//! This module implements the ratchet tree component of MLS.
2//!
3//! It exposes the [`Node`] enum that can contain either a [`LeafNode`] or a [`ParentNode`].
4
5// # Internal documentation
6//
7// This module provides the [`TreeSync`] struct, which contains the state
8// shared between a group of MLS clients in the shape of a tree, where each
9// non-blank leaf corresponds to one group member. The functions provided by
10// its implementation allow the creation of a [`TreeSyncDiff`] instance, which
11// in turn can be mutably operated on and merged back into the original
12// [`TreeSync`] instance.
13//
14// The submodules of this module define the nodes of the tree (`nodes`),
15// helper functions and structs for the algorithms used to sync the tree across
16// the group ([`hashes`]) and the diff functionality ([`diff`]).
17//
18// Finally, this module contains the [`treekem`] module, which allows the
19// encryption and decryption of updates to the tree.
20
21#[cfg(any(feature = "test-utils", test))]
22use std::fmt;
23
24use openmls_traits::{
25    crypto::OpenMlsCrypto,
26    signatures::Signer,
27    types::{Ciphersuite, CryptoError},
28};
29use serde::{Deserialize, Serialize};
30use thiserror::Error;
31use tls_codec::{TlsDeserialize, TlsDeserializeBytes, TlsSerialize, TlsSize};
32
33use self::{
34    diff::{StagedTreeSyncDiff, TreeSyncDiff},
35    node::{
36        leaf_node::{
37            Capabilities, NewLeafNodeParams, TreeInfoTbs, TreePosition, VerifiableLeafNode,
38        },
39        NodeIn,
40    },
41    treesync_node::{TreeSyncLeafNode, TreeSyncNode, TreeSyncParentNode},
42};
43#[cfg(any(feature = "test-utils", test))]
44use crate::{binary_tree::array_representation::level, test_utils::bytes_to_hex};
45use crate::{
46    binary_tree::array_representation::ParentNodeIndex, treesync::node::leaf_node::LeafNodeIn,
47};
48use crate::{
49    binary_tree::{
50        array_representation::{is_node_in_tree, LeafNodeIndex, TreeSize},
51        MlsBinaryTree, MlsBinaryTreeError,
52    },
53    ciphersuite::{signable::Verifiable, Secret},
54    credentials::CredentialWithKey,
55    error::LibraryError,
56    extensions::Extensions,
57    group::{GroupId, Member},
58    key_packages::Lifetime,
59    messages::{PathSecret, PathSecretError},
60    schedule::CommitSecret,
61    storage::OpenMlsProvider,
62};
63
64// Private
65mod hashes;
66use errors::*;
67
68// Crate
69pub(crate) mod diff;
70pub(crate) mod node;
71pub(crate) mod treekem;
72pub(crate) mod treesync_node;
73
74use node::encryption_keys::EncryptionKeyPair;
75
76// Public
77pub mod errors;
78#[cfg(feature = "test-utils")]
79pub use node::encryption_keys::test_utils;
80pub use node::encryption_keys::EncryptionKey;
81
82// Public re-exports
83pub use node::{
84    leaf_node::{
85        LeafNode, LeafNodeParameters, LeafNodeParametersBuilder, LeafNodeSource,
86        LeafNodeUpdateError,
87    },
88    parent_node::ParentNode,
89    Node,
90};
91
92// Tests
93#[cfg(any(feature = "test-utils", test))]
94pub mod tests_and_kats;
95
96/// An exported ratchet tree as used in, e.g., [`GroupInfo`](crate::messages::group_info::GroupInfo).
97#[derive(PartialEq, Eq, Clone, Debug, Serialize, Deserialize, TlsSerialize, TlsSize)]
98pub struct RatchetTree(Vec<Option<Node>>);
99
100/// An error during processing of an incoming ratchet tree.
101#[derive(Error, Debug, PartialEq, Clone)]
102pub enum RatchetTreeError {
103    /// The ratchet tree is empty.
104    #[error("The ratchet tree has no nodes.")]
105    MissingNodes,
106    /// The ratchet tree has a trailing blank node.
107    #[error("The ratchet tree has trailing blank nodes.")]
108    TrailingBlankNodes,
109    /// Invalid node signature.
110    #[error("Invalid node signature.")]
111    InvalidNodeSignature,
112    /// Wrong node type.
113    #[error("Wrong node type.")]
114    WrongNodeType,
115}
116
117impl RatchetTree {
118    /// Create a [`RatchetTree`] from a vector of nodes stripping all trailing blank nodes.
119    ///
120    /// Note: The caller must ensure to call this with a vector that is *not* empty after removing all trailing blank nodes.
121    fn trimmed(mut nodes: Vec<Option<Node>>) -> Self {
122        // Remove all trailing blank nodes.
123        match nodes.iter().enumerate().rfind(|(_, node)| node.is_some()) {
124            Some((rightmost_nonempty_position, _)) => {
125                // We need to add 1 to `rightmost_nonempty_position` to keep the rightmost node.
126                nodes.resize(rightmost_nonempty_position + 1, None);
127            }
128            None => {
129                // If there is no rightmost non-blank node, the vector consist of blank nodes only.
130                nodes.clear();
131            }
132        }
133
134        debug_assert!(!nodes.is_empty(), "Caller should have ensured that `RatchetTree::trimmed` is not called with a vector that is empty after removing all trailing blank nodes.");
135        Self(nodes)
136    }
137
138    /// Create a new [`RatchetTree`] from a vector of nodes.
139    pub(crate) fn try_from_nodes(
140        ciphersuite: Ciphersuite,
141        crypto: &impl OpenMlsCrypto,
142        nodes: Vec<Option<NodeIn>>,
143        group_id: &GroupId,
144    ) -> Result<Self, RatchetTreeError> {
145        // ValSem300: "Exported ratchet trees must not have trailing blank nodes."
146        //
147        // We can check this by only looking at the last node (if any).
148        match nodes.last() {
149            Some(None) => {
150                // The ratchet tree is not empty, i.e., has a last node, *but* the last node *is* blank.
151                Err(RatchetTreeError::TrailingBlankNodes)
152            }
153            None => {
154                // The ratchet tree is empty.
155                Err(RatchetTreeError::MissingNodes)
156            }
157            Some(Some(_)) => {
158                // The ratchet tree is not empty, i.e., has a last node, and the last node is not blank.
159
160                // Verify the nodes.
161                // https://validation.openmls.tech/#valn1407
162                let mut verified_nodes = Vec::new();
163                for (index, node) in nodes.into_iter().enumerate() {
164                    let verified_node = match (index % 2, node) {
165                        // Even indices must be leaf nodes.
166                        (0, Some(NodeIn::LeafNode(leaf_node))) => {
167                            let tree_position = TreePosition::new(
168                                group_id.clone(),
169                                LeafNodeIndex::new((index / 2) as u32),
170                            );
171                            let verifiable_leaf_node = leaf_node.into_verifiable_leaf_node();
172                            let signature_key = verifiable_leaf_node
173                                .signature_key()
174                                .clone()
175                                .into_signature_public_key_enriched(
176                                    ciphersuite.signature_algorithm(),
177                                );
178                            Some(Node::leaf_node(match verifiable_leaf_node {
179                                VerifiableLeafNode::KeyPackage(leaf_node) => leaf_node
180                                    .verify(crypto, &signature_key)
181                                    .map_err(|_| RatchetTreeError::InvalidNodeSignature)?,
182                                VerifiableLeafNode::Update(mut leaf_node) => {
183                                    leaf_node.add_tree_position(tree_position);
184                                    leaf_node
185                                        .verify(crypto, &signature_key)
186                                        .map_err(|_| RatchetTreeError::InvalidNodeSignature)?
187                                }
188                                VerifiableLeafNode::Commit(mut leaf_node) => {
189                                    leaf_node.add_tree_position(tree_position);
190                                    leaf_node
191                                        .verify(crypto, &signature_key)
192                                        .map_err(|_| RatchetTreeError::InvalidNodeSignature)?
193                                }
194                            }))
195                        }
196                        // Odd indices must be parent nodes.
197                        (1, Some(NodeIn::ParentNode(parent_node))) => {
198                            Some(Node::ParentNode(parent_node))
199                        }
200                        // Blank nodes.
201                        (_, None) => None,
202                        // All other cases are invalid.
203                        _ => {
204                            return Err(RatchetTreeError::WrongNodeType);
205                        }
206                    };
207                    verified_nodes.push(verified_node);
208                }
209                Ok(Self::trimmed(verified_nodes))
210            }
211        }
212    }
213
214    /// Returns an iterator over all nodes in the ratchet tree.
215    pub fn nodes(&self) -> impl Iterator<Item = &Node> {
216        self.0.iter().flatten()
217    }
218
219    /// Returns an iterator over all leaf nodes in the ratchet tree.
220    pub fn leaves(&self) -> impl Iterator<Item = &LeafNode> {
221        self.nodes().filter_map(|node| match node {
222            Node::LeafNode(leaf_node) => Some(&**leaf_node),
223            Node::ParentNode(_parent_node) => None,
224        })
225    }
226
227    /// Returns an iterator over all parent nodes in the ratchet tree.
228    pub fn parents(&self) -> impl Iterator<Item = &ParentNode> {
229        self.nodes().filter_map(|node| match node {
230            Node::ParentNode(parent_node) => Some(&**parent_node),
231            Node::LeafNode(_leaf_node) => None,
232        })
233    }
234}
235
236/// A ratchet tree made of unverified nodes. This is used for deserialization
237/// and verification.
238#[derive(
239    PartialEq,
240    Eq,
241    Clone,
242    Debug,
243    Serialize,
244    Deserialize,
245    TlsDeserialize,
246    TlsDeserializeBytes,
247    TlsSerialize,
248    TlsSize,
249)]
250pub struct RatchetTreeIn(Vec<Option<NodeIn>>);
251
252impl RatchetTreeIn {
253    /// Create a new [`RatchetTreeIn`] from a vector of nodes after verifying
254    /// the nodes.
255    pub fn into_verified(
256        self,
257        ciphersuite: Ciphersuite,
258        crypto: &impl OpenMlsCrypto,
259        group_id: &GroupId,
260    ) -> Result<RatchetTree, RatchetTreeError> {
261        RatchetTree::try_from_nodes(ciphersuite, crypto, self.0, group_id)
262    }
263
264    /// Returns an iterator over all nodes in the ratchet tree.
265    pub fn nodes(&self) -> impl Iterator<Item = &NodeIn> {
266        self.0.iter().flatten()
267    }
268
269    /// Returns an iterator over all leaf nodes in the ratchet tree.
270    pub fn leaves(&self) -> impl Iterator<Item = &LeafNodeIn> {
271        self.nodes().filter_map(|node| match node {
272            NodeIn::LeafNode(leaf_node) => Some(&**leaf_node),
273            NodeIn::ParentNode(_parent_node) => None,
274        })
275    }
276
277    /// Returns an iterator over all non-blank leaf nodes together with their
278    /// real [`LeafNodeIndex`], unlike [`Self::leaves`] whose positions are
279    /// compacted by skipping blank tree slots.
280    pub fn full_leaves(&self) -> impl Iterator<Item = (LeafNodeIndex, &LeafNodeIn)> {
281        self.0
282            .iter()
283            .enumerate()
284            .filter_map(|(node_index, slot)| match slot {
285                Some(NodeIn::LeafNode(leaf_node)) => {
286                    Some((LeafNodeIndex::new((node_index / 2) as u32), &**leaf_node))
287                }
288                _ => None,
289            })
290    }
291
292    /// Returns an iterator over all parent nodes in the ratchet tree.
293    pub fn parents(&self) -> impl Iterator<Item = &ParentNode> {
294        self.nodes().filter_map(|node| match node {
295            NodeIn::ParentNode(parent_node) => Some(&**parent_node),
296            NodeIn::LeafNode(_leaf_node) => None,
297        })
298    }
299
300    fn from_ratchet_tree(ratchet_tree: RatchetTree) -> Self {
301        let nodes = ratchet_tree
302            .0
303            .into_iter()
304            .map(|node| node.map(NodeIn::from))
305            .collect();
306        Self(nodes)
307    }
308
309    #[cfg(test)]
310    pub(crate) fn from_nodes(nodes: Vec<Option<NodeIn>>) -> Self {
311        Self(nodes)
312    }
313}
314
315impl From<RatchetTree> for RatchetTreeIn {
316    fn from(ratchet_tree: RatchetTree) -> Self {
317        RatchetTreeIn::from_ratchet_tree(ratchet_tree)
318    }
319}
320
321// The following `From` implementation breaks abstraction layers and MUST
322// NOT be made available outside of tests or "test-utils".
323#[cfg(any(feature = "test-utils", test))]
324impl From<RatchetTreeIn> for RatchetTree {
325    fn from(ratchet_tree_in: RatchetTreeIn) -> Self {
326        Self(
327            ratchet_tree_in
328                .0
329                .into_iter()
330                .map(|node| node.map(Node::from))
331                .collect(),
332        )
333    }
334}
335
336#[cfg(any(feature = "test-utils", test))]
337fn log2(x: u32) -> usize {
338    if x == 0 {
339        return 0;
340    }
341    (31 - x.leading_zeros()) as usize
342}
343
344#[cfg(any(feature = "test-utils", test))]
345pub(crate) fn root(size: u32) -> u32 {
346    (1 << log2(size)) - 1
347}
348
349#[cfg(any(feature = "test-utils", test))]
350impl fmt::Display for RatchetTree {
351    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
352        let factor = 3;
353        let nodes = &self.0;
354        let tree_size = nodes.len() as u32;
355
356        for (i, node) in nodes.iter().enumerate() {
357            let level = level(i as u32);
358            write!(f, "{i:04}")?;
359            if let Some(node) = node {
360                let (key_bytes, parent_hash_bytes) = match node {
361                    Node::LeafNode(leaf_node) => {
362                        write!(f, "\tL      ")?;
363                        let key_bytes = leaf_node.encryption_key().as_slice();
364                        let parent_hash_bytes = leaf_node
365                            .parent_hash()
366                            .map(bytes_to_hex)
367                            .unwrap_or_default();
368                        (key_bytes, parent_hash_bytes)
369                    }
370                    Node::ParentNode(parent_node) => {
371                        if root(tree_size) == i as u32 {
372                            write!(f, "\tP (*)  ")?;
373                        } else {
374                            write!(f, "\tP      ")?;
375                        }
376                        let key_bytes = parent_node.public_key().as_slice();
377                        let parent_hash_string = bytes_to_hex(parent_node.parent_hash());
378                        (key_bytes, parent_hash_string)
379                    }
380                };
381                write!(
382                    f,
383                    "PK: {}  PH: {} | ",
384                    bytes_to_hex(key_bytes),
385                    if !parent_hash_bytes.is_empty() {
386                        parent_hash_bytes
387                    } else {
388                        str::repeat("  ", 32)
389                    }
390                )?;
391
392                write!(f, "{}◼︎", str::repeat(" ", level * factor))?;
393            } else {
394                if root(tree_size) == i as u32 {
395                    write!(
396                        f,
397                        "\t_ (*)  PK: {}  PH: {} | ",
398                        str::repeat("__", 32),
399                        str::repeat("__", 32)
400                    )?;
401                } else {
402                    write!(
403                        f,
404                        "\t_      PK: {}  PH: {} | ",
405                        str::repeat("__", 32),
406                        str::repeat("__", 32)
407                    )?;
408                }
409
410                write!(f, "{}❑", str::repeat(" ", level * factor))?;
411            }
412            writeln!(f)?;
413        }
414
415        Ok(())
416    }
417}
418
419/// The [`TreeSync`] struct holds an `MlsBinaryTree` instance, which contains
420/// the state that is synced across the group, as well as the [`LeafNodeIndex`]
421/// pointing to the leaf of this group member and the current hash of the tree.
422///
423/// It follows the same pattern of tree and diff as the underlying
424/// `MlsBinaryTree`, where the [`TreeSync`] instance is immutable safe for
425/// merging a `TreeSyncDiff`, which can be created, staged and merged (see
426/// `TreeSyncDiff`).
427///
428/// [`TreeSync`] instance guarantee a few invariants that are checked upon
429/// creating a new instance from an imported set of nodes, as well as when
430/// merging a diff.
431#[derive(Debug, Serialize, Deserialize)]
432#[cfg_attr(any(test, feature = "test-utils"), derive(PartialEq, Clone))]
433pub struct TreeSync {
434    tree: MlsBinaryTree<TreeSyncLeafNode, TreeSyncParentNode>,
435    tree_hash: Vec<u8>,
436}
437
438impl TreeSync {
439    /// Create a new tree with an own leaf for the given credential.
440    ///
441    /// Returns the resulting [`TreeSync`] instance, as well as the
442    /// corresponding [`CommitSecret`].
443    pub(crate) fn new(
444        provider: &impl OpenMlsProvider,
445        signer: &impl Signer,
446        ciphersuite: Ciphersuite,
447        credential_with_key: CredentialWithKey,
448        life_time: Lifetime,
449        capabilities: Capabilities,
450        extensions: Extensions<LeafNode>,
451    ) -> Result<(Self, CommitSecret, EncryptionKeyPair), LibraryError> {
452        let new_leaf_node_params = NewLeafNodeParams {
453            ciphersuite,
454            credential_with_key,
455            // Creation of a group is considered to be from a key package.
456            leaf_node_source: LeafNodeSource::KeyPackage(life_time),
457            capabilities,
458            extensions,
459            tree_info_tbs: TreeInfoTbs::KeyPackage,
460        };
461        let (leaf, encryption_key_pair) = LeafNode::new(provider, signer, new_leaf_node_params)?;
462
463        let node = Node::leaf_node(leaf);
464        let path_secret: PathSecret = Secret::random(ciphersuite, provider.rand())
465            .map_err(LibraryError::unexpected_crypto_error)?
466            .into();
467        let commit_secret: CommitSecret = path_secret
468            .derive_path_secret(provider.crypto(), ciphersuite)?
469            .into();
470        let nodes = vec![TreeSyncNode::from(node).into()];
471        let tree = MlsBinaryTree::new(nodes)
472            .map_err(|_| LibraryError::custom("Unexpected error creating the binary tree."))?;
473        let mut tree_sync = Self {
474            tree,
475            tree_hash: vec![],
476        };
477        // Populate tree hash caches.
478        tree_sync.populate_parent_hashes(provider.crypto(), ciphersuite)?;
479
480        Ok((tree_sync, commit_secret, encryption_key_pair))
481    }
482
483    /// Create a new single-leaf tree for a virtual-client-created group.
484    ///
485    /// The creator's leaf uses the caller-supplied `encryption_key_pair`
486    /// (derived from a `key_package` operation secret) and carries a
487    /// `key_package` `leaf_node_source` with `life_time`, matching the
488    /// `DerivationInfoTBE` selector a sibling uses to reconstruct the leaf. As
489    /// the sole leaf is also the root, it has an empty parent hash.
490    /// `leaf_extensions` already carries the VC derivation info. No commit
491    /// secret is returned: epoch-0 secrets come from the `epoch_secret` derived
492    /// from the KeyPackage seed, not the joiner key schedule.
493    #[cfg(feature = "virtual-clients-draft")]
494    #[allow(clippy::too_many_arguments)]
495    pub(crate) fn new_vc(
496        provider: &impl OpenMlsProvider,
497        signer: &impl Signer,
498        ciphersuite: Ciphersuite,
499        credential_with_key: CredentialWithKey,
500        life_time: Lifetime,
501        capabilities: Capabilities,
502        leaf_extensions: Extensions<LeafNode>,
503        encryption_key_pair: EncryptionKeyPair,
504    ) -> Result<(Self, EncryptionKeyPair), LibraryError> {
505        let new_leaf_node_params = NewLeafNodeParams {
506            ciphersuite,
507            credential_with_key,
508            // A virtual-client-created group's creator leaf is key_package-sourced,
509            // matching the non-VC creator path.
510            leaf_node_source: LeafNodeSource::KeyPackage(life_time),
511            capabilities,
512            extensions: leaf_extensions,
513            tree_info_tbs: TreeInfoTbs::KeyPackage,
514        };
515        let (leaf, encryption_key_pair) = LeafNode::new_with_encryption_key_pair(
516            signer,
517            new_leaf_node_params,
518            encryption_key_pair,
519        )?;
520
521        let node = Node::leaf_node(leaf);
522        let nodes = vec![TreeSyncNode::from(node).into()];
523        let tree = MlsBinaryTree::new(nodes)
524            .map_err(|_| LibraryError::custom("Unexpected error creating the binary tree."))?;
525        let mut tree_sync = Self {
526            tree,
527            tree_hash: vec![],
528        };
529        tree_sync.populate_parent_hashes(provider.crypto(), ciphersuite)?;
530
531        Ok((tree_sync, encryption_key_pair))
532    }
533
534    /// Return the full tree
535    pub(crate) fn tree(&self) -> &MlsBinaryTree<TreeSyncLeafNode, TreeSyncParentNode> {
536        &self.tree
537    }
538
539    /// Return the tree hash of the root node of the tree.
540    pub(crate) fn tree_hash(&self) -> &[u8] {
541        self.tree_hash.as_slice()
542    }
543
544    /// Merge the given diff into this `TreeSync` instance, refreshing the
545    /// `tree_hash` value in the process.
546    pub(crate) fn merge_diff(&mut self, tree_sync_diff: StagedTreeSyncDiff) {
547        let (diff, new_tree_hash) = tree_sync_diff.into_parts();
548        self.tree_hash = new_tree_hash;
549        self.tree.merge_diff(diff);
550    }
551
552    /// Create an empty diff based on this [`TreeSync`] instance all operations
553    /// are created based on an initial, empty [`TreeSyncDiff`].
554    pub(crate) fn empty_diff(&self) -> TreeSyncDiff<'_> {
555        self.into()
556    }
557
558    /// A helper function that generates a [`TreeSync`] instance from the given
559    /// slice of nodes. It verifies that the provided encryption key is present
560    /// in the tree and that the invariants documented in [`TreeSync`] hold.
561    pub(crate) fn from_ratchet_tree(
562        crypto: &impl OpenMlsCrypto,
563        ciphersuite: Ciphersuite,
564        ratchet_tree: RatchetTree,
565    ) -> Result<Self, TreeSyncFromNodesError> {
566        let total_nodes = ratchet_tree.0.len();
567        let mut leaf_nodes = Vec::with_capacity(total_nodes.div_ceil(2));
568        let mut parent_nodes = Vec::with_capacity(total_nodes / 2);
569
570        // Set the leaf indices in all the leaves and convert the node types.
571        for (node_index, node_option) in ratchet_tree.0.into_iter().enumerate() {
572            if node_index % 2 == 0 {
573                let leaf = match node_option {
574                    Some(node) => match TreeSyncNode::from(node) {
575                        TreeSyncNode::Leaf(l) => *l,
576                        TreeSyncNode::Parent(_) => {
577                            return Err(TreeSyncFromNodesError::from(
578                                PublicTreeError::MalformedTree,
579                            ))
580                        }
581                    },
582                    None => TreeSyncLeafNode::blank(),
583                };
584                leaf_nodes.push(leaf);
585            } else {
586                let parent = match node_option {
587                    Some(node) => match TreeSyncNode::from(node) {
588                        TreeSyncNode::Parent(p) => *p,
589                        TreeSyncNode::Leaf(_) => {
590                            return Err(TreeSyncFromNodesError::from(
591                                PublicTreeError::MalformedTree,
592                            ))
593                        }
594                    },
595                    None => TreeSyncParentNode::blank(),
596                };
597                parent_nodes.push(parent);
598            }
599        }
600
601        // Unmerged leaves must point into the tree
602        let leaf_count = leaf_nodes.len() as u32;
603        for parent in parent_nodes.iter() {
604            if let Some(parent_node) = parent.node() {
605                if let Some(last_unmerged_leaf_index) = parent_node.unmerged_leaves().last() {
606                    if last_unmerged_leaf_index.u32() >= leaf_count {
607                        return Err(TreeSyncFromNodesError::from(PublicTreeError::MalformedTree));
608                    }
609                }
610            }
611        }
612
613        let tree = MlsBinaryTree::from_components(leaf_nodes, parent_nodes)
614            .map_err(|_| PublicTreeError::MalformedTree)?;
615        let mut tree_sync = Self {
616            tree,
617            tree_hash: vec![],
618        };
619
620        // Verify all parent hashes.
621        tree_sync
622            .verify_parent_hashes(crypto, ciphersuite)
623            .map_err(|e| match e {
624                TreeSyncParentHashError::LibraryError(e) => e.into(),
625                TreeSyncParentHashError::InvalidParentHash => {
626                    TreeSyncFromNodesError::from(PublicTreeError::InvalidParentHash)
627                }
628            })?;
629
630        // Populate tree hash caches.
631        tree_sync.populate_parent_hashes(crypto, ciphersuite)?;
632        Ok(tree_sync)
633    }
634
635    /// Find the `LeafNodeIndex` which a new leaf would have if it were added to the
636    /// tree. This is either the left-most blank node or, if there are no blank
637    /// leaves, the leaf count, since adding a member would extend the tree by
638    /// one leaf.
639    pub(crate) fn free_leaf_index(&self) -> LeafNodeIndex {
640        let diff = self.empty_diff();
641        diff.free_leaf_index()
642    }
643
644    /// Populate the parent hash caches of all nodes in the tree.
645    fn populate_parent_hashes(
646        &mut self,
647        crypto: &impl OpenMlsCrypto,
648        ciphersuite: Ciphersuite,
649    ) -> Result<(), LibraryError> {
650        let diff = self.empty_diff();
651        // Make the diff into a staged diff. This implicitly computes the
652        // tree hashes and poulates the tree hash caches.
653        let staged_diff = diff.into_staged_diff(crypto, ciphersuite)?;
654        // Merge the diff.
655        self.merge_diff(staged_diff);
656        Ok(())
657    }
658
659    /// Verify the parent hashes of all parent nodes in the tree.
660    ///
661    /// Returns an error if one of the parent nodes in the tree has an invalid
662    /// parent hash.
663    fn verify_parent_hashes(
664        &self,
665        crypto: &impl OpenMlsCrypto,
666        ciphersuite: Ciphersuite,
667    ) -> Result<(), TreeSyncParentHashError> {
668        // The ability to verify parent hashes is required both for diffs and
669        // treesync instances. We choose the computationally slightly more
670        // expensive solution of implementing parent hash verification for the
671        // diff and creating an empty diff whenever we need to verify parent
672        // hashes for a `TreeSync` instance. At the time of writing, this
673        // happens only upon construction of a `TreeSync` instance from a vector
674        // of nodes. The alternative solution would be to create a `TreeLike`
675        // trait, which allows tree navigation and node access. We could then
676        // implement `TreeLike` for both `TreeSync` and `TreeSyncDiff` and
677        // finally implement parent hash verification for any struct that
678        // implements `TreeLike`. We choose the less complex version for now.
679        // Should this turn out to cause too much computational overhead, we
680        // should reconsider and choose the alternative sketched above
681        let diff = self.empty_diff();
682        // No need to merge the diff, since we didn't actually modify any state.
683        diff.verify_parent_hashes(crypto, ciphersuite)
684    }
685
686    /// Returns the tree size
687    pub(crate) fn tree_size(&self) -> TreeSize {
688        self.tree.tree_size()
689    }
690
691    /// Returns a vec of all leaf slots, including blanks.
692    pub fn leaves(&self) -> Vec<Option<&LeafNode>> {
693        self.tree
694            .leaves()
695            .map(|(_, tsn)| tsn.node().as_ref())
696            .collect()
697    }
698
699    /// Returns an iterator over the (non-blank) [`LeafNode`]s in the tree.
700    pub fn full_leaves(&self) -> impl Iterator<Item = (LeafNodeIndex, &LeafNode)> {
701        self.tree
702            .leaves()
703            .filter_map(|(index, tsn)| tsn.node().as_ref().map(|ln| (index, ln)))
704    }
705
706    /// Returns an iterator over the (non-blank) [`ParentNode`]s in the tree.
707    pub fn full_parents(&self) -> impl Iterator<Item = (ParentNodeIndex, &ParentNode)> {
708        self.tree
709            .parents()
710            .filter_map(|(index, tsn)| tsn.node().as_ref().map(|pn| (index, pn)))
711    }
712
713    /// Returns an iterator over the [`ParentNodeIndex`]es of blank [`ParentNode`]s in the tree.
714    pub fn blank_parents<'a>(&'a self) -> impl Iterator<Item = ParentNodeIndex> + 'a {
715        self.tree
716            .parents()
717            .filter_map(|(index, tsn)| tsn.node().as_ref().map_or(Some(index), |_| None))
718    }
719
720    /// Returns an iterator over the [`LeafNodeIndex`]es of blank [`LeafNode`]s in the tree.
721    pub fn blank_leaves<'a>(&'a self) -> impl Iterator<Item = LeafNodeIndex> + 'a {
722        self.tree
723            .leaves()
724            .filter_map(|(index, tsn)| tsn.node().as_ref().map_or(Some(index), |_| None))
725    }
726
727    /// Returns the index of the last full leaf in the tree.
728    fn rightmost_full_leaf(&self) -> LeafNodeIndex {
729        let mut index = LeafNodeIndex::new(0);
730        for (leaf_index, leaf) in self.tree.leaves() {
731            if leaf.node().as_ref().is_some() {
732                index = leaf_index;
733            }
734        }
735        index
736    }
737
738    /// Returns a list of [`Member`]s containing only full nodes.
739    ///
740    /// XXX: For performance reasons we probably want to have this in a borrowing
741    ///      version as well. But it might well go away again.
742    pub(crate) fn full_leaf_members(&self) -> impl Iterator<Item = Member> + '_ {
743        self.tree
744            .leaves()
745            // Filter out blank nodes
746            .filter_map(|(index, tsn)| tsn.node().as_ref().map(|node| (index, node)))
747            // Map to `Member`
748            .map(|(index, leaf_node)| {
749                Member::new(
750                    index,
751                    leaf_node.encryption_key().as_slice().to_vec(),
752                    leaf_node.signature_key().as_slice().to_vec(),
753                    leaf_node.credential().clone(),
754                )
755            })
756    }
757
758    /// Returns the nodes in the tree ordered according to the
759    /// array-representation of the underlying binary tree.
760    pub fn export_ratchet_tree(&self) -> RatchetTree {
761        let mut nodes = Vec::new();
762
763        // Determine the index of the rightmost full leaf.
764        let max_length = self.rightmost_full_leaf();
765
766        // We take all the leaves including the rightmost full leaf, blank
767        // leaves beyond that are trimmed.
768        let mut leaves = self
769            .tree
770            .leaves()
771            .map(|(_, leaf)| leaf)
772            .take(max_length.usize() + 1);
773
774        // Get the first leaf.
775        if let Some(leaf) = leaves.next() {
776            nodes.push(leaf.node().clone().map(Node::leaf_node));
777        } else {
778            // The tree was empty.
779            return RatchetTree::trimmed(vec![]);
780        }
781
782        // Blank parent node used for padding
783        let default_parent = TreeSyncParentNode::default();
784
785        // Get the parents.
786        let parents = self
787            .tree
788            .parents()
789            // Drop the index
790            .map(|(_, parent)| parent)
791            // Take the parents up to the max length
792            .take(max_length.usize())
793            // Pad the parents with blank nodes if needed
794            .chain(
795                (self.tree.parents().count()..self.tree.leaves().count() - 1)
796                    .map(|_| &default_parent),
797            );
798
799        // Interleave the leaves and parents.
800        for (leaf, parent) in leaves.zip(parents) {
801            nodes.push(parent.node().clone().map(Node::parent_node));
802            nodes.push(leaf.node().clone().map(Node::leaf_node));
803        }
804
805        RatchetTree::trimmed(nodes)
806    }
807
808    /// Return a reference to the leaf at the given `LeafNodeIndex` or `None` if the
809    /// leaf is blank.
810    pub(crate) fn leaf(&self, leaf_index: LeafNodeIndex) -> Option<&LeafNode> {
811        let tsn = self.tree.leaf(leaf_index);
812        tsn.node().as_ref()
813    }
814
815    /// Returns a [`TreeSyncError`] if the `leaf_index` is not a leaf in this
816    /// tree or empty.
817    pub(crate) fn is_leaf_in_tree(&self, leaf_index: LeafNodeIndex) -> bool {
818        is_node_in_tree(leaf_index.into(), self.tree.tree_size())
819    }
820
821    /// Return a vector containing all [`EncryptionKey`]s for which the owner of
822    /// the given `leaf_index` should have private key material.
823    pub(crate) fn owned_encryption_keys(&self, leaf_index: LeafNodeIndex) -> Vec<EncryptionKey> {
824        self.empty_diff()
825            .encryption_keys(leaf_index)
826            .cloned()
827            .collect::<Vec<EncryptionKey>>()
828    }
829
830    /// Derives [`EncryptionKeyPair`]s for the nodes in the shared direct path
831    /// of the leaves with index `leaf_index` and `sender_index`.  This function
832    /// also checks that the derived public keys match the existing public keys.
833    ///
834    /// Returns the `CommitSecret` derived from the path secret of the root
835    /// node, as well as the derived [`EncryptionKeyPair`]s. Returns an error if
836    /// the target leaf is outside of the tree.
837    ///
838    /// Returns TreeSyncSetPathError::PublicKeyMismatch if the derived keys don't
839    /// match with the existing ones.
840    ///
841    /// Returns TreeSyncSetPathError::LibraryError if the sender_index is not
842    /// in the tree.
843    pub(crate) fn derive_path_secrets(
844        &self,
845        crypto: &impl OpenMlsCrypto,
846        ciphersuite: Ciphersuite,
847        mut path_secret: PathSecret,
848        sender_index: LeafNodeIndex,
849        leaf_index: LeafNodeIndex,
850    ) -> Result<(Vec<EncryptionKeyPair>, CommitSecret), DerivePathError> {
851        // We assume both nodes are in the tree, since the sender_index must be in the tree
852        // Skip the nodes in the subtree path for which we are an unmerged leaf.
853        let subtree_path = self.tree.subtree_path(leaf_index, sender_index);
854        let mut keypairs = Vec::new();
855        for parent_index in subtree_path {
856            // We know the node is in the tree, since it is in the subtree path
857            let tsn = self.tree.parent_by_index(parent_index);
858            // We only care about non-blank nodes.
859            if let Some(ref parent_node) = tsn.node() {
860                // If our own leaf index is not in the list of unmerged leaves
861                // then we should have the secret for this node.
862                if !parent_node.unmerged_leaves().contains(&leaf_index) {
863                    let keypair = path_secret.derive_key_pair(crypto, ciphersuite)?;
864                    // The derived public key should match the one in the node.
865                    // If not, the tree is corrupt.
866                    if parent_node.encryption_key() != keypair.public_key() {
867                        return Err(DerivePathError::PublicKeyMismatch);
868                    } else {
869                        // If everything is ok, set the private key and derive
870                        // the next path secret.
871                        keypairs.push(keypair);
872                        path_secret = path_secret.derive_path_secret(crypto, ciphersuite)?;
873                    }
874                };
875                // If the leaf is blank or our index is in the list of unmerged
876                // leaves, go to the next node.
877            }
878        }
879        Ok((keypairs, path_secret.into()))
880    }
881
882    /// Return a reference to the parent node at the given `ParentNodeIndex` or
883    /// `None` if the node is blank.
884    pub(crate) fn parent(&self, node_index: ParentNodeIndex) -> Option<&ParentNode> {
885        let tsn = self.tree.parent(node_index);
886        tsn.node().as_ref()
887    }
888}
889
890#[cfg(test)]
891impl TreeSync {
892    pub(crate) fn leaf_count(&self) -> u32 {
893        self.tree.leaf_count()
894    }
895}
896
897#[cfg(test)]
898mod test {
899    use super::*;
900
901    #[cfg(debug_assertions)]
902    #[test]
903    #[should_panic]
904    /// This should only panic in debug-builds.
905    fn test_ratchet_tree_internal_empty() {
906        RatchetTree::trimmed(vec![]);
907    }
908
909    #[cfg(debug_assertions)]
910    #[test]
911    #[should_panic]
912    /// This should only panic in debug-builds.
913    fn test_ratchet_tree_internal_empty_after_trim() {
914        RatchetTree::trimmed(vec![None]);
915    }
916
917    #[openmls_test::openmls_test]
918    fn test_ratchet_tree_trailing_blank_nodes() {
919        let provider = &Provider::default();
920        let (key_package, _, _) = crate::key_packages::tests::key_package(ciphersuite, provider);
921        let node_in = NodeIn::from(Node::leaf_node(LeafNode::from(key_package)));
922        let tests = [
923            (vec![], false),
924            (vec![None], false),
925            (vec![None, None], false),
926            (vec![None, None, None], false),
927            (vec![Some(node_in.clone())], true),
928            (vec![Some(node_in.clone()), None], false),
929            (
930                vec![Some(node_in.clone()), None, Some(node_in.clone())],
931                true,
932            ),
933            (
934                vec![Some(node_in.clone()), None, Some(node_in), None],
935                false,
936            ),
937        ];
938
939        for (test, expected) in tests.into_iter() {
940            let got = RatchetTree::try_from_nodes(
941                ciphersuite,
942                provider.crypto(),
943                test,
944                &GroupId::random(provider.rand()),
945            )
946            .is_ok();
947            assert_eq!(got, expected);
948        }
949    }
950
951    #[cfg(not(debug_assertions))]
952    #[test]
953    /// This should not panic in release-builds.
954    fn test_ratchet_tree_internal_empty() {
955        RatchetTree::trimmed(vec![]);
956    }
957
958    #[cfg(not(debug_assertions))]
959    #[test]
960    /// This should not panic in release-builds.
961    fn test_ratchet_tree_internal_empty_after_trim() {
962        RatchetTree::trimmed(vec![None]);
963    }
964
965    #[openmls_test::openmls_test]
966    fn test_ratchet_tree_in_full_leaves_reports_real_index_past_a_blank_leaf() {
967        let provider = &Provider::default();
968        let (key_package, credential, _) =
969            crate::key_packages::tests::key_package(ciphersuite, provider);
970        let node_in = NodeIn::from(Node::leaf_node(LeafNode::from(key_package)));
971
972        // A 2-leaf tree whose first leaf (index 0) is blank and whose second
973        // leaf (index 1) is the only non-blank one: flat positions
974        // 0 = leaf 0 (blank), 1 = parent (blank), 2 = leaf 1 (node_in).
975        let ratchet_tree = RatchetTreeIn(vec![None, None, Some(node_in)]);
976
977        // `leaves()` skips the blank slot, so its position is compacted and no
978        // longer matches the leaf's real index -- exactly the bug `full_leaves()`
979        // fixes for callers that need the real `LeafNodeIndex`.
980        let compacted: Vec<_> = ratchet_tree.leaves().collect();
981        assert_eq!(compacted.len(), 1);
982
983        let full: Vec<_> = ratchet_tree.full_leaves().collect();
984        assert_eq!(full.len(), 1);
985        let (index, leaf) = full[0];
986        assert_eq!(index, LeafNodeIndex::new(1));
987        assert_eq!(leaf.credential(), &credential);
988    }
989}