Skip to content

Commit

Permalink
Remove debug output & fix test
Browse files Browse the repository at this point in the history
  • Loading branch information
@raphaelrobert
raphaelrobert committed Nov 14, 2024
1 parent 0e4ee20 commit c92b46c
Showing 1 changed file with 12 additions and 168 deletions.
180 changes: 12 additions & 168 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,23 +178,16 @@ impl WrappingTree {

// We populate the tree with the public nodes and decrypt them
// starting from the root
println!(
"Expanding nodes with root key {:?}",
derive_node_key(&init_secret, root)
);
match public_tree.resolution(root) {
Resolution::Empty => {
// The root node is blank, nothing to do
println!("Root node's resolution is empty");
}
Resolution::One(child) => {
// The root node has only one child, we start from there
println!("Root node has only one child: {:?}", child);
tree.expand_nodes(child, derive_node_key(&init_secret, child), &public_tree);
}
Resolution::Both(_, _) => {
// The root node has two children, we stay at the root
println!("Root node has two children");
tree.expand_nodes(root, derive_node_key(&init_secret, root), &public_tree);
}
}
Expand All @@ -210,25 +203,14 @@ impl WrappingTree {
// If the leaf node is present in the public tree, we decrypt it
// and store it in the tree
if let Some(ciphertext) = public_tree.leaf_nodes[l.usize()] {
println!("Decrypting leaf node {} with key {:?}", l.usize(), key);
self.leaf_nodes[l.usize()] = self.decrypt_leaf_node(ciphertext, &key);
debug_assert_eq!(
self.leaf_nodes[l.usize()]
.clone()
.content
.unwrap()
.plaintext
.key[..32],
[1u8; 32]
);
}
}
TreeNodeIndex::Parent(p) => {
// If the parent node is present in the public tree, we decrypt it
// and store it in the tree. We also extract the keys for the children.
let (left_key, right_key) =
if let Some(ciphertext) = public_tree.parent_nodes[p.usize()] {
println!("Decrypting parent node {} with key {:?}", p.usize(), key);
let node = self.decrypt_parent_node(ciphertext, &key);
let left_key = node.content.as_ref().map(|c| c.plaintext.left_key).unwrap();
let right_key = node
Expand Down Expand Up @@ -307,11 +289,7 @@ impl WrappingTree {

// Wrap the leaf node
let leaf_key = derive_node_key(&self.init_secret, TreeNodeIndex::Leaf(index));
println!(
"Encrypting leaf node {} with key {:?}",
index.usize(),
leaf_key
);

self.leaf_nodes[index.usize()] = encrypt_leaf_node(&leaf_key, leaf_node_plaintext);

// Wrap the parent nodes up to the root
Expand All @@ -322,13 +300,7 @@ impl WrappingTree {

/// Remove a node from the tree
pub fn remove(&mut self, index: LeafNodeIndex) {
println!(
"Removing node {:?}, leaf node length: {}",
index,
self.leaf_nodes.len()
);
if self.leaf_nodes.len() == 0 {
println!("Tree is empty, nothing to remove");
// The tree is empty, nothing to do
return;
}
Expand Down Expand Up @@ -356,10 +328,8 @@ impl WrappingTree {
}

let right_most_leaf_index = if let Some(right_most_leaf_index) = right_most_leaf_index {
println!("Right most leaf index: {:?}", right_most_leaf_index);
right_most_leaf_index
} else {
println!("Tree is empty, clearing it");
// The tree is empty, we can clear it
self.leaf_nodes.clear();
self.parent_nodes.clear();
Expand All @@ -368,11 +338,9 @@ impl WrappingTree {

// Wrap the parent nodes up to the root
if self.leaf_nodes.len() > 1 {
println!("Wrapping up from node {:?}", index);
self.wrap_up(index)
} else {
// If the tree only has one node, we rewrap it since it's the root
println!("Rewrapping root/leaf node after removal");
self.rewrap_leaf(LeafNodeIndex::new(0));
}

Expand All @@ -393,24 +361,19 @@ impl WrappingTree {

fn rewrap_leaf(&mut self, index: LeafNodeIndex) {
if index.usize() >= self.leaf_nodes.len() {
println!("Rewrap leaf: Index is outside the tree, nothing to do");
// The index is outside the tree, nothing to do
return;
}

let key = derive_node_key(&self.init_secret, TreeNodeIndex::Leaf(index));

println!("Rewrapping leaf node {:?}", index);

if let Some(ref leaf_node) = self.leaf_nodes[index.usize()].content {
self.leaf_nodes[index.usize()] = encrypt_leaf_node(&key, leaf_node.plaintext.clone());
}
}

fn rewrap_parent(&mut self, index: ParentNodeIndex) {
println!("Rewrapping parent node {:?}", index);
if index.usize() >= self.parent_nodes.len() {
println!("Rewrap parent: Index is outside the tree, nothing to do");
// The index is outside the tree, nothing to do
return;
}
Expand All @@ -419,25 +382,19 @@ impl WrappingTree {

match res {
Resolution::Empty => {
panic!("rewrap_parent: Resolution is empty");
// This should never happen
}
Resolution::One(child) => {
println!("Rewrapping partial parent node {:?}", child);

match child {
TreeNodeIndex::Leaf(leaf_node_index) => {
self.rewrap_leaf(leaf_node_index);
}
TreeNodeIndex::Parent(parent_node_index) => {
self.rewrap_parent(parent_node_index);
}
Resolution::One(child) => match child {
TreeNodeIndex::Leaf(leaf_node_index) => {
self.rewrap_leaf(leaf_node_index);
}
}
TreeNodeIndex::Parent(parent_node_index) => {
self.rewrap_parent(parent_node_index);
}
},
Resolution::Both(_, _) => {
let key = derive_node_key(&self.init_secret, TreeNodeIndex::Parent(index));

println!("Rewrapping full parent node {:?}", index);

if let Some(ref parent_node) = self.parent_nodes[index.usize()].content {
self.parent_nodes[index.usize()] =
encrypt_parent_node(&key, parent_node.plaintext.clone());
Expand All @@ -453,7 +410,6 @@ impl WrappingTree {
let root = root(self.size());

if index.usize() >= self.leaf_nodes.len() {
println!("Wrap up: Index is outside the tree, nothing to do");
// The index is outside the tree, nothing to do
return;
}
Expand All @@ -470,7 +426,6 @@ impl WrappingTree {
continue;
}
Resolution::One(child) => {
println!("Wrap up: one with node {:?}", node_index);
// The node has only one child, we need to skip it
// and we blank it, unless it is the root node.
if TreeNodeIndex::from(node_index) != root {
Expand All @@ -480,26 +435,20 @@ impl WrappingTree {
{
last_key = if child.u32() < TreeNodeIndex::from(node_index).u32() {
// The child was on the left, hence we keep the left key
println!("\tThe child was on the left");
node_content.plaintext.left_key
} else {
// The child was on the right, hence we keep the right key
println!("\tThe child was on the right");
node_content.plaintext.right_key
};
println!("\tExtracting key from parent node {:?}", last_key);
}

println!("\tBlanking node {:?}", node_index);
self.parent_nodes[node_index.usize()] = ParentNode { content: None };

continue;
} else {
println!("Wrapped up to root node, but there is only one child");
// Blank the root node
self.parent_nodes[node_index.usize()] = ParentNode { content: None };
// Rewrap child
println!("\tRewrapping child node {:?}", child);
match child {
TreeNodeIndex::Leaf(leaf_node_index) => {
self.rewrap_leaf(leaf_node_index);
Expand All @@ -517,17 +466,11 @@ impl WrappingTree {
// want to derive a key for the node in the direct path and
// keep the key for the other one.

println!(
"Wrap up: both with node {:?}\n\tLeft: {:?}\n\tRight: {:?}",
node_index, left_index, right_index
);

let left_key;
let right_key;

if last_index.u32() < TreeNodeIndex::from(node_index).u32() {
// We came from the left, hence we keep the right key
println!("We came from the left");
left_key = last_key;
right_key = self
.parent_nodes
Expand All @@ -538,7 +481,6 @@ impl WrappingTree {
// find next key instead
} else {
// We came from the right, hence we keep the left key
println!("We came from the right");
left_key = self
.parent_nodes
.get(node_index.usize())
Expand All @@ -548,23 +490,14 @@ impl WrappingTree {
right_key = last_key;
}

println!(
"Parent node {:?}: \n\tLeft key: {:?}, \n\tRight key: {:?}",
node_index, left_key, right_key
);

let parent_node_plaintext = ParentNodePlaintext {
left_key,
right_key,
};

// Encrypt the two keys and store the new parent node
last_key = derive_node_key(&self.init_secret, node_index.into());
println!(
"Encrypting parent node {} with key {:?}",
node_index.usize(),
last_key
);

let new_node = encrypt_parent_node(&last_key, parent_node_plaintext);
self.parent_nodes[node_index.usize()] = new_node;
}
Expand Down Expand Up @@ -595,27 +528,17 @@ impl WrappingTree {
let root = root(TreeSize::from_leaf_count(self.leaf_nodes.len()));
let mut path_index = index;

println!("next_or_derive: {:?}", index);

while path_index != root {
let parent = parent(path_index);
println!("Looking for key in node {:?}", parent);
if let Some(content) = &self.parent_nodes[parent.usize()].content {
if path_index > TreeNodeIndex::Parent(parent) {
println!("Found right key");
return content.plaintext.right_key;
} else {
println!("Found left key");
return content.plaintext.left_key;
}
}
path_index = TreeNodeIndex::Parent(parent);
}
println!(
"Deriving key for index {:?}: {:?}",
index,
derive_node_key(&self.init_secret, index)
);
derive_node_key(&self.init_secret, index)
}

Expand Down Expand Up @@ -800,19 +723,14 @@ fn fuzz() {
mac: [2u8; 32],
};

for epoch in 0..EPOCH_COUNT {
println!("--- Epoch: {}", epoch);

for operation in 0..OPERATION_COUNT {
println!("-- Operation: {}", operation);
for _epoch in 0..EPOCH_COUNT {
for _operation in 0..OPERATION_COUNT {
let index = LeafNodeIndex::new(rand::random::<u32>() % LEAF_COUNT);
let add = rand::random::<bool>();

if add {
println!("+ {:?}", index);
tree.add(index, leaf_node_plaintext.clone());
} else {
println!("- {:?}", index);
tree.remove(index);
}

Expand All @@ -825,77 +743,3 @@ fn fuzz() {
eval_tree(&tree);
}
}

#[test]
fn pathological_case2() {
let init_secret = [3u8; 32];
let mut tree = WrappingTree::new(init_secret);
let leaf_node_plaintext = LeafNodePlaintext {
key: [1u8; 32],
mac: [2u8; 32],
};

tree.add(LeafNodeIndex::new(3), leaf_node_plaintext.clone());
tree.add(LeafNodeIndex::new(1), leaf_node_plaintext.clone());
tree.add(LeafNodeIndex::new(0), leaf_node_plaintext.clone());

println!("First tree: {:?}", tree.export_public_tree().leaf_nodes);

let public_tree = tree.export_public_tree();
let root_secret = tree.export_root_secret();
let new_tree = WrappingTree::from_public_tree(public_tree, *root_secret);
let leaf_node_plaintexts = tree.export_leaf_node_plaintexts();
let new_leaf_node_plaintexts = new_tree.export_leaf_node_plaintexts();
assert_eq!(leaf_node_plaintexts, new_leaf_node_plaintexts);

println!("Second epoch");
let new_init_secret = [5u8; 32];
tree.new_epoch(new_init_secret);

println!("Exporting tree");

let public_tree = tree.export_public_tree();
let root_secret = tree.export_root_secret();
let new_tree = WrappingTree::from_public_tree(public_tree, *root_secret);
let leaf_node_plaintexts = tree.export_leaf_node_plaintexts();
let new_leaf_node_plaintexts = new_tree.export_leaf_node_plaintexts();
assert_eq!(leaf_node_plaintexts, new_leaf_node_plaintexts);

println!("Second tree: {:?}", tree.export_public_tree().leaf_nodes);

let remove_node = 0;

println!("Removing node {}", remove_node);

tree.remove(LeafNodeIndex::new(remove_node));

println!("Node {} removed", remove_node);

println!("Third tree: {:?}", tree.export_public_tree().leaf_nodes);

let public_tree = tree.export_public_tree();
let root_secret = tree.export_root_secret();
let new_tree = WrappingTree::from_public_tree(public_tree, *root_secret);
let leaf_node_plaintexts = tree.export_leaf_node_plaintexts();
let new_leaf_node_plaintexts = new_tree.export_leaf_node_plaintexts();
assert_eq!(leaf_node_plaintexts, new_leaf_node_plaintexts);

let new_init_secret = [7u8; 32];
tree.new_epoch(new_init_secret);

let public_tree = tree.export_public_tree();
let root_secret = tree.export_root_secret();
let new_tree = WrappingTree::from_public_tree(public_tree, *root_secret);
let leaf_node_plaintexts = tree.export_leaf_node_plaintexts();
let new_leaf_node_plaintexts = new_tree.export_leaf_node_plaintexts();
assert_eq!(leaf_node_plaintexts, new_leaf_node_plaintexts);
}

#[test]
fn root_test() {
assert_eq!(root(TreeSize::new(1)).u32(), 0);
assert_eq!(TreeSize::new(1).parent_count(), 0);
assert_eq!(TreeSize::new(1).leaf_count(), 1);
assert_eq!(TreeSize::new(0).parent_count(), 0);
assert_eq!(TreeSize::new(0).leaf_count(), 0);
}

0 comments on commit c92b46c

Please sign in to comment.