From 4a1f6e0260b608d1a8d0706b4259eecab89828fa Mon Sep 17 00:00:00 2001 From: behzad nouri Date: Sat, 31 Dec 2022 17:08:25 +0000 Subject: [PATCH] generalizes the return type of Shred::get_signed_data (#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. --- .../broadcast_duplicates_run.rs | 4 +- ledger/src/shred.rs | 49 ++++++++++--------- ledger/src/shred/legacy.rs | 16 +++--- ledger/src/shred/merkle.rs | 31 +++++++----- ledger/src/shred/shred_code.rs | 12 +++-- ledger/src/shred/shred_data.rs | 10 +++- ledger/src/shred/traits.rs | 10 ++-- 7 files changed, 80 insertions(+), 52 deletions(-) diff --git a/core/src/broadcast_stage/broadcast_duplicates_run.rs b/core/src/broadcast_stage/broadcast_duplicates_run.rs index c617ede200e459..0c414c38eb99fb 100644 --- a/core/src/broadcast_stage/broadcast_duplicates_run.rs +++ b/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 diff --git a/ledger/src/shred.rs b/ledger/src/shred.rs index 07460273d98e81..22da633adbc27e 100644 --- a/ledger/src/shred.rs +++ b/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); // Returns the portion of the shred's payload which is erasure coded. dispatch!(pub(crate) fn erasure_shard(self) -> Result, 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 { 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 => { diff --git a/ledger/src/shred/legacy.rs b/ledger/src/shred/legacy.rs index 76b8eca36cd735..2a6b72b339822b 100644 --- a/ledger/src/shred/legacy.rs +++ b/ledger/src/shred/legacy.rs @@ -48,7 +48,9 @@ pub struct ShredCode { payload: Vec, } -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 { 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 { debug_assert_eq!(self.payload.len(), Self::SIZE_OF_PAYLOAD); - &self.payload[SIZE_OF_SIGNATURE..] + Ok(&self.payload[SIZE_OF_SIGNATURE..]) } } diff --git a/ledger/src/shred/merkle.rs b/ledger/src/shred/merkle.rs index 8a7cdef6534e34..0c6746398a6b3e 100644 --- a/ledger/src/shred/merkle.rs +++ b/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); #[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) -> Result { @@ -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.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.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(())); diff --git a/ledger/src/shred/shred_code.rs b/ledger/src/shred/shred_code.rs index fd7715d130f40d..165ab80d310f51 100644 --- a/ledger/src/shred/shred_code.rs +++ b/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); 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 { + 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, diff --git a/ledger/src/shred/shred_data.rs b/ledger/src/shred/shred_data.rs index 7e1e1e2bb0baaa..915ae29d4e3f2e 100644 --- a/ledger/src/shred/shred_data.rs +++ b/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); 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 { + 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, diff --git a/ledger/src/shred/traits.rs b/ledger/src/shred/traits.rs index 3cb9b01403e1ff..35a6c0a617af64 100644 --- a/ledger/src/shred/traits.rs +++ b/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) -> Result; 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; // 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 { @@ -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 {