Make secret tree work with arbitrary indices, not only usize

mls-rs/src/group/ciphertext_processor.rs

@@ -16,7 +16,10 @@ use super::{
     secret_tree::{KeyType, MessageKeyData},
     GroupContext,
 };
-use crate::{client::MlsError, tree_kem::node::LeafIndex};
+use crate::{
+    client::MlsError,
+    tree_kem::node::{LeafIndex, NodeIndex},
+};
 use mls_rs_codec::MlsEncode;
 use mls_rs_core::{crypto::CipherSuiteProvider, error::IntoAnyError};
 use zeroize::Zeroizing;
@@ -67,7 +70,7 @@ where
         &mut self,
         key_type: KeyType,
     ) -> Result<MessageKeyData, MlsError> {
-        let self_index = self.group_state.self_index();
+        let self_index = NodeIndex::from(self.group_state.self_index());
 
         self.group_state
             .epoch_secrets_mut()
@@ -83,6 +86,8 @@ where
         key_type: KeyType,
         generation: u32,
     ) -> Result<MessageKeyData, MlsError> {
+        let sender = NodeIndex::from(sender);
+
         self.group_state
             .epoch_secrets_mut()
             .secret_tree

mls-rs/src/group/epoch.rs

@@ -4,6 +4,7 @@
 
 #[cfg(feature = "psk")]
 use crate::psk::PreSharedKey;
+use crate::tree_kem::node::NodeIndex;
 #[cfg(feature = "prior_epoch")]
 use crate::{crypto::SignaturePublicKey, group::GroupContext, tree_kem::node::LeafIndex};
 use alloc::vec::Vec;
@@ -66,7 +67,7 @@ pub(crate) struct EpochSecrets {
     #[mls_codec(with = "mls_rs_codec::byte_vec")]
     pub(crate) sender_data_secret: SenderDataSecret,
     #[cfg(any(feature = "secret_tree_access", feature = "private_message"))]
-    pub(crate) secret_tree: SecretTree,
+    pub(crate) secret_tree: SecretTree<NodeIndex>,
 }
 
 #[derive(Clone, Debug, PartialEq, MlsEncode, MlsDecode, MlsSize)]

mls-rs/src/group/external_commit.rs

@@ -190,7 +190,7 @@ impl<C: ClientConfig> ExternalCommitBuilder<C> {
             resumption_secret: PreSharedKey::new(vec![]),
             sender_data_secret: SenderDataSecret::from(vec![]),
             #[cfg(any(feature = "secret_tree_access", feature = "private_message"))]
-            secret_tree: SecretTree::empty(),
+            secret_tree: SecretTree::empty(0),
         };
 
         let (mut group, _) = Group::join_with(

mls-rs/src/group/interop_test_vectors/tree_kem.rs

@@ -102,11 +102,13 @@ async fn tree_kem() {
             let mut tree_private = TreeKemPrivate::new(LeafIndex(leaf.index));
 
             // Set and validate HPKE keys on direct path
-            let path = tree.nodes.direct_path(tree_private.self_index).unwrap();
+            let path = tree.nodes.direct_copath(tree_private.self_index);
 
             tree_private.secret_keys = Vec::new();
 
             for dp in path {
+                let dp = dp.path;
+
                 let secret = leaf
                     .path_secrets
                     .iter()

mls-rs/src/group/mod.rs

@@ -719,14 +719,14 @@ where
         let path = provisional_state
             .public_tree
             .nodes
-            .direct_path(self_index)?;
+            .direct_copath(self_index);
 
         provisional_private_tree
             .secret_keys
             .resize(path.len() + 1, None);
 
         for (i, n) in path.iter().enumerate() {
-            if provisional_state.public_tree.nodes.is_blank(*n)? {
+            if provisional_state.public_tree.nodes.is_blank(n.path)? {
                 provisional_private_tree.secret_keys[i + 1] = None;
             }
         }
@@ -1512,7 +1512,7 @@ where
     pub fn next_encryption_key(&mut self) -> Result<MessageKey, MlsError> {
         self.epoch_secrets.secret_tree.next_message_key(
             &self.cipher_suite_provider,
-            self.private_tree.self_index,
+            crate::tree_kem::node::NodeIndex::from(self.private_tree.self_index),
             KeyType::Application,
         )
     }
@@ -1525,7 +1525,7 @@ where
     ) -> Result<MessageKey, MlsError> {
         self.epoch_secrets.secret_tree.message_key_generation(
             &self.cipher_suite_provider,
-            LeafIndex(sender),
+            crate::tree_kem::node::NodeIndex::from(sender),
             KeyType::Application,
             generation,
         )