shred: expose chained merkle root

anza-xyz · Mar 26, 2024 · 2578849 · 2578849
1 parent bcf3d80
commit 2578849
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 1 deletion.
diff --git a/ledger/src/shred.rs b/ledger/src/shred.rs
@@ -334,6 +334,16 @@ macro_rules! dispatch {
                 Self::ShredData(shred) => shred.$name($($arg, )?),
             }
         }
+    };
+    (allow_dead $vis:vis fn $name:ident(&self $(, $arg:ident : $ty:ty)?) $(-> $out:ty)?) => {
+        #[allow(dead_code)]
+        #[inline]
+        $vis fn $name(&self $(, $arg:$ty)?) $(-> $out)? {
+            match self {
+                Self::ShredCode(shred) => shred.$name($($arg, )?),
+                Self::ShredData(shred) => shred.$name($($arg, )?),
+            }
+        }
     }
 }
 
@@ -344,6 +354,7 @@ impl Shred {
     dispatch!(fn set_signature(&mut self, signature: Signature));
     dispatch!(fn signed_data(&self) -> Result<SignedData, Error>);
 
+    dispatch!(allow_dead pub(crate) fn chained_merkle_root(&self) -> Result<Hash, Error>);
     // Returns the portion of the shred's payload which is erasure coded.
     dispatch!(pub(crate) fn erasure_shard(self) -> Result<Vec<u8>, Error>);
     // Like Shred::erasure_shard but returning a slice.
@@ -726,6 +737,27 @@ pub mod layout {
         }
     }
 
+    #[allow(dead_code)]
+    pub(crate) fn get_chained_merkle_root(shred: &[u8]) -> Option<Hash> {
+        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(),
+            _ => 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<R: Rng>(

diff --git 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(crate) fn get_chained_merkle_root_offset(
+        proof_size: u8,
+        resigned: bool,
+    ) -> Result<usize, Error> {
         Ok(Self::SIZE_OF_HEADERS + Self::capacity(proof_size, /*chained:*/ true, resigned)?)
     }
 
+    pub(super) fn chained_merkle_root(&self) -> Result<Hash, Error> {
+        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(crate) fn get_chained_merkle_root_offset(
+        proof_size: u8,
+        resigned: bool,
+    ) -> Result<usize, Error> {
         Ok(Self::SIZE_OF_HEADERS + Self::capacity(proof_size, /*chained:*/ true, resigned)?)
     }
 
-    fn chained_merkle_root(&self) -> Result<Hash, Error> {
+    pub(super) fn chained_merkle_root(&self) -> Result<Hash, Error> {
         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
@@ -47,6 +47,13 @@ impl ShredCode {
         }
     }
 
+    pub(super) fn chained_merkle_root(&self) -> Result<Hash, Error> {
+        match self {
+            Self::Legacy(_) => Err(Error::InvalidShredType),
+            Self::Merkle(shred) => shred.chained_merkle_root(),
+        }
+    }
+
     pub(super) fn merkle_root(&self) -> Result<Hash, Error> {
         match self {
             Self::Legacy(_) => Err(Error::InvalidShredType),

diff --git 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<Hash, Error> {
+        match self {
+            Self::Legacy(_) => Err(Error::InvalidShredType),
+            Self::Merkle(shred) => shred.chained_merkle_root(),
+        }
+    }
+
     pub(super) fn merkle_root(&self) -> Result<Hash, Error> {
         match self {
             Self::Legacy(_) => Err(Error::InvalidShredType),