Skip to content

Commit

Permalink
[Epoch Sync] Strengthen the storage access guarantees of block merkle…
Browse files Browse the repository at this point in the history 
… proof. (#12308)

This serves as the foundation for solving
#12255

For Epoch Sync, we need to be able to generate block merkle proofs
despite not having older blocks.

Luckily, after much [banging my head against the
wall](https://github.com/user-attachments/assets/74954ab8-a80e-409a-85d1-0a48655274d3)
on the merkle tree proof algorithm mostly without much luck, I
eventually found a way to modify the algorithm (without changing the
result it generates) to guarantee that proving block X against block Y
will never access anything older than X or newer than Y (the latter is
not important and also already true, but we may as well say it neatly
like this).

Added a test to assert that this guarantee is met for all pairs 0 <= X <
Y < 100.

The rigorous proof of this guarantee is left as an exercise to the
reader (good luck :) ).

Update: actually I rewrote the whole algorithm. Not only is there less
code, but also it is now O(log N) lookups as opposed to O(log^2 N) looks
before.
  • Loading branch information
@robin-near
robin-near authored Oct 28, 2024
1 parent 54f53c8 commit 244422d
Show file tree
Hide file tree
Showing 2 changed files with 190 additions and 150 deletions.
322 changes: 172 additions & 150 deletions chain/chain/src/store/merkle_proof.rs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
use std::{collections::HashMap, sync::Arc};
use std::sync::Arc;

use near_chain_primitives::Error;
use near_primitives::{
hash::CryptoHash,
merkle::{combine_hash, Direction, MerklePath, MerklePathItem, PartialMerkleTree},
types::{MerkleHash, NumBlocks},
types::NumBlocks,
utils::index_to_bytes,
};
use near_store::{DBCol, Store};
Expand All @@ -23,6 +23,9 @@ pub trait MerkleProofAccess {
fn get_block_hash_from_ordinal(&self, block_ordinal: NumBlocks) -> Result<CryptoHash, Error>;

/// Get merkle proof for block with hash `block_hash` in the merkle tree of `head_block_hash`.
///
/// Guarantees that no block data (PartialMerkleTree or block hash) for any block older than
/// `block_hash` or newer than `head_block_hash` will be accessed.
fn compute_past_block_proof_in_merkle_tree_of_later_block(
&self,
block_hash: &CryptoHash,
Expand All @@ -40,39 +43,69 @@ pub trait MerkleProofAccess {
block_hash, head_block_hash
)));
}
let mut level = 0;
let mut counter = 1;
let mut cur_index = leaf_index;

