From c4fa78abb5a7bf6eb170c345196141de0b2cb95d Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Thu, 28 Mar 2024 03:15:04 +0000
Subject: [PATCH] v1.18: shred: expose chained merkle root (backport of #435)
 (#459)

shred: expose chained merkle root (#435)

* shred: expose chained merkle root

* pr feedback: macro, pub(super), _=> none

(cherry picked from commit cfd5b71b2896b625dae75189b7f43e60fcd53f07)

Co-authored-by: Ashwin Sekar <ashwin@anza.xyz>
---
 ledger/src/shred.rs            | 32 ++++++++++++++++++++++++++++++++
 ledger/src/shred/merkle.rs     | 24 +++++++++++++++++++++++-
 ledger/src/shred/shred_code.rs |  7 +++++++
 ledger/src/shred/shred_data.rs |  7 +++++++
 4 files changed, 69 insertions(+), 1 deletion(-)

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<SignedData, Error>);
 
+    dispatch!(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 +728,36 @@ 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(),
+            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<R: Rng>(
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<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(super) 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
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<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
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<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),