generalizes the return type of Shred::get_signed_data (solana-labs#29446

)

The commit adds an associated SignedData type to Shred trait so that
merkle and legacy shreds can return different types for signed_data
method.
This would allow legacy shreds to point to a section of the shred
payload, whereas merkle shreds would compute and return the merkle root.
Ultimately this would allow to remove the merkle root from the shreds
binary.

core/src/broadcast_stage/broadcast_duplicates_run.rs

@@ -312,7 +312,7 @@ impl BroadcastRun for BroadcastDuplicatesRun {
                     .original_last_data_shreds
                     .lock()
                     .unwrap()
-                    .remove(&shred.signature())
+                    .remove(shred.signature())
                 {
                     if cluster_partition.contains(&node.id) {
                         info!(
@@ -327,7 +327,7 @@ impl BroadcastRun for BroadcastDuplicatesRun {
                     .partition_last_data_shreds
                     .lock()
                     .unwrap()
-                    .remove(&shred.signature())
+                    .remove(shred.signature())
                 {
                     // If the shred is part of the partition, broadcast it directly to the
                     // partition node. This is to account for cases when the partition stake

ledger/src/shred.rs

@@ -52,7 +52,7 @@
 #[cfg(test)]
 pub(crate) use shred_code::MAX_CODE_SHREDS_PER_SLOT;
 use {
-    self::{shred_code::ShredCode, traits::Shred as _},
+    self::{merkle::MerkleRoot, shred_code::ShredCode, traits::Shred as _},
     crate::blockstore::{self, MAX_DATA_SHREDS_PER_SLOT},
     bitflags::bitflags,
     num_enum::{IntoPrimitive, TryFromPrimitive},
@@ -230,6 +230,20 @@ pub enum Shred {
     ShredData(ShredData),
 }
 
+pub(crate) enum SignedData<'a> {
+    Chunk(&'a [u8]), // Chunk of payload past signature.
+    MerkleRoot(MerkleRoot),
+}
+
+impl<'a> AsRef<[u8]> for SignedData<'a> {
+    fn as_ref(&self) -> &[u8] {
+        match self {
+            Self::Chunk(chunk) => chunk,
+            Self::MerkleRoot(root) => root,
+        }
+    }
+}
+
 /// Tuple which uniquely identifies a shred should it exists.
 #[derive(Clone, Copy, Eq, Debug, Hash, PartialEq)]
 pub struct ShredId(Slot, /*shred index:*/ u32, ShredType);
@@ -309,7 +323,7 @@ use dispatch;
 impl Shred {
     dispatch!(fn common_header(&self) -> &ShredCommonHeader);
     dispatch!(fn set_signature(&mut self, signature: Signature));
-    dispatch!(fn signed_data(&self) -> &[u8]);
+    dispatch!(fn signed_data(&self) -> Result<SignedData, Error>);
 
     // Returns the portion of the shred's payload which is erasure coded.
     dispatch!(pub(crate) fn erasure_shard(self) -> Result<Vec<u8>, Error>);
@@ -455,12 +469,13 @@ impl Shred {
         ErasureSetId(self.slot(), self.fec_set_index())
     }
 
-    pub fn signature(&self) -> Signature {
-        self.common_header().signature
+    pub fn signature(&self) -> &Signature {
+        &self.common_header().signature
     }
 
     pub fn sign(&mut self, keypair: &Keypair) {
-        let signature = keypair.sign_message(self.signed_data());
+        let data = self.signed_data().unwrap();
+        let signature = keypair.sign_message(data.as_ref());
         self.set_signature(signature);
     }
 
@@ -508,8 +523,10 @@ impl Shred {
 
     #[must_use]
     pub fn verify(&self, pubkey: &Pubkey) -> bool {
-        let data = self.signed_data();
-        self.signature().verify(pubkey.as_ref(), data)
+        match self.signed_data() {
+            Ok(data) => self.signature().verify(pubkey.as_ref(), data.as_ref()),
+            Err(_) => false,
+        }
     }
 
     // Returns true if the erasure coding of the two shreds mismatch.
@@ -538,27 +555,13 @@ impl Shred {
 // Helper methods to extract pieces of the shred from the payload
 // without deserializing the entire payload.
 pub mod layout {
-    use {super::*, crate::shred::merkle::MerkleRoot, std::ops::Range};
+    use {super::*, std::ops::Range};
     #[cfg(test)]
     use {
         rand::{seq::SliceRandom, Rng},
         std::collections::HashMap,
     };
 
-    pub(crate) enum SignedData<'a> {
-        Chunk(&'a [u8]),
-        MerkleRoot(MerkleRoot),
-    }
-
-    impl<'a> AsRef<[u8]> for SignedData<'a> {
-        fn as_ref(&self) -> &[u8] {
-            match self {
-                Self::Chunk(chunk) => chunk,
-                Self::MerkleRoot(root) => root,
-            }
-        }
-    }
-
     fn get_shred_size(packet: &Packet) -> Option<usize> {
         let size = packet.data(..)?.len();
         if packet.meta().repair() {
@@ -1467,7 +1470,7 @@ mod tests {
         assert_eq!(layout::get_index(data), Some(shred.index()));
         assert_eq!(layout::get_version(data), Some(shred.version()));
         assert_eq!(layout::get_shred_id(data), Some(shred.id()));
-        assert_eq!(layout::get_signature(data), Some(shred.signature()));
+        assert_eq!(layout::get_signature(data), Some(*shred.signature()));
         assert_eq!(layout::get_shred_type(data).unwrap(), shred.shred_type());
         match shred.shred_type() {
             ShredType::Code => {

ledger/src/shred/legacy.rs

@@ -48,7 +48,9 @@ pub struct ShredCode {
     payload: Vec<u8>,
 }
 
-impl Shred for ShredData {
+impl<'a> Shred<'a> for ShredData {
+    type SignedData = &'a [u8];
+
     impl_shred_common!();
     // Legacy data shreds are always zero padded and
     // the same size as coding shreds.
@@ -109,13 +111,15 @@ impl Shred for ShredData {
         shred_data::sanitize(self)
     }
 
-    fn signed_data(&self) -> &[u8] {
+    fn signed_data(&'a self) -> Result<Self::SignedData, Error> {
         debug_assert_eq!(self.payload.len(), Self::SIZE_OF_PAYLOAD);
-        &self.payload[SIZE_OF_SIGNATURE..]
+        Ok(&self.payload[SIZE_OF_SIGNATURE..])
     }
 }
 
-impl Shred for ShredCode {
+impl<'a> Shred<'a> for ShredCode {
+    type SignedData = &'a [u8];
+
     impl_shred_common!();
     const SIZE_OF_PAYLOAD: usize = shred_code::ShredCode::SIZE_OF_PAYLOAD;
     const SIZE_OF_HEADERS: usize = SIZE_OF_CODING_SHRED_HEADERS;
@@ -171,9 +175,9 @@ impl Shred for ShredCode {
         shred_code::sanitize(self)
     }
 
-    fn signed_data(&self) -> &[u8] {
+    fn signed_data(&'a self) -> Result<Self::SignedData, Error> {
         debug_assert_eq!(self.payload.len(), Self::SIZE_OF_PAYLOAD);
-        &self.payload[SIZE_OF_SIGNATURE..]
+        Ok(&self.payload[SIZE_OF_SIGNATURE..])
     }
 }
 

ledger/src/shred/merkle.rs

@@ -91,16 +91,18 @@ impl Shred {
     dispatch!(fn sanitize(&self, verify_merkle_proof: bool) -> Result<(), Error>);
     dispatch!(fn set_merkle_branch(&mut self, merkle_branch: &MerkleBranch) -> Result<(), Error>);
     dispatch!(fn set_signature(&mut self, signature: Signature));
-    dispatch!(fn signed_data(&self) -> &[u8]);
+    dispatch!(fn signed_data(&self) -> Result<MerkleRoot, Error>);
 
     #[must_use]
     fn verify(&self, pubkey: &Pubkey) -> bool {
-        let data = self.signed_data();
-        self.signature().verify(pubkey.as_ref(), data)
+        match self.signed_data() {
+            Ok(data) => self.signature().verify(pubkey.as_ref(), data.as_ref()),
+            Err(_) => false,
+        }
     }
 
-    fn signature(&self) -> Signature {
-        self.common_header().signature
+    fn signature(&self) -> &Signature {
+        &self.common_header().signature
     }
 
     fn from_payload(shred: Vec<u8>) -> Result<Self, Error> {
@@ -419,7 +421,9 @@ impl ShredCode {
     }
 }
 
-impl ShredTrait for ShredData {
+impl<'a> ShredTrait<'a> for ShredData {
+    type SignedData = MerkleRoot;
+
     impl_shred_common!();
 
     // Also equal to:
@@ -486,12 +490,14 @@ impl ShredTrait for ShredData {
         self.sanitize(/*verify_merkle_proof:*/ true)
     }
 
-    fn signed_data(&self) -> &[u8] {
-        self.merkle_root().map(AsRef::as_ref).unwrap_or_default()
+    fn signed_data(&'a self) -> Result<Self::SignedData, Error> {
+        self.merkle_root().copied()
     }
 }
 
-impl ShredTrait for ShredCode {
+impl<'a> ShredTrait<'a> for ShredCode {
+    type SignedData = MerkleRoot;
+
     impl_shred_common!();
     const SIZE_OF_PAYLOAD: usize = shred_code::ShredCode::SIZE_OF_PAYLOAD;
     const SIZE_OF_HEADERS: usize = SIZE_OF_CODING_SHRED_HEADERS;
@@ -551,8 +557,8 @@ impl ShredTrait for ShredCode {
         self.sanitize(/*verify_merkle_proof:*/ true)
     }
 
-    fn signed_data(&self) -> &[u8] {
-        self.merkle_root().map(AsRef::as_ref).unwrap_or_default()
+    fn signed_data(&'a self) -> Result<Self::SignedData, Error> {
+        self.merkle_root().copied()
     }
 }
 
@@ -1344,7 +1350,8 @@ mod test {
             let merkle_branch = make_merkle_branch(index, num_shreds, &tree).unwrap();
             assert_eq!(merkle_branch.proof.len(), usize::from(proof_size));
             shred.set_merkle_branch(&merkle_branch).unwrap();
-            let signature = keypair.sign_message(shred.signed_data());
+            let data = shred.signed_data().unwrap();
+            let signature = keypair.sign_message(data.as_ref());
             shred.set_signature(signature);
             assert!(shred.verify(&keypair.pubkey()));
             assert_matches!(shred.sanitize(/*verify_merkle_proof:*/ true), Ok(()));

ledger/src/shred/shred_code.rs

@@ -4,8 +4,8 @@ use {
             common::dispatch,
             legacy, merkle,
             traits::{Shred, ShredCode as ShredCodeTrait},
-            CodingShredHeader, Error, ShredCommonHeader, ShredType, DATA_SHREDS_PER_FEC_BLOCK,
-            MAX_DATA_SHREDS_PER_SLOT, SIZE_OF_NONCE,
+            CodingShredHeader, Error, ShredCommonHeader, ShredType, SignedData,
+            DATA_SHREDS_PER_FEC_BLOCK, MAX_DATA_SHREDS_PER_SLOT, SIZE_OF_NONCE,
         },
         shredder::ERASURE_BATCH_SIZE,
     },
@@ -39,12 +39,18 @@ impl ShredCode {
     dispatch!(pub(super) fn payload(&self) -> &Vec<u8>);
     dispatch!(pub(super) fn sanitize(&self) -> Result<(), Error>);
     dispatch!(pub(super) fn set_signature(&mut self, signature: Signature));
-    dispatch!(pub(super) fn signed_data(&self) -> &[u8]);
 
     // Only for tests.
     dispatch!(pub(super) fn set_index(&mut self, index: u32));
     dispatch!(pub(super) fn set_slot(&mut self, slot: Slot));
 
+    pub(super) fn signed_data(&self) -> Result<SignedData, Error> {
+        match self {
+            Self::Legacy(shred) => Ok(SignedData::Chunk(shred.signed_data()?)),
+            Self::Merkle(shred) => Ok(SignedData::MerkleRoot(shred.signed_data()?)),
+        }
+    }
+
     pub(super) fn new_from_parity_shard(
         slot: Slot,
         index: u32,

ledger/src/shred/shred_data.rs

@@ -4,7 +4,7 @@ use {
         common::dispatch,
         legacy, merkle,
         traits::{Shred as _, ShredData as ShredDataTrait},
-        DataShredHeader, Error, ShredCommonHeader, ShredFlags, ShredType, ShredVariant,
+        DataShredHeader, Error, ShredCommonHeader, ShredFlags, ShredType, ShredVariant, SignedData,
         MAX_DATA_SHREDS_PER_SLOT,
     },
     solana_sdk::{clock::Slot, signature::Signature},
@@ -29,12 +29,18 @@ impl ShredData {
     dispatch!(pub(super) fn payload(&self) -> &Vec<u8>);
     dispatch!(pub(super) fn sanitize(&self) -> Result<(), Error>);
     dispatch!(pub(super) fn set_signature(&mut self, signature: Signature));
-    dispatch!(pub(super) fn signed_data(&self) -> &[u8]);
 
     // Only for tests.
     dispatch!(pub(super) fn set_index(&mut self, index: u32));
     dispatch!(pub(super) fn set_slot(&mut self, slot: Slot));
 
+    pub(super) fn signed_data(&self) -> Result<SignedData, Error> {
+        match self {
+            Self::Legacy(shred) => Ok(SignedData::Chunk(shred.signed_data()?)),
+            Self::Merkle(shred) => Ok(SignedData::MerkleRoot(shred.signed_data()?)),
+        }
+    }
+
     pub(super) fn new_from_data(
         slot: Slot,
         index: u32,

ledger/src/shred/traits.rs

@@ -3,13 +3,15 @@ use {
     solana_sdk::{clock::Slot, signature::Signature},
 };
 
-pub(super) trait Shred: Sized {
+pub(super) trait Shred<'a>: Sized {
     // Total size of payload including headers, merkle
     // branches (if any), zero paddings, etc.
     const SIZE_OF_PAYLOAD: usize;
     // Size of common and code/data headers.
     const SIZE_OF_HEADERS: usize;
 
+    type SignedData: AsRef<[u8]>;
+
     fn from_payload(shred: Vec<u8>) -> Result<Self, Error>;
     fn common_header(&self) -> &ShredCommonHeader;
     fn sanitize(&self) -> Result<(), Error>;
@@ -27,14 +29,14 @@ pub(super) trait Shred: Sized {
     fn erasure_shard_as_slice(&self) -> Result<&[u8], Error>;
 
     // Portion of the payload which is signed.
-    fn signed_data(&self) -> &[u8];
+    fn signed_data(&'a self) -> Result<Self::SignedData, Error>;
 
     // Only for tests.
     fn set_index(&mut self, index: u32);
     fn set_slot(&mut self, slot: Slot);
 }
 
-pub(super) trait ShredData: Shred {
+pub(super) trait ShredData: for<'a> Shred<'a> {
     fn data_header(&self) -> &DataShredHeader;
 
     fn parent(&self) -> Result<Slot, Error> {
@@ -56,7 +58,7 @@ pub(super) trait ShredData: Shred {
     fn data(&self) -> Result<&[u8], Error>;
 }
 
-pub(super) trait ShredCode: Shred {
+pub(super) trait ShredCode: for<'a> Shred<'a> {
     fn coding_header(&self) -> &CodingShredHeader;
 
     fn first_coding_index(&self) -> Option<u32> {