let mut path = vec![];
let mut tree_nodes = HashMap::new();
let mut iter = tree_size;
while iter > 1 {
if cur_index % 2 == 0 {
cur_index += 1
} else {
cur_index -= 1;
let mut level: u64 = 0;
let mut index = leaf_index;
let mut remaining_size = tree_size;

while remaining_size > 1 {
// Walk left.
{
let cur_index = index;
let cur_level = level;
// Go up as long as we're the right child. This finds us a largest subtree for
// which our current node is the rightmost descendant of at the current level.
while remaining_size > 1 && index % 2 == 1 {
index /= 2;
remaining_size = (remaining_size + 1) / 2;
level += 1;
}
if level > cur_level {
// To prove this subtree, get the partial tree for the rightmost leaf of the
// subtree. It's OK if we can only go as far as the tree size; we'll aggregate
// whatever we can. Once we have the partial tree, we push in whatever is in
// between the levels we just traversed.
let ordinal = ((cur_index + 1) * (1 << cur_level) - 1).min(tree_size - 1);
let partial_tree_for_node = get_block_merkle_tree_from_ordinal(self, ordinal)?;
partial_tree_for_node.iter_path_from_bottom(|hash, l| {
if l >= cur_level && l < level {
path.push(MerklePathItem { hash, direction: Direction::Left });
}
});
}
}
let direction = if cur_index % 2 == 0 { Direction::Left } else { Direction::Right };
let maybe_hash = if cur_index % 2 == 1 {
// node not immediately available. Needs to be reconstructed
reconstruct_merkle_tree_node(
self,
cur_index,
level,
counter,
tree_size,
&mut tree_nodes,
)?
} else {
get_merkle_tree_node(self, cur_index, level, counter, tree_size, &mut tree_nodes)?
};
if let Some(hash) = maybe_hash {
path.push(MerklePathItem { hash, direction });
// Walk right.
if remaining_size > 1 {
let right_sibling_index = index + 1;
let ordinal = ((right_sibling_index + 1) * (1 << level) - 1).min(tree_size - 1);
// It's possible the right sibling is actually zero, in which case we don't push
// anything to the path.
if ordinal >= right_sibling_index * (1 << level) {
// To prove a right sibling, get the partial tree for the rightmost leaf of the
// subtree, and also get the block hash of the rightmost leaf; combining these
// two will give us the root of the subtree, i.e. the right sibling.
let leaf_hash = self.get_block_hash_from_ordinal(ordinal)?;
let mut subtree_root_hash = leaf_hash;
if level > 0 {
let partial_tree_for_sibling =
get_block_merkle_tree_from_ordinal(self, ordinal)?;
partial_tree_for_sibling.iter_path_from_bottom(|hash, l| {
if l < level {
subtree_root_hash = combine_hash(&hash, &subtree_root_hash);
}
});
}
path.push(MerklePathItem {
hash: subtree_root_hash,
direction: Direction::Right,
});
}

index = (index + 1) / 2;
remaining_size = (remaining_size + 1) / 2;
level += 1;
}
cur_index /= 2;
iter = (iter + 1) / 2;
level += 1;
counter *= 2;
}
Ok(path)
}
Expand All @@ -86,124 +119,6 @@ fn get_block_merkle_tree_from_ordinal(
this.get_block_merkle_tree(&block_hash)
}

/// Get node at given position (index, level). If the node does not exist, return `None`.
fn get_merkle_tree_node(
this: &(impl MerkleProofAccess + ?Sized),
index: u64,
level: u64,
counter: u64,
tree_size: u64,
tree_nodes: &mut HashMap<(u64, u64), Option<MerkleHash>>,
) -> Result<Option<MerkleHash>, Error> {
if let Some(hash) = tree_nodes.get(&(index, level)) {
Ok(*hash)
} else {
if level == 0 {
let maybe_hash = if index >= tree_size {
None
} else {
Some(this.get_block_hash_from_ordinal(index)?)
};
tree_nodes.insert((index, level), maybe_hash);
Ok(maybe_hash)
} else {
let cur_tree_size = (index + 1) * counter;
let maybe_hash = if cur_tree_size > tree_size {
if index * counter <= tree_size {
let left_hash = get_merkle_tree_node(
this,
index * 2,
level - 1,
counter / 2,
tree_size,
tree_nodes,
)?;
let right_hash = reconstruct_merkle_tree_node(
this,
index * 2 + 1,
level - 1,
counter / 2,
tree_size,
tree_nodes,
)?;
combine_maybe_hashes(left_hash, right_hash)
} else {
None
}
} else {
Some(
*get_block_merkle_tree_from_ordinal(this, cur_tree_size)?
.get_path()
.last()
.ok_or_else(|| Error::Other("Merkle tree node missing".to_string()))?,
)
};
tree_nodes.insert((index, level), maybe_hash);
Ok(maybe_hash)
}
}
}

/// Reconstruct node at given position (index, level). If the node does not exist, return `None`.
fn reconstruct_merkle_tree_node(
this: &(impl MerkleProofAccess + ?Sized),
index: u64,
level: u64,
counter: u64,
tree_size: u64,
tree_nodes: &mut HashMap<(u64, u64), Option<MerkleHash>>,
) -> Result<Option<MerkleHash>, Error> {
if let Some(hash) = tree_nodes.get(&(index, level)) {
Ok(*hash)
} else {
if level == 0 {
let maybe_hash = if index >= tree_size {
None
} else {
Some(this.get_block_hash_from_ordinal(index)?)
};
tree_nodes.insert((index, level), maybe_hash);
Ok(maybe_hash)
} else {
let left_hash = get_merkle_tree_node(
this,
index * 2,
level - 1,
counter / 2,
tree_size,
tree_nodes,
)?;
let right_hash = reconstruct_merkle_tree_node(
this,
index * 2 + 1,
level - 1,
counter / 2,
tree_size,
tree_nodes,
)?;
let maybe_hash = combine_maybe_hashes(left_hash, right_hash);
tree_nodes.insert((index, level), maybe_hash);

Ok(maybe_hash)
}
}
}

fn combine_maybe_hashes(
hash1: Option<MerkleHash>,
hash2: Option<MerkleHash>,
) -> Option<MerkleHash> {
match (hash1, hash2) {
(Some(h1), Some(h2)) => Some(combine_hash(&h1, &h2)),
(Some(h1), None) => Some(h1),
(None, Some(_)) => {
debug_assert!(false, "Inconsistent state in merkle proof computation: left node is None but right node exists");
None
}
_ => None,
}
}

impl MerkleProofAccess for Store {
fn get_block_merkle_tree(
&self,
Expand All @@ -226,3 +141,110 @@ impl MerkleProofAccess for Store {
)
}
}

#[cfg(test)]
mod tests {
use super::MerkleProofAccess;
use near_o11y::testonly::init_test_logger;
use near_primitives::hash::{hash, CryptoHash};
use near_primitives::merkle::{verify_hash, MerklePath, PartialMerkleTree};
use near_primitives::types::NumBlocks;
use std::collections::HashMap;
use std::ops::RangeInclusive;
use std::sync::Arc;

struct MerkleProofTestFixture {
block_merkle_trees: HashMap<CryptoHash, PartialMerkleTree>,
block_hashes: Vec<CryptoHash>,
block_merkle_roots: Vec<CryptoHash>,
last_partial_merkle_tree: PartialMerkleTree,
tree_size: NumBlocks,
allowed_access_range: RangeInclusive<NumBlocks>,
}

impl MerkleProofAccess for MerkleProofTestFixture {
fn get_block_merkle_tree(
&self,
block_hash: &CryptoHash,
) -> Result<Arc<PartialMerkleTree>, near_chain_primitives::Error> {
let tree = self.block_merkle_trees.get(block_hash).unwrap().clone();
if !self.allowed_access_range.contains(&tree.size()) {
panic!("Block partial merkle tree for ordinal {} is not available", tree.size());
}
Ok(tree.into())
}

fn get_block_hash_from_ordinal(
&self,
block_ordinal: NumBlocks,
) -> Result<CryptoHash, near_chain_primitives::Error> {
if !self.allowed_access_range.contains(&block_ordinal) {
panic!("Block hash for ordinal {} is not available", block_ordinal);
}
Ok(self.block_hashes[block_ordinal as usize])
}
}

impl MerkleProofTestFixture {
fn new() -> Self {
MerkleProofTestFixture {
block_merkle_trees: HashMap::new(),
block_hashes: vec![],
block_merkle_roots: vec![],
last_partial_merkle_tree: PartialMerkleTree::default(),
tree_size: 0,
allowed_access_range: 0..=NumBlocks::MAX,
}
}

fn append_block(&mut self) {
let hash = hash(&self.tree_size.to_be_bytes());
self.block_hashes.push(hash);
self.block_merkle_roots.push(self.last_partial_merkle_tree.root());
self.block_merkle_trees.insert(hash, self.last_partial_merkle_tree.clone());
self.last_partial_merkle_tree.insert(hash);
self.tree_size += 1;
}

fn append_n_blocks(&mut self, n: u64) {
for _ in 0..n {
self.append_block();
}
}

fn make_proof(&mut self, index: u64, against: u64) -> MerklePath {
self.allowed_access_range = index..=against;
self.compute_past_block_proof_in_merkle_tree_of_later_block(
&self.block_hashes[index as usize],
&self.block_hashes[against as usize],
)
.unwrap()
}

fn verify_proof(&self, index: u64, against: u64, proof: &MerklePath) {
let provee = self.block_hashes[index as usize];
let root = self.block_merkle_roots[against as usize];
assert!(verify_hash(root, proof, provee));
}
}

/// Tests that deriving a merkle proof for block X against merkle root of block Y
/// requires no PartialMerkleTree or BlockHash data for any block earlier than X
/// or later than Y.
///
/// This is useful for Epoch Sync where earlier block data are not available.
#[test]
fn test_no_dependency_on_blocks_outside_range() {
init_test_logger();
let mut f = MerkleProofTestFixture::new();
const MAX_BLOCKS: u64 = 100;
f.append_n_blocks(MAX_BLOCKS + 1);
for i in 0..MAX_BLOCKS {
for j in i + 1..MAX_BLOCKS {
println!("Testing proof of block {} against merkle root at {}", i, j);
let proof = f.make_proof(i, j);
f.verify_proof(i, j, &proof);
}
}
}
}
18 changes: 18 additions & 0 deletions core/primitives/src/merkle.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,24 @@ impl PartialMerkleTree {
pub fn get_path(&self) -> &[MerkleHash] {
&self.path
}

/// Iterate over the path from the bottom to the top, calling `f` with the hash and the level.
/// The level is 0 for the leaf and increases by 1 for each level in the actual tree.
pub fn iter_path_from_bottom(&self, mut f: impl FnMut(MerkleHash, u64)) {
let mut level = 0;
let mut index = self.size;
for node in self.path.iter().rev() {
if index == 0 {
// shouldn't happen
return;
}
let trailing_zeros = index.trailing_zeros();
level += trailing_zeros;
index >>= trailing_zeros;
index -= 1;
f(*node, level as u64);
}
}
}

#[cfg(test)]
Expand Down

0 comments on commit 244422d

Please sign in to comment.