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