Revision control
Copy as Markdown
Other Tools
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// Copyright by contributors to this project.
// SPDX-License-Identifier: (Apache-2.0 OR MIT)
use alloc::{vec, vec::Vec};
use mls_rs_codec::{MlsDecode, MlsEncode, MlsSize};
use mls_rs_core::crypto::HpkeSecretKey;
use crate::{client::MlsError, crypto::CipherSuiteProvider};
use super::{
math::leaf_lca_level,
node::LeafIndex,
path_secret::{PathSecret, PathSecretGenerator},
TreeKemPublic,
};
#[derive(Clone, Debug, MlsEncode, MlsDecode, MlsSize, Eq, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[non_exhaustive]
pub struct TreeKemPrivate {
pub self_index: LeafIndex,
pub secret_keys: Vec<Option<HpkeSecretKey>>,
}
impl TreeKemPrivate {
pub fn new_self_leaf(self_index: LeafIndex, leaf_secret: HpkeSecretKey) -> Self {
TreeKemPrivate {
self_index,
secret_keys: vec![Some(leaf_secret)],
}
}
pub fn new_for_external() -> Self {
TreeKemPrivate {
self_index: LeafIndex(0),
secret_keys: Default::default(),
}
}
#[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
pub async fn update_secrets<P: CipherSuiteProvider>(
&mut self,
cipher_suite_provider: &P,
signer_index: LeafIndex,
path_secret: PathSecret,
public_tree: &TreeKemPublic,
) -> Result<(), MlsError> {
// Identify the lowest common
// ancestor of the leaves at index and at GroupInfo.signer_index. Set the private key
// for this node to the private key derived from the path_secret.
let lca_index = leaf_lca_level(self.self_index.into(), signer_index.into()) as usize - 2;
// For each parent of the common ancestor, up to the root of the tree, derive a new
// path secret and set the private key for the node to the private key derived from the
// path secret. The private key MUST be the private key that corresponds to the public
// key in the node.
let mut node_secret_gen =
PathSecretGenerator::starting_with(cipher_suite_provider, path_secret);
let path = public_tree.nodes.direct_copath(self.self_index);
let filtered = &public_tree.nodes.filtered(self.self_index)?;
self.secret_keys.resize(path.len() + 1, None);
for (i, (n, f)) in path.iter().zip(filtered).enumerate().skip(lca_index) {
if *f {
continue;
}
let secret = node_secret_gen.next_secret().await?;
let expected_pub_key = public_tree
.nodes
.borrow_node(n.path)?
.as_ref()
.map(|n| n.public_key())
.ok_or(MlsError::PubKeyMismatch)?;
let (secret_key, public_key) = secret.to_hpke_key_pair(cipher_suite_provider).await?;
if expected_pub_key != &public_key {
return Err(MlsError::PubKeyMismatch);
}
// It's ok to use index directly because of the resize above
self.secret_keys[i + 1] = Some(secret_key);
}
Ok(())
}
#[cfg(feature = "by_ref_proposal")]
pub fn update_leaf(&mut self, new_leaf: HpkeSecretKey) {
self.secret_keys = vec![None; self.secret_keys.len()];
self.secret_keys[0] = Some(new_leaf);
}
}
#[cfg(test)]
impl TreeKemPrivate {
pub fn new(self_index: LeafIndex) -> Self {
TreeKemPrivate {
self_index,
secret_keys: Default::default(),
}
}
}
#[cfg(test)]
mod tests {
use assert_matches::assert_matches;
use crate::{
cipher_suite::CipherSuite,
client::test_utils::TEST_CIPHER_SUITE,
crypto::test_utils::test_cipher_suite_provider,
group::test_utils::{get_test_group_context, random_bytes},
identity::basic::BasicIdentityProvider,
tree_kem::{
kem::TreeKem,
leaf_node::test_utils::{
default_properties, get_basic_test_node, get_basic_test_node_sig_key,
},
math::TreeIndex,
node::LeafIndex,
},
};
use super::*;
#[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
async fn random_hpke_secret_key() -> HpkeSecretKey {
let (secret, _) = test_cipher_suite_provider(TEST_CIPHER_SUITE)
.kem_derive(&random_bytes(32))
.await
.unwrap();
secret
}
#[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
async fn test_create_self_leaf() {
let secret = random_hpke_secret_key().await;
let self_index = LeafIndex(42);
let private_key = TreeKemPrivate::new_self_leaf(self_index, secret.clone());
assert_eq!(private_key.self_index, self_index);
assert_eq!(private_key.secret_keys.len(), 1);
assert_eq!(private_key.secret_keys[0].as_ref().unwrap(), &secret)
}
// Create a ratchet tree for Alice, Bob and Charlie. Alice generates an update path for
// Charlie. Return (Public Tree, Charlie's private key, update path, path secret)
// The ratchet tree returned has leaf indexes as [alice, bob, charlie]
#[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
async fn update_secrets_setup(
cipher_suite: CipherSuite,
) -> (TreeKemPublic, TreeKemPrivate, TreeKemPrivate, PathSecret) {
let cipher_suite_provider = test_cipher_suite_provider(cipher_suite);
let (alice_leaf, alice_hpke_secret, alice_signing) =
get_basic_test_node_sig_key(cipher_suite, "alice").await;
let bob_leaf = get_basic_test_node(cipher_suite, "bob").await;
let (charlie_leaf, charlie_hpke_secret, _charlie_signing) =
get_basic_test_node_sig_key(cipher_suite, "charlie").await;
// Create a new public tree with Alice
let (mut public_tree, mut alice_private) = TreeKemPublic::derive(
alice_leaf,
alice_hpke_secret,
&BasicIdentityProvider,
&Default::default(),
)
.await
.unwrap();
// Add bob and charlie to the tree
public_tree
.add_leaves(
vec![bob_leaf, charlie_leaf],
&BasicIdentityProvider,
&cipher_suite_provider,
)
.await
.unwrap();
// Alice's secret key is longer now
alice_private.secret_keys.resize(3, None);
// Generate an update path for Alice
let encap_gen = TreeKem::new(&mut public_tree, &mut alice_private)
.encap(
&mut get_test_group_context(42, cipher_suite).await,
&[],
&alice_signing,
default_properties(),
None,
&cipher_suite_provider,
#[cfg(test)]
&Default::default(),
)
.await
.unwrap();
// Get a path secret from Alice for Charlie
let path_secret = encap_gen.path_secrets[1].clone().unwrap();
// Private key for Charlie
let charlie_private = TreeKemPrivate::new_self_leaf(LeafIndex(2), charlie_hpke_secret);
(public_tree, charlie_private, alice_private, path_secret)
}
#[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
async fn test_update_secrets() {
let cipher_suite = TEST_CIPHER_SUITE;
let (public_tree, mut charlie_private, alice_private, path_secret) =
update_secrets_setup(cipher_suite).await;
let existing_private = charlie_private.secret_keys.first().cloned().unwrap();
// Add the secrets for Charlie to his private key
charlie_private
.update_secrets(
&test_cipher_suite_provider(cipher_suite),
LeafIndex(0),
path_secret,
&public_tree,
)
.await
.unwrap();
// Make sure that Charlie's private key didn't lose keys
assert_eq!(charlie_private.secret_keys.len(), 3);
// Check that the intersection of the secret keys of Alice and Charlie matches.
// The intersection contains only the root.
assert_eq!(alice_private.secret_keys[2], charlie_private.secret_keys[2]);
assert_eq!(
charlie_private.secret_keys[0].as_ref(),
existing_private.as_ref()
);
}
#[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
async fn test_update_secrets_key_mismatch() {
let cipher_suite = TEST_CIPHER_SUITE;
let (mut public_tree, mut charlie_private, _, path_secret) =
update_secrets_setup(cipher_suite).await;
// Sabotage the public tree
public_tree
.nodes
.borrow_as_parent_mut(public_tree.total_leaf_count().root())
.unwrap()
.public_key = random_bytes(32).into();
// Add the secrets for Charlie to his private key
let res = charlie_private
.update_secrets(
&test_cipher_suite_provider(cipher_suite),
LeafIndex(0),
path_secret,
&public_tree,
)
.await;
assert_matches!(res, Err(MlsError::PubKeyMismatch));
}
#[cfg(feature = "by_ref_proposal")]
#[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
async fn setup_direct_path(self_index: LeafIndex, leaf_count: u32) -> TreeKemPrivate {
let secret = random_hpke_secret_key().await;
let mut private_key = TreeKemPrivate::new_self_leaf(self_index, secret.clone());
private_key.secret_keys = (0..0.direct_copath(&leaf_count).len() + 1)
.map(|_| Some(secret.clone()))
.collect();
private_key
}
#[cfg(feature = "by_ref_proposal")]
#[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
async fn test_update_leaf() {
let self_leaf = LeafIndex(42);
let mut private_key = setup_direct_path(self_leaf, 128).await;
let new_secret = random_hpke_secret_key().await;
private_key.update_leaf(new_secret.clone());
// The update operation should have removed all the other keys in our direct path we
// previously added
assert!(private_key.secret_keys.iter().skip(1).all(|n| n.is_none()));
// The secret key for our leaf should have been updated accordingly
assert_eq!(private_key.secret_keys.first().unwrap(), &Some(new_secret));
}
}