1#[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
65mod hashes;
67use errors::*;
68
69pub(crate) mod diff;
71pub(crate) mod node;
72pub(crate) mod treekem;
73pub(crate) mod treesync_node;
74
75use node::encryption_keys::EncryptionKeyPair;
76
77pub mod errors;
79#[cfg(feature = "test-utils")]
80pub use node::encryption_keys::test_utils;
81pub use node::encryption_keys::EncryptionKey;
82
83pub use node::{
85 leaf_node::{
86 CapabilitiesPolicy, LeafNode, LeafNodeBuildError, LeafNodeParameters,
87 LeafNodeParametersBuilder, LeafNodeSource, LeafNodeUpdateError,
88 },
89 parent_node::ParentNode,
90 Node,
91};
92
93#[cfg(any(feature = "test-utils", test))]
95pub mod tests_and_kats;
96
97#[derive(PartialEq, Eq, Clone, Debug, Serialize, Deserialize, TlsSerialize, TlsSize)]
99pub struct RatchetTree(Vec<Option<Node>>);
100
101#[derive(Error, Debug, PartialEq, Clone)]
103pub enum RatchetTreeError {
104 #[error("The ratchet tree has no nodes.")]
106 MissingNodes,
107 #[error("The ratchet tree has trailing blank nodes.")]
109 TrailingBlankNodes,
110 #[error("Invalid node signature.")]
112 InvalidNodeSignature,
113 #[error("Wrong node type.")]
115 WrongNodeType,
116}
117
118impl RatchetTree {
119 fn trimmed(mut nodes: Vec<Option<Node>>) -> Self {
123 match nodes.iter().enumerate().rfind(|(_, node)| node.is_some()) {
125 Some((rightmost_nonempty_position, _)) => {
126 nodes.resize(rightmost_nonempty_position + 1, None);
128 }
129 None => {
130 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 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 match nodes.last() {
150 Some(None) => {
151 Err(RatchetTreeError::TrailingBlankNodes)
153 }
154 None => {
155 Err(RatchetTreeError::MissingNodes)
157 }
158 Some(Some(_)) => {
159 let mut verified_nodes = Vec::new();
164 for (index, node) in nodes.into_iter().enumerate() {
165 let verified_node = match (index % 2, node) {
166 (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 (1, Some(NodeIn::ParentNode(parent_node))) => {
199 Some(Node::ParentNode(parent_node))
200 }
201 (_, None) => None,
203 _ => {
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 pub fn nodes(&self) -> impl Iterator<Item = &Node> {
217 self.0.iter().flatten()
218 }
219
220 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 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#[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 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 pub fn nodes(&self) -> impl Iterator<Item = &NodeIn> {
267 self.0.iter().flatten()
268 }
269
270 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 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 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#[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#[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 #[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 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 tree_sync.populate_parent_hashes(provider.crypto(), ciphersuite)?;
485
486 Ok((tree_sync, commit_secret, encryption_key_pair))
487 }
488
489 #[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 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 pub(crate) fn tree(&self) -> &MlsBinaryTree<TreeSyncLeafNode, TreeSyncParentNode> {
546 &self.tree
547 }
548
549 pub(crate) fn tree_hash(&self) -> &[u8] {
551 self.tree_hash.as_slice()
552 }
553
554 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 pub(crate) fn empty_diff(&self) -> TreeSyncDiff<'_> {
565 self.into()
566 }
567
568 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 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 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 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 tree_sync.populate_parent_hashes(crypto, ciphersuite)?;
642 Ok(tree_sync)
643 }
644
645 pub(crate) fn free_leaf_index(&self) -> LeafNodeIndex {
650 let diff = self.empty_diff();
651 diff.free_leaf_index()
652 }
653
654 fn populate_parent_hashes(
656 &mut self,
657 crypto: &impl OpenMlsCrypto,
658 ciphersuite: Ciphersuite,
659 ) -> Result<(), LibraryError> {
660 let diff = self.empty_diff();
661 let staged_diff = diff.into_staged_diff(crypto, ciphersuite)?;
664 self.merge_diff(staged_diff);
666 Ok(())
667 }
668
669 fn verify_parent_hashes(
674 &self,
675 crypto: &impl OpenMlsCrypto,
676 ciphersuite: Ciphersuite,
677 ) -> Result<(), TreeSyncParentHashError> {
678 let diff = self.empty_diff();
692 diff.verify_parent_hashes(crypto, ciphersuite)
694 }
695
696 pub(crate) fn tree_size(&self) -> TreeSize {
698 self.tree.tree_size()
699 }
700
701 pub fn leaves(&self) -> Vec<Option<&LeafNode>> {
703 self.tree
704 .leaves()
705 .map(|(_, tsn)| tsn.node().as_ref())
706 .collect()
707 }
708
709 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 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 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 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 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 pub(crate) fn full_leaf_members(&self) -> impl Iterator<Item = Member> + '_ {
753 self.tree
754 .leaves()
755 .filter_map(|(index, tsn)| tsn.node().as_ref().map(|node| (index, node)))
757 .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 pub fn export_ratchet_tree(&self) -> RatchetTree {
771 let mut nodes = Vec::new();
772
773 let max_length = self.rightmost_full_leaf();
775
776 let mut leaves = self
779 .tree
780 .leaves()
781 .map(|(_, leaf)| leaf)
782 .take(max_length.usize() + 1);
783
784 if let Some(leaf) = leaves.next() {
786 nodes.push(leaf.node().clone().map(Node::leaf_node));
787 } else {
788 return RatchetTree::trimmed(vec![]);
790 }
791
792 let default_parent = TreeSyncParentNode::default();
794
795 let parents = self
797 .tree
798 .parents()
799 .map(|(_, parent)| parent)
801 .take(max_length.usize())
803 .chain(
805 (self.tree.parents().count()..self.tree.leaves().count() - 1)
806 .map(|_| &default_parent),
807 );
808
809 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 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 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 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 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 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 let tsn = self.tree.parent_by_index(parent_index);
868 if let Some(ref parent_node) = tsn.node() {
870 if !parent_node.unmerged_leaves().contains(&leaf_index) {
873 let keypair = path_secret.derive_key_pair(crypto, ciphersuite)?;
874 if parent_node.encryption_key() != keypair.public_key() {
877 return Err(DerivePathError::PublicKeyMismatch);
878 } else {
879 keypairs.push(keypair);
882 path_secret = path_secret.derive_path_secret(crypto, ciphersuite)?;
883 }
884 };
885 }
888 }
889 Ok((keypairs, path_secret.into()))
890 }
891
892 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 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 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 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 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 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 fn test_ratchet_tree_internal_empty() {
994 RatchetTree::trimmed(vec![]);
995 }
996
997 #[cfg(debug_assertions)]
998 #[test]
999 #[should_panic]
1000 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 fn test_ratchet_tree_internal_empty() {
1043 RatchetTree::trimmed(vec![]);
1044 }
1045
1046 #[cfg(not(debug_assertions))]
1047 #[test]
1048 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 let ratchet_tree = RatchetTreeIn(vec![None, None, Some(node_in)]);
1064
1065 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}