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

implements shred::layout::get_merkle_root (backport #29437) #29847

Merged
merged 1 commit into from
Jan 24, 2023
Merged

implements shred::layout::get_merkle_root (backport #29437) #29847

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 20 additions & 8 deletions ledger/src/shred.rs
View file
Open in desktop
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
Open in desktop
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