Skip to content

Commit

Permalink
chore: use decode in range proof verification (#303)
Browse files Browse the repository at this point in the history
  • Loading branch information
xinifinity authored Oct 3, 2023
1 parent 87f5599 commit 772035e
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 121 deletions.
1 change: 0 additions & 1 deletion firewood/src/merkle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@ mod node;
mod partial_path;
mod trie_hash;

pub(crate) use node::Encoded;
pub use node::{BranchNode, Data, ExtNode, LeafNode, Node, NodeType, NBRANCH};
pub use partial_path::PartialPath;
pub use trie_hash::{TrieHash, TRIE_HASH_LEN};
Expand Down
153 changes: 33 additions & 120 deletions firewood/src/proof.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,23 +4,18 @@
use std::cmp::Ordering;
use std::ops::Deref;

use bincode::Options;
use nix::errno::Errno;
use sha3::Digest;
use shale::disk_address::DiskAddress;
use shale::ShaleError;
use shale::ShaleStore;
use thiserror::Error;

use crate::merkle::Encoded;
use crate::nibbles::Nibbles;
use crate::nibbles::NibblesIterator;
use crate::{
db::DbError,
merkle::{
to_nibble_array, BranchNode, ExtNode, LeafNode, Merkle, MerkleError, Node, NodeType,
PartialPath, NBRANCH,
},
merkle::{to_nibble_array, Merkle, MerkleError, Node, NodeType},
merkle_util::{new_merkle, DataStoreError, MerkleSetup},
v2::api::Proof,
};
Expand Down Expand Up @@ -95,9 +90,6 @@ impl From<DbError> for ProofError {
}
}

const EXT_NODE_SIZE: usize = 2;
const BRANCH_NODE_SIZE: usize = 17;

