Skip to content

Commit

Permalink
implements shred::layout::get_merkle_root (backport #29437) (#29847)
Browse files Browse the repository at this point in the history 
implements shred::layout::get_merkle_root (#29437)

In preparation of removing merkle roots form shreds binary, the commit
adds api to recover the root from the merkle proof embedded in shreds
payload.

(cherry picked from commit db3d926)

Co-authored-by: behzad nouri <[email protected]>
  • Loading branch information
@mergify @behzadnouri
mergify[bot] and behzadnouri authored Jan 24, 2023
1 parent 1ef6b1f commit 1343d2b
Show file tree
Hide file tree
Showing 2 changed files with 117 additions and 21 deletions.
28 changes: 20 additions & 8 deletions ledger/src/shred.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,6 @@ const SIZE_OF_CODING_SHRED_HEADERS: usize = 89;
const SIZE_OF_SIGNATURE: usize = SIGNATURE_BYTES;
const SIZE_OF_SHRED_VARIANT: usize = 1;
const SIZE_OF_SHRED_SLOT: usize = 8;
const SIZE_OF_SHRED_INDEX: usize = 4;

const OFFSET_OF_SHRED_VARIANT: usize = SIZE_OF_SIGNATURE;
const OFFSET_OF_SHRED_SLOT: usize = SIZE_OF_SIGNATURE + SIZE_OF_SHRED_VARIANT;
Expand Down Expand Up @@ -539,6 +538,8 @@ impl Shred {
// Helper methods to extract pieces of the shred from the payload
// without deserializing the entire payload.
pub mod layout {
#[cfg(test)]
use crate::shred::merkle::MerkleRoot;
use {super::*, std::ops::Range};

fn get_shred_size(packet: &Packet) -> Option<usize> {
Expand Down Expand Up @@ -592,17 +593,15 @@ pub mod layout {
}

pub fn get_version(shred: &[u8]) -> Option<u16> {
const OFFSET_OF_SHRED_VERSION: usize = OFFSET_OF_SHRED_INDEX + SIZE_OF_SHRED_INDEX;
<[u8; 2]>::try_from(shred.get(OFFSET_OF_SHRED_VERSION..)?.get(..2)?)
<[u8; 2]>::try_from(shred.get(77..79)?)
.map(u16::from_le_bytes)
.ok()
}

// The caller should verify first that the shred is data and not code!
pub(super) fn get_parent_offset(shred: &[u8]) -> Option<u16> {
const OFFSET_OF_SHRED_PARENT: usize = SIZE_OF_COMMON_SHRED_HEADER;
debug_assert_eq!(get_shred_type(shred).unwrap(), ShredType::Data);
<[u8; 2]>::try_from(shred.get(OFFSET_OF_SHRED_PARENT..)?.get(..2)?)
<[u8; 2]>::try_from(shred.get(83..85)?)
.map(u16::from_le_bytes)
.ok()
}
Expand Down Expand Up @@ -631,17 +630,28 @@ pub mod layout {
}

pub(crate) fn get_reference_tick(shred: &[u8]) -> Result<u8, Error> {
const SIZE_OF_PARENT_OFFSET: usize = std::mem::size_of::<u16>();
const OFFSET_OF_SHRED_FLAGS: usize = SIZE_OF_COMMON_SHRED_HEADER + SIZE_OF_PARENT_OFFSET;
if get_shred_type(shred)? != ShredType::Data {
return Err(Error::InvalidShredType);
}
let flags = match shred.get(OFFSET_OF_SHRED_FLAGS) {
let flags = match shred.get(85) {
None => return Err(Error::InvalidPayloadSize(shred.len())),
Some(flags) => flags,
};
Ok(flags & ShredFlags::SHRED_TICK_REFERENCE_MASK.bits())
}

#[cfg(test)]
pub(crate) fn get_merkle_root(shred: &[u8]) -> Option<MerkleRoot> {
match get_shred_variant(shred).ok()? {
ShredVariant::LegacyCode | ShredVariant::LegacyData => None,
ShredVariant::MerkleCode(proof_size) => {
merkle::ShredCode::get_merkle_root(shred, proof_size)
}
ShredVariant::MerkleData(proof_size) => {
merkle::ShredData::get_merkle_root(shred, proof_size)
}
}
}
}

impl From<ShredCode> for Shred {
Expand Down Expand Up @@ -948,6 +958,8 @@ mod tests {
solana_sdk::{shred_version, signature::Signer},
};

const SIZE_OF_SHRED_INDEX: usize = 4;

fn bs58_decode<T: AsRef<[u8]>>(data: T) -> Vec<u8> {
bs58::decode(data).into_vec().unwrap()
}
Expand Down
110 changes: 97 additions & 13 deletions ledger/src/shred/merkle.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ const_assert_eq!(ShredData::SIZE_OF_PAYLOAD, 1203);
const MERKLE_HASH_PREFIX_LEAF: &[u8] = &[0x00];
const MERKLE_HASH_PREFIX_NODE: &[u8] = &[0x01];

type MerkleRoot = MerkleProofEntry;
pub(crate) type MerkleRoot = MerkleProofEntry;
type MerkleProofEntry = [u8; 20];

// Layout: {common, data} headers | data buffer | merkle branch
Expand Down Expand Up @@ -246,6 +246,27 @@ impl ShredData {
}
shred_data::sanitize(self)
}

#[cfg(test)]
pub(super) fn get_merkle_root(shred: &[u8], proof_size: u8) -> Option<MerkleRoot> {
debug_assert_eq!(
shred::layout::get_shred_variant(shred).unwrap(),
ShredVariant::MerkleData(proof_size)
);
// Shred index in the erasure batch.
let index = {
let fec_set_index = <[u8; 4]>::try_from(shred.get(79..83)?)
.map(u32::from_le_bytes)
.ok()?;
shred::layout::get_index(shred)?
.checked_sub(fec_set_index)
.map(usize::try_from)?
.ok()?
};
// Where the merkle branch starts in the shred binary.
let offset = Self::SIZE_OF_HEADERS + Self::capacity(proof_size).ok()?;
get_merkle_root(shred, proof_size, index, offset)
}
}

impl ShredCode {
Expand Down Expand Up @@ -375,6 +396,29 @@ impl ShredCode {
}
shred_code::sanitize(self)
}

#[cfg(test)]
pub(super) fn get_merkle_root(shred: &[u8], proof_size: u8) -> Option<MerkleRoot> {
debug_assert_eq!(
shred::layout::get_shred_variant(shred).unwrap(),
ShredVariant::MerkleCode(proof_size)
);
// Shred index in the erasure batch.
let index = {
let num_data_shreds = <[u8; 2]>::try_from(shred.get(83..85)?)
.map(u16::from_le_bytes)
.map(usize::from)
.ok()?;
let position = <[u8; 2]>::try_from(shred.get(87..89)?)
.map(u16::from_le_bytes)
.map(usize::from)
.ok()?;
num_data_shreds.checked_add(position)?
};
// Where the merkle branch starts in the shred binary.
let offset = Self::SIZE_OF_HEADERS + Self::capacity(proof_size).ok()?;
get_merkle_root(shred, proof_size, index, offset)
}
}

impl ShredTrait for ShredData {
Expand Down Expand Up @@ -569,17 +613,51 @@ fn join_nodes<S: AsRef<[u8]>, T: AsRef<[u8]>>(node: S, other: T) -> Hash {
}

fn verify_merkle_proof(index: usize, node: Hash, merkle_branch: &MerkleBranch) -> bool {
let proof = merkle_branch.proof.iter();
let (index, root) = proof.fold((index, node), |(index, node), other| {
let parent = if index % 2 == 0 {
join_nodes(node, other)
} else {
join_nodes(other, node)
};
(index >> 1, parent)
});
let root = &root.as_ref()[..SIZE_OF_MERKLE_ROOT];
(index, root) == (0usize, &merkle_branch.root[..])
let proof = merkle_branch.proof.iter().copied();
let root = fold_merkle_proof(index, node, proof);
root.as_ref() == Some(merkle_branch.root)
}

// Recovers root of the merkle tree from a leaf node
// at the given index and the respective proof.
fn fold_merkle_proof<'a, I>(index: usize, node: Hash, proof: I) -> Option<MerkleRoot>
where
I: IntoIterator<Item = &'a MerkleProofEntry>,
{
let (index, root) = proof
.into_iter()
.fold((index, node), |(index, node), other| {
let parent = if index % 2 == 0 {
join_nodes(node, other)
} else {
join_nodes(other, node)
};
(index >> 1, parent)
});
(index == 0).then(|| {
let root = &root.as_ref()[..SIZE_OF_MERKLE_ROOT];
MerkleRoot::try_from(root).ok()
})?
}

#[cfg(test)]
fn get_merkle_root(
shred: &[u8],
proof_size: u8,
index: usize, // Shred index in the erasure batch.
offset: usize, // Where the merkle branch starts in the shred binary.
) -> Option<MerkleRoot> {
let node = shred.get(SIZE_OF_SIGNATURE..offset)?;
let node = hashv(&[MERKLE_HASH_PREFIX_LEAF, node]);
// Merkle proof embedded in the payload.
let offset = offset + SIZE_OF_MERKLE_ROOT;
let size = usize::from(proof_size) * SIZE_OF_MERKLE_PROOF_ENTRY;
let proof = shred
.get(offset..offset + size)?
.chunks(SIZE_OF_MERKLE_PROOF_ENTRY)
.map(<&MerkleProofEntry>::try_from)
.map(Result::unwrap);
fold_merkle_proof(index, node, proof)
}

fn make_merkle_tree(mut nodes: Vec<Hash>) -> Vec<Hash> {
Expand Down Expand Up @@ -1059,7 +1137,7 @@ fn make_erasure_batch(
mod test {
use {
super::*,
crate::shred::ShredFlags,
crate::shred::{ShredFlags, ShredId},
itertools::Itertools,
matches::assert_matches,
rand::{seq::SliceRandom, CryptoRng, Rng},
Expand Down Expand Up @@ -1412,15 +1490,21 @@ mod test {
version,
fec_set_index: _,
} = *shred.common_header();
let shred_type = ShredType::from(shred_variant);
let key = ShredId::new(slot, index, shred_type);
let merkle_root = shred.merkle_root().copied().ok();
let shred = shred.payload();
assert_eq!(shred::layout::get_signature(shred), Some(signature));
assert_eq!(
shred::layout::get_shred_variant(shred).unwrap(),
shred_variant
);
assert_eq!(shred::layout::get_shred_type(shred).unwrap(), shred_type);
assert_eq!(shred::layout::get_slot(shred), Some(slot));
assert_eq!(shred::layout::get_index(shred), Some(index));
assert_eq!(shred::layout::get_version(shred), Some(version));
assert_eq!(shred::layout::get_shred_id(shred), Some(key));
assert_eq!(shred::layout::get_merkle_root(shred), merkle_root);
let slice = shred::layout::get_signed_message_range(shred).unwrap();
let message = shred.get(slice).unwrap();
assert!(signature.verify(keypair.pubkey().as_ref(), message));
Expand Down

0 comments on commit 1343d2b

Please sign in to comment.