Skip to content

Commit

Permalink
adds get_proof_offset for Merkle shreds (solana-labs#34798)
Browse files Browse the repository at this point in the history 
In preparation of adding chained Merkle shreds variant, the commit
reworks api for proof-offset within the shred binary.
  • Loading branch information
@behzadnouri
behzadnouri authored Jan 17, 2024
1 parent d878899 commit 586c794
Showing 1 changed file with 30 additions and 22 deletions.
52 changes: 30 additions & 22 deletions ledger/src/shred/merkle.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,28 +150,34 @@ impl ShredData {
}

// Where the merkle proof starts in the shred binary.
fn proof_offset(proof_size: u8) -> Result<usize, Error> {
fn proof_offset(&self) -> Result<usize, Error> {
let ShredVariant::MerkleData(proof_size) = self.common_header.shred_variant else {
return Err(Error::InvalidShredVariant);
};
Self::get_proof_offset(proof_size)
}

fn get_proof_offset(proof_size: u8) -> Result<usize, Error> {
Ok(Self::SIZE_OF_HEADERS + Self::capacity(proof_size)?)
}

pub(super) fn merkle_root(&self) -> Result<Hash, Error> {
let proof_size = self.proof_size()?;
let index = self.erasure_shard_index()?;
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
let proof = get_merkle_proof(&self.payload, proof_offset, proof_size)?;
let node = get_merkle_node(&self.payload, SIZE_OF_SIGNATURE..proof_offset)?;
get_merkle_root(index, node, proof)
}

fn merkle_proof(&self) -> Result<impl Iterator<Item = &MerkleProofEntry>, Error> {
let proof_size = self.proof_size()?;
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
get_merkle_proof(&self.payload, proof_offset, proof_size)
}

fn merkle_node(&self) -> Result<Hash, Error> {
let proof_size = self.proof_size()?;
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
get_merkle_node(&self.payload, SIZE_OF_SIGNATURE..proof_offset)
}

Expand Down Expand Up @@ -207,7 +213,7 @@ impl ShredData {
if proof.len() != usize::from(proof_size) {
return Err(Error::InvalidMerkleProof);
}
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
let mut cursor = Cursor::new(
self.payload
.get_mut(proof_offset..)
Expand All @@ -234,7 +240,7 @@ impl ShredData {
.map(usize::try_from)?
.ok()?
};
let proof_offset = Self::proof_offset(proof_size).ok()?;
let proof_offset = Self::get_proof_offset(proof_size).ok()?;
let proof = get_merkle_proof(shred, proof_offset, proof_size).ok()?;
let node = get_merkle_node(shred, SIZE_OF_SIGNATURE..proof_offset).ok()?;
get_merkle_root(index, node, proof).ok()
Expand Down Expand Up @@ -262,28 +268,34 @@ impl ShredCode {
}

// Where the merkle proof starts in the shred binary.
fn proof_offset(proof_size: u8) -> Result<usize, Error> {
fn proof_offset(&self) -> Result<usize, Error> {
let ShredVariant::MerkleCode(proof_size) = self.common_header.shred_variant else {
return Err(Error::InvalidShredVariant);
};
Self::get_proof_offset(proof_size)
}

fn get_proof_offset(proof_size: u8) -> Result<usize, Error> {
Ok(Self::SIZE_OF_HEADERS + Self::capacity(proof_size)?)
}

pub(super) fn merkle_root(&self) -> Result<Hash, Error> {
let proof_size = self.proof_size()?;
let index = self.erasure_shard_index()?;
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
let proof = get_merkle_proof(&self.payload, proof_offset, proof_size)?;
let node = get_merkle_node(&self.payload, SIZE_OF_SIGNATURE..proof_offset)?;
get_merkle_root(index, node, proof)
}

fn merkle_proof(&self) -> Result<impl Iterator<Item = &MerkleProofEntry>, Error> {
let proof_size = self.proof_size()?;
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
get_merkle_proof(&self.payload, proof_offset, proof_size)
}

fn merkle_node(&self) -> Result<Hash, Error> {
let proof_size = self.proof_size()?;
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
get_merkle_node(&self.payload, SIZE_OF_SIGNATURE..proof_offset)
}

Expand Down Expand Up @@ -321,7 +333,7 @@ impl ShredCode {
if proof.len() != usize::from(proof_size) {
return Err(Error::InvalidMerkleProof);
}
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
let mut cursor = Cursor::new(
self.payload
.get_mut(proof_offset..)
Expand Down Expand Up @@ -350,7 +362,7 @@ impl ShredCode {
.ok()?;
num_data_shreds.checked_add(position)?
};
let proof_offset = Self::proof_offset(proof_size).ok()?;
let proof_offset = Self::get_proof_offset(proof_size).ok()?;
let proof = get_merkle_proof(shred, proof_offset, proof_size).ok()?;
let node = get_merkle_node(shred, SIZE_OF_SIGNATURE..proof_offset).ok()?;
get_merkle_root(index, node, proof).ok()
Expand Down Expand Up @@ -402,8 +414,7 @@ impl<'a> ShredTrait<'a> for ShredData {
if self.payload.len() != Self::SIZE_OF_PAYLOAD {
return Err(Error::InvalidPayloadSize(self.payload.len()));
}
let proof_size = self.proof_size()?;
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
let mut shard = self.payload;
shard.truncate(proof_offset);
shard.drain(0..SIZE_OF_SIGNATURE);
Expand All @@ -414,8 +425,7 @@ impl<'a> ShredTrait<'a> for ShredData {
if self.payload.len() != Self::SIZE_OF_PAYLOAD {
return Err(Error::InvalidPayloadSize(self.payload.len()));
}
let proof_size = self.proof_size()?;
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
self.payload
.get(SIZE_OF_SIGNATURE..proof_offset)
.ok_or(Error::InvalidPayloadSize(self.payload.len()))
Expand Down Expand Up @@ -474,8 +484,7 @@ impl<'a> ShredTrait<'a> for ShredCode {
if self.payload.len() != Self::SIZE_OF_PAYLOAD {
return Err(Error::InvalidPayloadSize(self.payload.len()));
}
let proof_size = self.proof_size()?;
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
let mut shard = self.payload;
shard.truncate(proof_offset);
shard.drain(..Self::SIZE_OF_HEADERS);
Expand All @@ -486,8 +495,7 @@ impl<'a> ShredTrait<'a> for ShredCode {
if self.payload.len() != Self::SIZE_OF_PAYLOAD {
return Err(Error::InvalidPayloadSize(self.payload.len()));
}
let proof_size = self.proof_size()?;
let proof_offset = Self::proof_offset(proof_size)?;
let proof_offset = self.proof_offset()?;
self.payload
.get(Self::SIZE_OF_HEADERS..proof_offset)
.ok_or(Error::InvalidPayloadSize(self.payload.len()))
Expand Down

0 comments on commit 586c794

Please sign in to comment.