/// SubProof contains the encoded value and the hash value of a node that maps
/// to a single proof step. If reaches an end step during proof verification,
/// the hash value will be none, and the encoded value will be the value of the
Expand Down Expand Up @@ -455,93 +447,50 @@ impl<N: AsRef<[u8]> + Send> Proof<N> {
buf: &[u8],
end_node: bool,
) -> Result<(DiskAddress, Option<SubProof>, usize), ProofError> {
let mut items: Vec<Encoded<Vec<u8>>> = bincode::DefaultOptions::new().deserialize(buf)?;
let size = items.len();

match size {
EXT_NODE_SIZE => {
let mut items = items.into_iter();

let cur_key_path: Vec<u8> = items
.next()
.unwrap()
.decode()?
.into_iter()
.flat_map(to_nibble_array)
.collect();

let (cur_key_path, term) = PartialPath::decode(&cur_key_path);
let cur_key = cur_key_path.into_inner();

let data: Vec<u8> = items.next().unwrap().decode()?;

let node = NodeType::decode(buf)?;
let new_node = merkle
.new_node(Node::new(node))
.map_err(ProofError::InvalidNode)?;
let addr = new_node.as_ptr();
match new_node.inner() {
NodeType::Leaf(n) => {
let cur_key = &n.path().0;
// Check if the key of current node match with the given key.
if key.len() < cur_key.len() || key[..cur_key.len()] != cur_key {
let ext_ptr = get_ext_ptr(merkle, term, Data(data), CurKey(cur_key))?;

return Ok((ext_ptr, None, 0));
if key.len() < cur_key.len() || &key[..cur_key.len()] != cur_key {
return Ok((addr, None, 0));
}

let subproof = if term {
Some(SubProof {
encoded: data.clone(),
hash: None,
})
} else {
generate_subproof(data.clone()).map(Some)?
};

let cur_key_len = cur_key.len();

let ext_ptr = get_ext_ptr(merkle, term, Data(data), CurKey(cur_key))?;

Ok((ext_ptr, subproof, cur_key_len))
let subproof = Some(SubProof {
encoded: n.data().to_vec(),
hash: None,
});
Ok((addr, subproof, cur_key.len()))
}
NodeType::Extension(n) => {
let cur_key = &n.path().0;

BRANCH_NODE_SIZE => {
// we've already validated the size, that's why we can safely unwrap
let data = items.pop().unwrap().decode()?;
// Extract the value of the branch node and set to None if it's an empty Vec
let value = Some(data).filter(|data| !data.is_empty());

// Record encoded values of all children.
let mut chd_encoded: [Option<Vec<u8>>; NBRANCH] = Default::default();

// we popped the last element, so their should only be NBRANCH items left
for (i, chd) in items.into_iter().enumerate() {
let data = chd.decode()?;
chd_encoded[i] = Some(data).filter(|data| !data.is_empty());
}

// If the node is the last one to be decoded, then no subproof to be extracted.
if end_node {
let branch_ptr = build_branch_ptr(merkle, value, chd_encoded)?;

return Ok((branch_ptr, None, 1));
// Check if the key of current node match with the given key.
if key.len() < cur_key.len() || &key[..cur_key.len()] != cur_key {
return Ok((addr, None, 0));
}

if key.is_empty() {
return Err(ProofError::NoSuchNode);
}
let encoded = n.chd_encoded().ok_or(ProofError::InvalidData)?.to_vec();
let subproof = generate_subproof(encoded).map(Some)?;

Ok((addr, subproof, cur_key.len()))
}
// If the node is the last one to be decoded, then no subproof to be extracted.
NodeType::Branch(_) if end_node => Ok((addr, None, 1)),
NodeType::Branch(_) if key.is_empty() => Err(ProofError::NoSuchNode),
NodeType::Branch(n) => {
// Check if the subproof with the given key exist.
let index = key[0] as usize;

let Some(data) = chd_encoded[index].clone() else {
let branch_ptr = build_branch_ptr(merkle, value, chd_encoded)?;

return Ok((branch_ptr, None, 1));
let Some(data) = &n.chd_encode()[index] else {
return Ok((addr, None, 1));
};

let branch_ptr = build_branch_ptr(merkle, value, chd_encoded)?;
let subproof = generate_subproof(data)?;

Ok((branch_ptr, Some(subproof), 1))
let subproof = generate_subproof(data.to_vec())?;
Ok((addr, Some(subproof), 1))
}

_ => Err(ProofError::DecodeError(Box::new(
bincode::ErrorKind::Custom(String::from("")),
))),
}
}
}
Expand Down Expand Up @@ -625,42 +574,6 @@ fn generate_subproof(data: Vec<u8>) -> Result<SubProof, ProofError> {
}
}

struct CurKey(Vec<u8>);
struct Data(Vec<u8>);

fn get_ext_ptr<S: ShaleStore<Node> + Send + Sync>(
merkle: &Merkle<S>,
term: bool,
Data(data): Data,
CurKey(cur_key): CurKey,
) -> Result<DiskAddress, ProofError> {
let node = if term {
NodeType::Leaf(LeafNode::new(cur_key, data))
} else {
NodeType::Extension(ExtNode::new(cur_key, DiskAddress::null(), Some(data)))
};

merkle
.new_node(Node::new(node))
.map(|node| node.as_ptr())
.map_err(ProofError::InvalidNode)
}

fn build_branch_ptr<S: ShaleStore<Node> + Send + Sync>(
merkle: &Merkle<S>,
value: Option<Vec<u8>>,
chd_encoded: [Option<Vec<u8>>; NBRANCH],
) -> Result<DiskAddress, ProofError> {
let node = BranchNode::new([None; NBRANCH], value, chd_encoded);
let node = NodeType::Branch(node);
let node = Node::new(node);

merkle
.new_node(node)
.map_err(|_| ProofError::ProofNodeMissing)
.map(|node| node.as_ptr())
}

// unset_internal removes all internal node references.
// It should be called after a trie is constructed with two edge paths. Also
// the given boundary keys must be the one used to construct the edge paths.
Expand Down

0 comments on commit 772035e

Please sign in to comment.