1
0
mirror of https://github.com/logos-co/nomos-pocs.git synced 2025-01-12 18:34:10 +00:00

fix off-by-one in sparse merkle tree

This commit is contained in:
David Rusu 2024-12-06 11:36:06 +04:00
parent bd1f928e00
commit 0056486a6c
3 changed files with 147 additions and 37 deletions
emmarin/cl/cl/src

@ -25,7 +25,7 @@ pub struct NullifierCommitment([u8; 32]);
// The nullifier attached to input notes to prove an input has not
// already been spent.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct Nullifier([u8; 32]);
pub struct Nullifier(pub [u8; 32]);
impl NullifierSecret {
pub fn random(mut rng: impl RngCore) -> Self {

@ -1,7 +1,13 @@
use crate::cl::{merkle, mmr::MMR, Nullifier};
use std::collections::BTreeSet;
use crate::cl::{
merkle,
mmr::{MMRProof, MMR},
NoteCommitment, Nullifier,
};
use serde::{Deserialize, Serialize};
const MAX_NULL: usize = 256;
use super::sparse_merkle_tree;
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct Ledger {
@ -12,23 +18,59 @@ pub struct Ledger {
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct LedgerWitness {
pub commitments: MMR,
pub nullifiers: Vec<Nullifier>,
pub nf_root: [u8; 32],
}
impl LedgerWitness {
pub fn commit(&self) -> Ledger {
Ledger {
cm_root: self.commitments.commit(),
nf_root: self.nf_root,
}
}
pub fn assert_nf_update(&mut self, nf: Nullifier, path: &[merkle::PathNode]) {
// verify that the path corresponds to the nullifier
assert_eq!(sparse_merkle_tree::path_key(path), nf.0);
// verify that the nullifier was not already present
assert_eq!(
merkle::path_root(sparse_merkle_tree::ABSENT, path),
self.nf_root
);
// update the nullifer root with the nullifier inserted into the tree
self.nf_root = merkle::path_root(sparse_merkle_tree::PRESENT, path);
}
}
pub struct LedgerState {
commitments: MMR,
nullifiers: BTreeSet<[u8; 32]>,
}
impl LedgerState {
pub fn to_witness(&self) -> LedgerWitness {
LedgerWitness {
commitments: self.commitments.clone(),
nf_root: self.nf_root(),
}
}
pub fn nf_root(&self) -> [u8; 32] {
let bytes = self
.nullifiers
.iter()
.map(|i| i.as_bytes().to_vec())
.collect::<Vec<_>>();
merkle::root(merkle::padded_leaves::<MAX_NULL>(&bytes))
sparse_merkle_tree::sparse_root(&self.nullifiers)
}
pub fn add_commitment(&mut self, cm: NoteCommitment) -> MMRProof {
self.commitments.push(&cm.0)
}
pub fn add_nullifier(&mut self, nf: Nullifier) -> Vec<merkle::PathNode> {
let path = sparse_merkle_tree::sparse_path(nf.0, &self.nullifiers);
assert!(!self.nullifiers.contains(&nf.0));
self.nullifiers.insert(nf.0);
path
}
}

@ -4,17 +4,17 @@ use crate::cl::merkle;
use lazy_static::lazy_static;
/// absence of element is marked with all 0's
static ABSENT: [u8; 32] = [0u8; 32];
pub static ABSENT: [u8; 32] = [0u8; 32];
/// presence of element is marked with all 1's
static PRESENT: [u8; 32] = [255u8; 32];
pub static PRESENT: [u8; 32] = [255u8; 32];
lazy_static! {
// the roots of empty merkle trees of diffent heights
// i.e. all leafs are ABSENT
static ref EMPTY_ROOTS: [[u8; 32]; 256] = {
let mut roots = [ABSENT; 256];
for h in 1..256 {
static ref EMPTY_ROOTS: [[u8; 32]; 257] = {
let mut roots = [ABSENT; 257];
for h in 1..257 {
roots[h] = merkle::node(roots[h - 1], roots[h - 1]);
}
@ -26,16 +26,17 @@ pub fn sparse_root(elems: &BTreeSet<[u8; 32]>) -> [u8; 32] {
sparse_root_rec(0, elems)
}
fn sparse_root_rec(prefix: u8, elems: &BTreeSet<[u8; 32]>) -> [u8; 32] {
fn sparse_root_rec(prefix: u64, elems: &BTreeSet<[u8; 32]>) -> [u8; 32] {
if elems.is_empty() {
return empty_tree_root(255 - prefix);
return empty_tree_root(256 - prefix);
}
if prefix == 255 {
if prefix == 256 {
assert_eq!(elems.len(), 1);
return PRESENT;
}
// partition the elements
let (left, right): (BTreeSet<_>, BTreeSet<_>) = elems.iter().partition(|e| !bit(prefix, **e));
let (left, right): (BTreeSet<_>, BTreeSet<_>) =
elems.iter().partition(|e| !bit(prefix as u8, **e));
merkle::node(
sparse_root_rec(prefix + 1, &left),
@ -45,18 +46,18 @@ fn sparse_root_rec(prefix: u8, elems: &BTreeSet<[u8; 32]>) -> [u8; 32] {
pub fn sparse_path(elem: [u8; 32], elems: &BTreeSet<[u8; 32]>) -> Vec<merkle::PathNode> {
fn sparse_path_rec(
prefix: u8,
prefix: u64,
elem: [u8; 32],
elems: &BTreeSet<[u8; 32]>,
) -> Vec<merkle::PathNode> {
if prefix == 255 {
if prefix == 256 {
return Vec::new();
}
// partition the elements
let (left, right): (BTreeSet<_>, BTreeSet<_>) =
elems.iter().partition(|e| !bit(prefix, **e));
elems.iter().partition(|e| !bit(prefix as u8, **e));
match bit(prefix, elem) {
match bit(prefix as u8, elem) {
true => {
let left_root = sparse_root_rec(prefix + 1, &left);
let mut path = sparse_path_rec(prefix + 1, elem, &right);
@ -77,7 +78,27 @@ pub fn sparse_path(elem: [u8; 32], elems: &BTreeSet<[u8; 32]>) -> Vec<merkle::Pa
sparse_path_rec(0, elem, elems)
}
fn empty_tree_root(height: u8) -> [u8; 32] {
pub fn path_key(path: &[merkle::PathNode]) -> [u8; 32] {
assert_eq!(path.len(), 256);
let mut key = [0u8; 32];
for byte_i in (0..32).rev() {
let mut byte = 0u8;
for bit_i in 0..8 {
byte <<= 1;
match path[byte_i * 8 + bit_i] {
merkle::PathNode::Left(_) => byte += 1,
merkle::PathNode::Right(_) => byte += 0,
};
}
key[31 - byte_i] = byte;
}
key
}
fn empty_tree_root(height: u64) -> [u8; 32] {
assert!(height <= 256);
EMPTY_ROOTS[height as usize]
}
@ -85,7 +106,7 @@ fn bit(idx: u8, elem: [u8; 32]) -> bool {
let byte = idx / 8;
let bit_in_byte = idx - byte * 8;
(elem[byte as usize] & (1 << bit_in_byte)) >> bit_in_byte == 1
(elem[byte as usize] & (1 << bit_in_byte)) != 0
}
#[cfg(test)]
@ -96,6 +117,55 @@ mod tests {
rand::random()
}
#[test]
fn test_neighbour_paths() {
let elems = BTreeSet::from_iter([[0u8; 32]]);
let path_0 = sparse_path([0u8; 32], &elems);
let mut key_1 = [0u8; 32];
key_1[31] = 128;
let path_1 = sparse_path(key_1, &elems);
assert_ne!(path_0, path_1);
}
#[test]
fn test_path_bit_agreement() {
fn path_bit(idx: u8, path: &[merkle::PathNode]) -> bool {
match path[255 - idx as usize] {
merkle::PathNode::Left(_) => true,
merkle::PathNode::Right(_) => false,
}
}
let key = random_hash();
let path = sparse_path(key, &BTreeSet::new());
for i in 0..=255 {
let b = bit(i, key);
let pb = path_bit(i, &path);
assert_eq!(b, pb, "{}!={}@{}", b, pb, i);
}
}
#[test]
fn test_path_key() {
let elems = BTreeSet::from_iter(std::iter::repeat_with(random_hash).take(10));
// membership proofs
for e in elems.iter() {
let path = sparse_path(*e, &elems);
assert_eq!(path_key(&path), *e);
}
// non-membership proofs
for _ in 0..10 {
let elem = random_hash();
let path = sparse_path(elem, &elems);
assert_eq!(path_key(&path), elem);
}
}
#[test]
fn test_sparse_path() {
let elems = BTreeSet::from_iter(std::iter::repeat_with(random_hash).take(10));
@ -105,7 +175,6 @@ mod tests {
// membership proofs
for e in elems.iter() {
let path = sparse_path(*e, &elems);
assert_eq!(path.len(), 255);
assert_eq!(merkle::path_root(PRESENT, &path), root);
}
@ -114,7 +183,6 @@ mod tests {
let elem = random_hash();
let path = sparse_path(elem, &elems);
assert!(!elems.contains(&elem));
assert_eq!(path.len(), 255);
assert_eq!(merkle::path_root(ABSENT, &path), root);
}
}
@ -130,7 +198,7 @@ mod tests {
for (h, node) in path.into_iter().enumerate() {
match node {
merkle::PathNode::Left(hash) | merkle::PathNode::Right(hash) => {
assert_eq!(hash, empty_tree_root(h as u8))
assert_eq!(hash, empty_tree_root(h as u64))
}
}
}
@ -147,7 +215,7 @@ mod tests {
// / \ 0 subtree
// 1 0
let mut expected_root = PRESENT;
for h in 0..=254 {
for h in 0..=255 {
expected_root = merkle::node(expected_root, empty_tree_root(h))
}
@ -165,7 +233,7 @@ mod tests {
// 0 /\
// 0 1
let mut expected_root = PRESENT;
for h in 0..=254 {
for h in 0..=255 {
expected_root = merkle::node(empty_tree_root(h), expected_root)
}
@ -196,10 +264,10 @@ mod tests {
// \
// 1
let mut expected_root = PRESENT;
for h in 0..=253 {
for h in 0..=254 {
expected_root = merkle::node(empty_tree_root(h), expected_root)
}
expected_root = merkle::node(expected_root, empty_tree_root(254));
expected_root = merkle::node(expected_root, empty_tree_root(255));
assert_eq!(root, expected_root)
}
@ -223,11 +291,11 @@ mod tests {
// 0 1
let mut expected_root = PRESENT;
for h in 0..=254 {
for h in 0..=255 {
if h % 2 == 0 {
expected_root = merkle::node(empty_tree_root(h), expected_root)
} else {
expected_root = merkle::node(expected_root, empty_tree_root(h))
} else {
expected_root = merkle::node(empty_tree_root(h), expected_root)
}
}
assert_eq!(root, expected_root)
@ -245,12 +313,12 @@ mod tests {
// 1 0 0 1
let mut left_root = PRESENT;
for h in 0..=253 {
for h in 0..=254 {
left_root = merkle::node(left_root, empty_tree_root(h))
}
let mut right_root = PRESENT;
for h in 0..=253 {
for h in 0..=254 {
right_root = merkle::node(empty_tree_root(h), right_root)
}
let expected_root = merkle::node(left_root, right_root);