mls_rs/tree_kem/
node.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// Copyright by contributors to this project.
3// SPDX-License-Identifier: (Apache-2.0 OR MIT)
4
5use super::leaf_node::LeafNode;
6use crate::client::MlsError;
7use crate::crypto::HpkePublicKey;
8use crate::tree_kem::math as tree_math;
9use crate::tree_kem::parent_hash::ParentHash;
10use alloc::vec;
11use alloc::vec::Vec;
12use core::hash::Hash;
13use core::ops::{Deref, DerefMut};
14use mls_rs_codec::{MlsDecode, MlsEncode, MlsSize};
15use tree_math::{CopathNode, TreeIndex};
16
17#[cfg(feature = "serde")]
18use mls_rs_core::error::IntoAnyError;
19
20// Restrict leaf index to 24bits
21pub(crate) const MAX_LEAF_INDEX: u32 = (1 << 24) - 1;
22
23#[derive(Clone, Debug, PartialEq, MlsSize, MlsEncode, MlsDecode)]
24#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
25pub struct Parent {
26    pub public_key: HpkePublicKey,
27    pub parent_hash: ParentHash,
28    pub unmerged_leaves: Vec<LeafIndex>,
29}
30
31#[derive(Clone, Copy, Debug, Ord, PartialEq, PartialOrd, Hash, Eq, MlsSize, MlsEncode)]
32#[cfg_attr(feature = "serde", derive(serde::Serialize))]
33pub struct LeafIndex(u32);
34
35#[cfg(feature = "arbitrary")]
36impl<'a> arbitrary::Arbitrary<'a> for LeafIndex {
37    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
38        let value = u.int_in_range(0..=MAX_LEAF_INDEX)?;
39        Ok(LeafIndex(value))
40    }
41}
42
43impl TryFrom<u32> for LeafIndex {
44    type Error = MlsError;
45
46    fn try_from(value: u32) -> Result<Self, Self::Error> {
47        if value > MAX_LEAF_INDEX {
48            return Err(MlsError::InvalidTreeIndex);
49        }
50
51        Ok(Self(value))
52    }
53}
54
55impl LeafIndex {
56    pub(crate) fn from_node_index_unchecked(index: NodeIndex) -> Self {
57        LeafIndex(index >> 1)
58    }
59
60    pub(crate) fn unchecked(value: u32) -> Self {
61        Self(value)
62    }
63
64    pub(crate) fn next_unchecked(&self) -> Self {
65        LeafIndex(self.0 + 1)
66    }
67}
68
69impl Deref for LeafIndex {
70    type Target = u32;
71
72    fn deref(&self) -> &Self::Target {
73        &self.0
74    }
75}
76
77impl From<&LeafIndex> for NodeIndex {
78    fn from(leaf_index: &LeafIndex) -> Self {
79        leaf_index.0 * 2
80    }
81}
82
83impl From<LeafIndex> for NodeIndex {
84    fn from(leaf_index: LeafIndex) -> Self {
85        leaf_index.0 * 2
86    }
87}
88
89impl MlsDecode for LeafIndex {
90    fn mls_decode(reader: &mut &[u8]) -> Result<Self, mls_rs_codec::Error> {
91        let val = u32::mls_decode(reader)?;
92        LeafIndex::try_from(val).map_err(|_| mls_rs_codec::Error::Custom(6))
93    }
94}
95
96#[cfg(feature = "serde")]
97impl<'de> serde::Deserialize<'de> for LeafIndex {
98    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
99    where
100        D: serde::Deserializer<'de>,
101    {
102        let val = u32::deserialize(deserializer)?;
103
104        LeafIndex::try_from(val).map_err(|e| serde::de::Error::custom(e.into_any_error()))
105    }
106}
107
108pub type NodeIndex = u32;
109
110#[derive(Clone, Debug, PartialEq, MlsSize, MlsEncode, MlsDecode)]
111#[allow(clippy::large_enum_variant)]
112#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
113#[repr(u8)]
114//TODO: Research if this should actually be a Box<Leaf> for memory / performance reasons
115pub enum Node {
116    Leaf(LeafNode) = 1u8,
117    Parent(Parent) = 2u8,
118}
119
120impl Node {
121    pub fn public_key(&self) -> &HpkePublicKey {
122        match self {
123            Node::Parent(p) => &p.public_key,
124            Node::Leaf(l) => &l.public_key,
125        }
126    }
127}
128
129impl From<Parent> for Option<Node> {
130    fn from(p: Parent) -> Self {
131        Node::from(p).into()
132    }
133}
134
135impl From<LeafNode> for Option<Node> {
136    fn from(l: LeafNode) -> Self {
137        Node::from(l).into()
138    }
139}
140
141impl From<Parent> for Node {
142    fn from(p: Parent) -> Self {
143        Node::Parent(p)
144    }
145}
146
147impl From<LeafNode> for Node {
148    fn from(l: LeafNode) -> Self {
149        Node::Leaf(l)
150    }
151}
152
153pub(crate) trait NodeTypeResolver {
154    fn as_parent(&self) -> Result<&Parent, MlsError>;
155    fn as_parent_mut(&mut self) -> Result<&mut Parent, MlsError>;
156    fn as_leaf(&self) -> Result<&LeafNode, MlsError>;
157    fn as_leaf_mut(&mut self) -> Result<&mut LeafNode, MlsError>;
158    fn as_non_empty(&self) -> Result<&Node, MlsError>;
159}
160
161impl NodeTypeResolver for Option<Node> {
162    fn as_parent(&self) -> Result<&Parent, MlsError> {
163        self.as_ref()
164            .and_then(|n| match n {
165                Node::Parent(p) => Some(p),
166                Node::Leaf(_) => None,
167            })
168            .ok_or(MlsError::ExpectedNode)
169    }
170
171    fn as_parent_mut(&mut self) -> Result<&mut Parent, MlsError> {
172        self.as_mut()
173            .and_then(|n| match n {
174                Node::Parent(p) => Some(p),
175                Node::Leaf(_) => None,
176            })
177            .ok_or(MlsError::ExpectedNode)
178    }
179
180    fn as_leaf(&self) -> Result<&LeafNode, MlsError> {
181        self.as_ref()
182            .and_then(|n| match n {
183                Node::Parent(_) => None,
184                Node::Leaf(l) => Some(l),
185            })
186            .ok_or(MlsError::ExpectedNode)
187    }
188
189    fn as_leaf_mut(&mut self) -> Result<&mut LeafNode, MlsError> {
190        self.as_mut()
191            .and_then(|n| match n {
192                Node::Parent(_) => None,
193                Node::Leaf(l) => Some(l),
194            })
195            .ok_or(MlsError::ExpectedNode)
196    }
197
198    fn as_non_empty(&self) -> Result<&Node, MlsError> {
199        self.as_ref().ok_or(MlsError::UnexpectedEmptyNode)
200    }
201}
202
203#[derive(Clone, Debug, PartialEq, MlsSize, MlsEncode, MlsDecode, Default)]
204#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
205pub struct NodeVec(Vec<Option<Node>>);
206
207impl From<Vec<Option<Node>>> for NodeVec {
208    fn from(x: Vec<Option<Node>>) -> Self {
209        NodeVec(x)
210    }
211}
212
213impl Deref for NodeVec {
214    type Target = Vec<Option<Node>>;
215
216    fn deref(&self) -> &Self::Target {
217        &self.0
218    }
219}
220
221impl DerefMut for NodeVec {
222    fn deref_mut(&mut self) -> &mut Self::Target {
223        &mut self.0
224    }
225}
226
227impl NodeVec {
228    #[cfg(any(test, all(feature = "custom_proposal", feature = "tree_index")))]
229    pub fn occupied_leaf_count(&self) -> u32 {
230        self.non_empty_leaves().count() as u32
231    }
232
233    pub fn total_leaf_count(&self) -> u32 {
234        (self.len() as u32 / 2 + 1).next_power_of_two()
235    }
236
237    #[inline]
238    pub fn borrow_node(&self, index: NodeIndex) -> Result<&Option<Node>, MlsError> {
239        Ok(self.get(self.validate_index(index)?).unwrap_or(&None))
240    }
241
242    fn validate_index(&self, index: NodeIndex) -> Result<usize, MlsError> {
243        if (index as usize) >= self.len().next_power_of_two() {
244            Err(MlsError::InvalidNodeIndex(index))
245        } else {
246            Ok(index as usize)
247        }
248    }
249
250    #[cfg(test)]
251    fn empty_leaves(&mut self) -> impl Iterator<Item = (LeafIndex, &mut Option<Node>)> {
252        self.iter_mut()
253            .step_by(2)
254            .enumerate()
255            .filter(|(_, n)| n.is_none())
256            .map(|(i, n)| (LeafIndex::unchecked(i as u32), n))
257    }
258
259    pub fn non_empty_leaves(&self) -> impl Iterator<Item = (LeafIndex, &LeafNode)> + '_ {
260        self.leaves()
261            .enumerate()
262            .filter_map(|(i, l)| l.map(|l| (LeafIndex::unchecked(i as u32), l)))
263    }
264
265    pub fn non_empty_parents(&self) -> impl Iterator<Item = (NodeIndex, &Parent)> + '_ {
266        self.iter()
267            .enumerate()
268            .skip(1)
269            .step_by(2)
270            .map(|(i, n)| (i as NodeIndex, n))
271            .filter_map(|(i, n)| n.as_parent().ok().map(|p| (i, p)))
272    }
273
274    pub fn leaves(&self) -> impl Iterator<Item = Option<&LeafNode>> + '_ {
275        self.iter().step_by(2).map(|n| n.as_leaf().ok())
276    }
277
278    pub fn direct_copath(&self, index: LeafIndex) -> Vec<CopathNode<NodeIndex>> {
279        NodeIndex::from(index).direct_copath(&self.total_leaf_count())
280    }
281
282    // Section 8.4
283    // The filtered direct path of a node is obtained from the node's direct path by removing
284    // all nodes whose child on the nodes's copath has an empty resolution
285    pub fn filtered(&self, index: LeafIndex) -> Result<Vec<bool>, MlsError> {
286        Ok(NodeIndex::from(index)
287            .direct_copath(&self.total_leaf_count())
288            .into_iter()
289            .map(|cp| self.is_resolution_empty(cp.copath))
290            .collect())
291    }
292
293    #[inline]
294    pub fn is_blank(&self, index: NodeIndex) -> Result<bool, MlsError> {
295        self.borrow_node(index).map(|n| n.is_none())
296    }
297
298    #[inline]
299    pub fn is_leaf(&self, index: NodeIndex) -> bool {
300        index % 2 == 0
301    }
302
303    // Blank a previously filled leaf node, and return the existing leaf
304    pub fn blank_leaf_node(&mut self, leaf_index: LeafIndex) -> Result<LeafNode, MlsError> {
305        let node_index = self.validate_index(leaf_index.into())?;
306
307        match self.get_mut(node_index).and_then(Option::take) {
308            Some(Node::Leaf(l)) => Ok(l),
309            _ => Err(MlsError::RemovingNonExistingMember),
310        }
311    }
312
313    pub fn blank_direct_path(&mut self, leaf: LeafIndex) -> Result<(), MlsError> {
314        for i in self.direct_copath(leaf) {
315            if let Some(n) = self.get_mut(i.path as usize) {
316                *n = None
317            }
318        }
319
320        Ok(())
321    }
322
323    // Remove elements until the last node is non-blank
324    pub fn trim(&mut self) {
325        while self.last() == Some(&None) {
326            self.pop();
327        }
328    }
329
330    pub fn borrow_as_parent(&self, node_index: NodeIndex) -> Result<&Parent, MlsError> {
331        self.borrow_node(node_index).and_then(|n| n.as_parent())
332    }
333
334    pub fn borrow_as_parent_mut(&mut self, node_index: NodeIndex) -> Result<&mut Parent, MlsError> {
335        let index = self.validate_index(node_index)?;
336
337        self.get_mut(index)
338            .ok_or(MlsError::InvalidNodeIndex(node_index))?
339            .as_parent_mut()
340    }
341
342    pub fn borrow_as_leaf_mut(&mut self, index: LeafIndex) -> Result<&mut LeafNode, MlsError> {
343        let node_index = NodeIndex::from(index);
344        let index = self.validate_index(node_index)?;
345
346        self.get_mut(index)
347            .ok_or(MlsError::InvalidNodeIndex(node_index))?
348            .as_leaf_mut()
349    }
350
351    pub fn borrow_as_leaf(&self, index: LeafIndex) -> Result<&LeafNode, MlsError> {
352        let node_index = NodeIndex::from(index);
353        self.borrow_node(node_index).and_then(|n| n.as_leaf())
354    }
355
356    pub fn borrow_or_fill_node_as_parent(
357        &mut self,
358        node_index: NodeIndex,
359        public_key: &HpkePublicKey,
360    ) -> Result<&mut Parent, MlsError> {
361        let index = self.validate_index(node_index)?;
362
363        while self.len() <= index {
364            self.push(None);
365        }
366
367        self.get_mut(index)
368            .ok_or(MlsError::InvalidNodeIndex(node_index))
369            .and_then(|n| {
370                if n.is_none() {
371                    *n = Parent {
372                        public_key: public_key.clone(),
373                        parent_hash: ParentHash::empty(),
374                        unmerged_leaves: vec![],
375                    }
376                    .into();
377                }
378                n.as_parent_mut()
379            })
380    }
381
382    pub fn get_resolution_index(&self, index: NodeIndex) -> Result<Vec<NodeIndex>, MlsError> {
383        let mut indexes = vec![index];
384        let mut resolution = vec![];
385
386        while let Some(index) = indexes.pop() {
387            if let Some(Some(node)) = self.get(index as usize) {
388                resolution.push(index);
389
390                if let Node::Parent(p) = node {
391                    resolution.extend(p.unmerged_leaves.iter().map(NodeIndex::from));
392                }
393            } else if !index.is_leaf() {
394                indexes.push(index.right_unchecked());
395                indexes.push(index.left_unchecked());
396            }
397        }
398
399        Ok(resolution)
400    }
401
402    pub fn find_in_resolution(
403        &self,
404        index: NodeIndex,
405        to_find: Option<NodeIndex>,
406    ) -> Option<usize> {
407        let mut indexes = vec![index];
408        let mut resolution_len = 0;
409
410        while let Some(index) = indexes.pop() {
411            if let Some(Some(node)) = self.get(index as usize) {
412                if Some(index) == to_find || to_find.is_none() {
413                    return Some(resolution_len);
414                }
415
416                resolution_len += 1;
417
418                if let Node::Parent(p) = node {
419                    indexes.extend(p.unmerged_leaves.iter().map(NodeIndex::from));
420                }
421            } else if !index.is_leaf() {
422                indexes.push(index.right_unchecked());
423                indexes.push(index.left_unchecked());
424            }
425        }
426
427        None
428    }
429
430    pub fn is_resolution_empty(&self, index: NodeIndex) -> bool {
431        self.find_in_resolution(index, None).is_none()
432    }
433
434    pub(crate) fn next_empty_leaf(&self, start: LeafIndex) -> LeafIndex {
435        let mut n = NodeIndex::from(start) as usize;
436
437        while n < self.len() {
438            if self.0[n].is_none() {
439                return LeafIndex::from_node_index_unchecked(n as NodeIndex);
440            }
441
442            n += 2;
443        }
444
445        LeafIndex::from_node_index_unchecked(self.len() as NodeIndex + 1)
446    }
447
448    /// If `index` fits in the current tree, inserts `leaf` at `index`. Else, inserts `leaf` as the
449    /// last leaf
450    pub fn insert_leaf(&mut self, index: LeafIndex, leaf: LeafNode) {
451        let node_index = (*index as usize) << 1;
452
453        if node_index > self.len() {
454            self.push(None);
455            self.push(None);
456        } else if self.is_empty() {
457            self.push(None);
458        }
459
460        self.0[node_index] = Some(leaf.into());
461    }
462}
463
464#[cfg(test)]
465pub(crate) mod test_utils {
466    use super::*;
467    use crate::{
468        client::test_utils::TEST_CIPHER_SUITE, tree_kem::leaf_node::test_utils::get_basic_test_node,
469    };
470
471    #[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
472    pub(crate) async fn get_test_node_vec() -> NodeVec {
473        let mut nodes = vec![None; 7];
474
475        nodes[0] = get_basic_test_node(TEST_CIPHER_SUITE, "A").await.into();
476        nodes[4] = get_basic_test_node(TEST_CIPHER_SUITE, "C").await.into();
477
478        nodes[5] = Parent {
479            public_key: b"CD".to_vec().into(),
480            parent_hash: ParentHash::empty(),
481            unmerged_leaves: vec![LeafIndex::unchecked(2)],
482        }
483        .into();
484
485        nodes[6] = get_basic_test_node(TEST_CIPHER_SUITE, "D").await.into();
486
487        NodeVec::from(nodes)
488    }
489}
490
491#[cfg(test)]
492mod tests {
493    use assert_matches::assert_matches;
494
495    use super::*;
496    use crate::{
497        client::test_utils::TEST_CIPHER_SUITE,
498        tree_kem::{
499            leaf_node::test_utils::get_basic_test_node, node::test_utils::get_test_node_vec,
500        },
501    };
502
503    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
504    async fn node_key_getters() {
505        let test_node_parent: Node = Parent {
506            public_key: b"pub".to_vec().into(),
507            parent_hash: ParentHash::empty(),
508            unmerged_leaves: vec![],
509        }
510        .into();
511
512        let test_leaf = get_basic_test_node(TEST_CIPHER_SUITE, "B").await;
513        let test_node_leaf: Node = test_leaf.clone().into();
514
515        assert_eq!(test_node_parent.public_key().as_ref(), b"pub");
516        assert_eq!(test_node_leaf.public_key(), &test_leaf.public_key);
517    }
518
519    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
520    async fn test_empty_leaves() {
521        let mut test_vec = get_test_node_vec().await;
522        let mut test_vec_clone = get_test_node_vec().await;
523        let empty_leaves: Vec<(LeafIndex, &mut Option<Node>)> = test_vec.empty_leaves().collect();
524        assert_eq!(
525            [(LeafIndex::unchecked(1), &mut test_vec_clone[2])].as_ref(),
526            empty_leaves.as_slice()
527        );
528    }
529
530    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
531    async fn test_direct_path() {
532        let test_vec = get_test_node_vec().await;
533        // Tree math is already tested in that module, just ensure equality
534        let expected = 0.direct_copath(&4);
535        let actual = test_vec.direct_copath(LeafIndex::unchecked(0));
536        assert_eq!(actual, expected);
537    }
538
539    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
540    async fn test_filtered_direct_path_co_path() {
541        let test_vec = get_test_node_vec().await;
542        let expected = [true, false];
543        let actual = test_vec.filtered(LeafIndex::unchecked(0)).unwrap();
544        assert_eq!(actual, expected);
545    }
546
547    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
548    async fn test_get_parent_node() {
549        let mut test_vec = get_test_node_vec().await;
550
551        // If the node is a leaf it should fail
552        assert!(test_vec.borrow_as_parent_mut(0).is_err());
553
554        // If the node index is out of range it should fail
555        assert!(test_vec
556            .borrow_as_parent_mut(test_vec.len() as u32)
557            .is_err());
558
559        // Otherwise it should succeed
560        let mut expected = Parent {
561            public_key: b"CD".to_vec().into(),
562            parent_hash: ParentHash::empty(),
563            unmerged_leaves: vec![LeafIndex::unchecked(2)],
564        };
565
566        assert_eq!(test_vec.borrow_as_parent_mut(5).unwrap(), &mut expected);
567    }
568
569    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
570    async fn test_get_resolution() {
571        let test_vec = get_test_node_vec().await;
572
573        let resolution_node_5 = test_vec.get_resolution_index(5).unwrap();
574        let resolution_node_2 = test_vec.get_resolution_index(2).unwrap();
575        let resolution_node_3 = test_vec.get_resolution_index(3).unwrap();
576
577        assert_eq!(&resolution_node_5, &[5, 4]);
578        assert!(resolution_node_2.is_empty());
579        assert_eq!(&resolution_node_3, &[0, 5, 4]);
580    }
581
582    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
583    async fn test_get_or_fill_existing() {
584        let mut test_vec = get_test_node_vec().await;
585        let mut test_vec2 = test_vec.clone();
586
587        let expected = test_vec[5].as_parent_mut().unwrap();
588        let actual = test_vec2
589            .borrow_or_fill_node_as_parent(5, &Vec::new().into())
590            .unwrap();
591
592        assert_eq!(actual, expected);
593    }
594
595    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
596    async fn test_get_or_fill_empty() {
597        let mut test_vec = get_test_node_vec().await;
598
599        let mut expected = Parent {
600            public_key: vec![0u8; 4].into(),
601            parent_hash: ParentHash::empty(),
602            unmerged_leaves: vec![],
603        };
604
605        let actual = test_vec
606            .borrow_or_fill_node_as_parent(1, &vec![0u8; 4].into())
607            .unwrap();
608
609        assert_eq!(actual, &mut expected);
610    }
611
612    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
613    async fn test_leaf_count() {
614        let test_vec = get_test_node_vec().await;
615        assert_eq!(test_vec.len(), 7);
616        assert_eq!(test_vec.occupied_leaf_count(), 3);
617        assert_eq!(
618            test_vec.non_empty_leaves().count(),
619            test_vec.occupied_leaf_count() as usize
620        );
621    }
622
623    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
624    async fn test_total_leaf_count() {
625        let test_vec = get_test_node_vec().await;
626        assert_eq!(test_vec.occupied_leaf_count(), 3);
627        assert_eq!(test_vec.total_leaf_count(), 4);
628    }
629
630    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
631    async fn max_leaf_index() {
632        let test_index = LeafIndex::try_from(1).unwrap();
633
634        let serialized = test_index.mls_encode_to_vec().unwrap();
635
636        LeafIndex::mls_decode(&mut &*serialized).unwrap();
637
638        #[cfg(feature = "serde")]
639        {
640            let serialized = serde_json::to_string(&test_index).unwrap();
641            serde_json::from_str::<LeafIndex>(&serialized).unwrap();
642        }
643    }
644
645    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
646    async fn max_leaf_index_failure() {
647        let res = LeafIndex::try_from(MAX_LEAF_INDEX + 1);
648        assert_matches!(res, Err(MlsError::InvalidTreeIndex));
649
650        let serialized = LeafIndex::unchecked(MAX_LEAF_INDEX + 1)
651            .mls_encode_to_vec()
652            .unwrap();
653
654        let res = LeafIndex::mls_decode(&mut &*serialized);
655        assert_matches!(res, Err(mls_rs_codec::Error::Custom(6)));
656
657        #[cfg(feature = "serde")]
658        {
659            let serialized =
660                serde_json::to_string(&LeafIndex::unchecked(MAX_LEAF_INDEX + 1)).unwrap();
661
662            let res: Result<LeafIndex, _> = serde_json::from_str(&serialized);
663
664            assert!(res.is_err())
665        }
666    }
667}
mls_rs/tree_kem/node.rs

mls_rs/tree_kem/
node.rs
