diff --git a/ledger/src/shred.rs b/ledger/src/shred.rs index 24d5000b65311b..2b6f6f136784c2 100644 --- a/ledger/src/shred.rs +++ b/ledger/src/shred.rs @@ -310,6 +310,7 @@ impl ErasureSetId { macro_rules! dispatch { ($vis:vis fn $name:ident(&self $(, $arg:ident : $ty:ty)?) $(-> $out:ty)?) => { #[inline] + #[allow(dead_code)] $vis fn $name(&self $(, $arg:$ty)?) $(-> $out)? { match self { Self::ShredCode(shred) => shred.$name($($arg, )?), @@ -344,6 +345,7 @@ impl Shred { dispatch!(fn set_signature(&mut self, signature: Signature)); dispatch!(fn signed_data(&self) -> Result); + dispatch!(pub(crate) fn chained_merkle_root(&self) -> Result); // Returns the portion of the shred's payload which is erasure coded. dispatch!(pub(crate) fn erasure_shard(self) -> Result, Error>); // Like Shred::erasure_shard but returning a slice. @@ -726,6 +728,36 @@ pub mod layout { } } + #[allow(dead_code)] + pub(crate) fn get_chained_merkle_root(shred: &[u8]) -> Option { + let offset = match get_shred_variant(shred).ok()? { + ShredVariant::LegacyCode | ShredVariant::LegacyData => None, + ShredVariant::MerkleCode { + proof_size, + chained: true, + resigned, + } => merkle::ShredCode::get_chained_merkle_root_offset(proof_size, resigned).ok(), + ShredVariant::MerkleData { + proof_size, + chained: true, + resigned, + } => merkle::ShredData::get_chained_merkle_root_offset(proof_size, resigned).ok(), + ShredVariant::MerkleCode { + proof_size: _, + chained: false, + resigned: _, + } => None, + ShredVariant::MerkleData { + proof_size: _, + chained: false, + resigned: _, + } => None, + }?; + shred + .get(offset..offset + SIZE_OF_MERKLE_ROOT) + .map(Hash::new) + } + // Minimally corrupts the packet so that the signature no longer verifies. #[cfg(test)] pub(crate) fn corrupt_packet( diff --git a/ledger/src/shred/merkle.rs b/ledger/src/shred/merkle.rs index b785eeb6dc32cc..a7cc134824cf31 100644 --- a/ledger/src/shred/merkle.rs +++ b/ledger/src/shred/merkle.rs @@ -190,9 +190,24 @@ impl ShredData { else { return Err(Error::InvalidShredVariant); }; + Self::get_chained_merkle_root_offset(proof_size, resigned) + } + + pub(super) fn get_chained_merkle_root_offset( + proof_size: u8, + resigned: bool, + ) -> Result { Ok(Self::SIZE_OF_HEADERS + Self::capacity(proof_size, /*chained:*/ true, resigned)?) } + pub(super) fn chained_merkle_root(&self) -> Result { + let offset = self.chained_merkle_root_offset()?; + self.payload + .get(offset..offset + SIZE_OF_MERKLE_ROOT) + .map(Hash::new) + .ok_or(Error::InvalidPayloadSize(self.payload.len())) + } + fn set_chained_merkle_root(&mut self, chained_merkle_root: &Hash) -> Result<(), Error> { let offset = self.chained_merkle_root_offset()?; let Some(buffer) = self.payload.get_mut(offset..offset + SIZE_OF_MERKLE_ROOT) else { @@ -361,10 +376,17 @@ impl ShredCode { else { return Err(Error::InvalidShredVariant); }; + Self::get_chained_merkle_root_offset(proof_size, resigned) + } + + pub(super) fn get_chained_merkle_root_offset( + proof_size: u8, + resigned: bool, + ) -> Result { Ok(Self::SIZE_OF_HEADERS + Self::capacity(proof_size, /*chained:*/ true, resigned)?) } - fn chained_merkle_root(&self) -> Result { + pub(super) fn chained_merkle_root(&self) -> Result { let offset = self.chained_merkle_root_offset()?; self.payload .get(offset..offset + SIZE_OF_MERKLE_ROOT) diff --git a/ledger/src/shred/shred_code.rs b/ledger/src/shred/shred_code.rs index 0ad97a0f729a77..067d7edaf437eb 100644 --- a/ledger/src/shred/shred_code.rs +++ b/ledger/src/shred/shred_code.rs @@ -47,6 +47,13 @@ impl ShredCode { } } + pub(super) fn chained_merkle_root(&self) -> Result { + match self { + Self::Legacy(_) => Err(Error::InvalidShredType), + Self::Merkle(shred) => shred.chained_merkle_root(), + } + } + pub(super) fn merkle_root(&self) -> Result { match self { Self::Legacy(_) => Err(Error::InvalidShredType), diff --git a/ledger/src/shred/shred_data.rs b/ledger/src/shred/shred_data.rs index 15f407172cfc4b..ac409376370420 100644 --- a/ledger/src/shred/shred_data.rs +++ b/ledger/src/shred/shred_data.rs @@ -41,6 +41,13 @@ impl ShredData { } } + pub(super) fn chained_merkle_root(&self) -> Result { + match self { + Self::Legacy(_) => Err(Error::InvalidShredType), + Self::Merkle(shred) => shred.chained_merkle_root(), + } + } + pub(super) fn merkle_root(&self) -> Result { match self { Self::Legacy(_) => Err(Error::InvalidShredType),