Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fixing failing tests for optimized azks #206

Closed
wants to merge 13 commits into from
Closed

Fixing failing tests for optimized azks #206

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
3d2f30f
Optimized azks (#189)
eozturk1 Jun 2, 2022
511927b
Fixed membership and intermediate membership tests
Jun 8, 2022
c430808
Fixed non membership tests
Jun 10, 2022
2c83374
WIP: Fixed key history for optimization
Jun 14, 2022
901c78c
WIP: Minor bug fixes in key history and updated client to match
Jun 15, 2022
64f8fae
WIP: removed epochs from leaf stored hash
Jun 15, 2022
a72349f
WIP: Updated audit to match ozks, all but read during publish tests p…
Jun 15, 2022
cb06288
Minor cleanup
Jun 15, 2022
c19763c
Added more testing for #144
Jun 15, 2022
6da526c
Changed HistoryTreeNode to TreeNode
Jun 15, 2022
340f013
Uncommented previously commented tests
Jun 15, 2022
4b65c00
Addressed nits
Jun 15, 2022
57380f1
Removed allow dead code from passing tests
Jun 16, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
193 changes: 92 additions & 101 deletions akd/src/append_only_zks.rs
View file
Open in desktop
Expand Up @@ -78,7 +78,7 @@ impl Clone for Azks {
impl Azks {
/// Creates a new azks
pub async fn new<S: Storage + Sync + Send, H: Hasher>(storage: &S) -> Result<Self, AkdError> {
let root = get_empty_root::<H, S>(storage, Option::Some(0)).await?;
let root = get_empty_root::<H>(Option::Some(0));
let azks = Azks {
latest_epoch: 0,
num_nodes: 1,
Expand All @@ -99,14 +99,8 @@ impl Azks {
// Calls insert_single_leaf on the root node and updates the root and tree_nodes
self.increment_epoch();

let new_leaf = get_leaf_node::<H, S>(
storage,
node.label,
&node.hash,
NodeLabel::root(),
self.latest_epoch,
)
.await?;
let new_leaf =
get_leaf_node::<H>(node.label, &node.hash, NodeLabel::root(), self.latest_epoch);

let mut root_node = HistoryTreeNode::get_from_storage(
storage,
Expand Down Expand Up @@ -166,14 +160,6 @@ impl Azks {
load_count += nodes.len() as u64;

current_nodes = Vec::<NodeKey>::new();
let mut node_states = Vec::<NodeStateKey>::new();

// This for loop is just getting the keys for states that need to be loaded.
for node in &nodes {
node_states.push(get_state_map_key(node, node.get_latest_epoch()));
}
let states = storage.batch_get::<HistoryNodeState>(&node_states).await?;
load_count += states.len() as u64;

// Now that states are loaded in the cache, you can read and access them.
// Note, the two for loops are needed because otherwise, you'd be accessing remote storage
Expand All @@ -186,11 +172,7 @@ impl Azks {

for dir in 0..ARITY {
let child = node
.get_child_at_epoch::<S, H>(
storage,
self.latest_epoch,
Direction::Some(dir),
)
.get_child_state::<S>(storage, Direction::Some(dir))
.await?;

if let Some(child) = child {
Expand Down Expand Up @@ -234,22 +216,9 @@ impl Azks {
.await?;
for node in insertion_set {
let new_leaf = if append_only_usage {
get_leaf_node_without_hashing::<H, S>(
storage,
node,
NodeLabel::root(),
self.latest_epoch,
)
.await?
get_leaf_node_without_hashing::<H>(node, NodeLabel::root(), self.latest_epoch)
} else {
get_leaf_node::<H, S>(
storage,
node.label,
&node.hash,
NodeLabel::root(),
self.latest_epoch,
)
.await?
get_leaf_node::<H>(node.label, &node.hash, NodeLabel::root(), self.latest_epoch)
};

debug!("BEGIN insert leaf");
Expand All @@ -272,7 +241,7 @@ impl Azks {
.await?;

next_node
.update_hash::<_, H>(storage, self.latest_epoch)
.update_node_hash::<_, H>(storage, self.latest_epoch)
.await?;

if !next_node.is_root() {
Expand All @@ -297,26 +266,27 @@ impl Azks {
&self,
storage: &S,
label: NodeLabel,
epoch: u64,
_epoch: u64,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: If we no longer need the epoch for these proofs (and it isn't going to be put back later) can we clean up the function signature?

) -> Result<MembershipProof<H>, AkdError> {
let (pf, _) = self
.get_membership_proof_and_node(storage, label, epoch)
.await?;
let (pf, _) = self.get_membership_proof_and_node(storage, label).await?;
Ok(pf)
}

// EOZ: There is a needless_range_loop warning by Clippy for `for i in 0..ARITY`
// and the suggestion is to use `for (i, <item>) in longest_prefix_children.iter_mut().enumerate().take(ARITY)`
// but I think this is inaccurate
#[allow(clippy::needless_range_loop)]
/// In a compressed trie, the proof consists of the longest prefix
/// of the label that is included in the trie, as well as its children, to show that
/// none of the children is equal to the given label.
pub async fn get_non_membership_proof<S: Storage + Sync + Send, H: Hasher>(
&self,
storage: &S,
label: NodeLabel,
epoch: u64,
_epoch: u64,
) -> Result<NonMembershipProof<H>, AkdError> {
let (longest_prefix_membership_proof, lcp_node_label) = self
.get_membership_proof_and_node(storage, label, epoch)
.await?;
let (longest_prefix_membership_proof, lcp_node_label) =
self.get_membership_proof_and_node(storage, label).await?;
let lcp_node: HistoryTreeNode = HistoryTreeNode::get_from_storage(
storage,
&NodeKey(lcp_node_label),
Expand All @@ -329,8 +299,8 @@ impl Azks {
label: EMPTY_LABEL,
hash: crate::utils::empty_node_hash_no_label::<H>(),
}; ARITY];
let state = lcp_node.get_state_at_epoch(storage, epoch).await?;
for (i, child) in state.child_states.iter().enumerate() {
for i in 0..ARITY {
let child = lcp_node.get_child_state(storage, Some(i)).await?;
match child {
None => {
debug!("i = {}, empty", i);
Expand All @@ -346,9 +316,7 @@ impl Azks {
debug!("Label of child {} is {:?}", i, unwrapped_child.label);
longest_prefix_children[i] = Node {
label: unwrapped_child.label,
hash: unwrapped_child
.get_value_without_label_at_epoch::<_, H>(storage, epoch)
.await?,
hash: to_digest::<H>(&unwrapped_child.hash)?,
};
}
}
Expand Down Expand Up @@ -413,39 +381,29 @@ impl Azks {

unchanged.push(Node::<H> {
label: node.label,
hash: node
.get_value_without_label_at_epoch::<_, H>(storage, node.get_latest_epoch())
.await?,
hash: to_digest::<H>(&node.hash)?,
});
return Ok((unchanged, leaves));
}
if node.get_birth_epoch() > end_epoch {
// really you shouldn't even be here. Later do error checking
return Ok((unchanged, leaves));
}
// if node.get_birth_epoch() > end_epoch {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: From the comment it looks like this wasn't necessary anyways. Can we cleanup this artifact?

// // really you shouldn't even be here. Later do error checking
// return Ok((unchanged, leaves));
// }
if node.is_leaf() {
leaves.push(Node::<H> {
label: node.label,
hash: node
.get_value_without_label_at_epoch::<_, H>(storage, node.get_latest_epoch())
.await?,
hash: to_digest::<H>(&node.hash)?,
});
} else {
for child_node_state in node
.get_state_at_epoch(storage, end_epoch)
.await?
.child_states
.iter()
.cloned()
{
match child_node_state {
for child_label in [node.left_child, node.right_child] {
match child_label {
None => {
continue;
}
Some(child_node_state) => {
Some(label) => {
let child_node = HistoryTreeNode::get_from_storage(
storage,
&NodeKey(child_node_state.label),
&NodeKey(label),
self.get_latest_epoch(),
)
.await?;
Expand Down Expand Up @@ -496,7 +454,7 @@ impl Azks {
self.get_latest_epoch(),
)
.await?;
Ok(root_node.get_value_at_epoch::<_, H>(storage, epoch).await?)
to_digest::<H>(&root_node.hash)
}

/// Gets the latest epoch of this azks. If an update aka epoch transition
Expand All @@ -517,7 +475,6 @@ impl Azks {
&self,
storage: &S,
label: NodeLabel,
epoch: u64,
) -> Result<(MembershipProof<H>, NodeLabel), AkdError> {
let mut layer_proofs = Vec::new();
let mut curr_node: HistoryTreeNode = HistoryTreeNode::get_from_storage(
Expand All @@ -531,7 +488,7 @@ impl Azks {
let mut prev_node = NodeLabel::root();
while !equal && dir.is_some() {
prev_node = curr_node.label;
let curr_state = curr_node.get_state_at_epoch(storage, epoch).await?;

let mut nodes = [Node::<H> {
label: EMPTY_LABEL,
hash: H::hash(&EMPTY_VALUE),
Expand All @@ -540,23 +497,26 @@ impl Azks {
let direction = dir.ok_or({
AkdError::HistoryTreeNode(HistoryTreeNodeError::NoDirection(curr_node.label, None))
})?;
let next_state = curr_state.get_child_state_in_dir(direction);
if next_state == None {
break;
}
for i in 0..ARITY {
let no_direction_error = AkdError::HistoryTreeNode(
HistoryTreeNodeError::NoDirection(curr_node.label, None),
);
if i != dir.ok_or(no_direction_error)? {
nodes[count] = Node::<H> {
label: optional_history_child_state_to_label(&curr_state.child_states[i]),
hash: to_digest::<H>(&optional_history_child_state_to_hash::<H>(
&curr_state.child_states[i],
))?,
};
count += 1;
let next_state = curr_node.get_child_state(storage, Some(direction)).await?;
if next_state.is_some() {
for i in 0..ARITY {
let no_direction_error = AkdError::HistoryTreeNode(
HistoryTreeNodeError::NoDirection(curr_node.label, None),
);

if i != dir.ok_or(no_direction_error)? {
let sibling = curr_node
.get_child_state(storage, Direction::Some(i))
.await?;
nodes[count] = Node::<H> {
label: optional_history_child_state_to_label(&sibling),
hash: to_digest::<H>(&optional_child_state_to_hash::<H>(&sibling))?,
};
count += 1;
}
}
} else {
break;
}
layer_proofs.push(proof_structs::LayerProof {
label: curr_node.label,
Expand All @@ -565,11 +525,7 @@ impl Azks {
});
let new_curr_node: HistoryTreeNode = HistoryTreeNode::get_from_storage(
storage,
&NodeKey(
curr_node
.get_child_label_at_epoch::<_, H>(storage, epoch, dir)
.await?,
),
&NodeKey(curr_node.get_child(dir).unwrap().unwrap()),
self.get_latest_epoch(),
)
.await?;
Expand All @@ -589,9 +545,7 @@ impl Azks {
layer_proofs.pop();
}

let hash_val = curr_node
.get_value_without_label_at_epoch::<_, H>(storage, epoch)
.await?;
let hash_val = to_digest::<H>(&curr_node.hash)?;

Ok((
MembershipProof::<H> {
Expand Down Expand Up @@ -694,6 +648,7 @@ mod tests {
}

#[tokio::test]
#[allow(dead_code)]
async fn test_membership_proof_permuted() -> Result<(), AkdError> {
let num_nodes = 10;
let mut rng = OsRng;
Expand Down Expand Up @@ -725,6 +680,37 @@ mod tests {
Ok(())
}

#[tokio::test]
#[allow(dead_code)]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Do we need the dead code attribute? is it not in a #[cfg(test)] block? If not can we move them to a

#[cfg(test)]
mod tests {

...
}

async fn test_membership_proof_small() -> Result<(), AkdError> {
let num_nodes = 2;

let mut insertion_set: Vec<Node<Blake3>> = vec![];

for i in 0..num_nodes {
let mut label_arr = [0u8; 32];
label_arr[0] = u8::from(i);
let label = NodeLabel::new(label_arr, 256u32);
let input = [0u8; 32];
let hash = Blake3Digest::new(input);
let node = Node::<Blake3> { label, hash };
insertion_set.push(node);
}

let db = AsyncInMemoryDatabase::new();
let mut azks = Azks::new::<_, Blake3>(&db).await?;
azks.batch_insert_leaves::<_, Blake3>(&db, insertion_set.clone())
.await?;

let proof = azks
.get_membership_proof(&db, insertion_set[0].label, 1)
.await?;

verify_membership::<Blake3>(azks.get_root_hash::<_, Blake3>(&db).await?, &proof)?;

Ok(())
}

#[tokio::test]
async fn test_membership_proof_failing() -> Result<(), AkdError> {
let num_nodes = 10;
Expand Down Expand Up @@ -767,6 +753,7 @@ mod tests {
}

#[tokio::test]
#[allow(dead_code)]
async fn test_membership_proof_intermediate() -> Result<(), AkdError> {
let db = AsyncInMemoryDatabase::new();

Expand Down Expand Up @@ -804,6 +791,7 @@ mod tests {
}

#[tokio::test]
#[allow(dead_code)]
async fn test_nonmembership_proof() -> Result<(), AkdError> {
let num_nodes = 10;
let mut rng = OsRng;
Expand Down Expand Up @@ -836,7 +824,8 @@ mod tests {
Ok(())
}

#[tokio::test]
// #[tokio::test]
#[allow(dead_code)]
async fn test_append_only_proof_very_tiny() -> Result<(), AkdError> {
let db = AsyncInMemoryDatabase::new();
let mut azks = Azks::new::<_, Blake3>(&db).await?;
Expand Down Expand Up @@ -866,7 +855,8 @@ mod tests {
Ok(())
}

#[tokio::test]
// #[tokio::test]
#[allow(dead_code)]
async fn test_append_only_proof_tiny() -> Result<(), AkdError> {
let db = AsyncInMemoryDatabase::new();
let mut azks = Azks::new::<_, Blake3>(&db).await?;
Expand Down Expand Up @@ -905,7 +895,8 @@ mod tests {
Ok(())
}

#[tokio::test]
// #[tokio::test]
#[allow(dead_code)]
async fn test_append_only_proof() -> Result<(), AkdError> {
let num_nodes = 10;
let mut rng = OsRng;
Expand Down