Skip to content

Commit

Permalink
generalizes the return type of Shred::get_signed_data (solana-labs#29446
Browse files Browse the repository at this point in the history 
)

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.
  • Loading branch information
@behzadnouri
behzadnouri authored Dec 31, 2022
1 parent 70c9017 commit 754ecf4
Show file tree
Hide file tree
Showing 7 changed files with 80 additions and 52 deletions.
4 changes: 2 additions & 2 deletions core/src/broadcast_stage/broadcast_duplicates_run.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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!(
Expand All @@ -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
Expand Down
49 changes: 26 additions & 23 deletions ledger/src/shred.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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>);
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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 => {
Expand Down
16 changes: 10 additions & 6 deletions ledger/src/shred/legacy.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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..])
}
}

Expand Down
31 changes: 19 additions & 12 deletions ledger/src/shred/merkle.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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()
}
}

Expand Down Expand Up @@ -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(()));
Expand Down
12 changes: 9 additions & 3 deletions ledger/src/shred/shred_code.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
Expand Down Expand Up @@ -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,
Expand Down
10 changes: 8 additions & 2 deletions ledger/src/shred/shred_data.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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,
Expand Down
10 changes: 6 additions & 4 deletions ledger/src/shred/traits.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>;
Expand All @@ -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> {
Expand All @@ -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> {
Expand Down

0 comments on commit 754ecf4

Please sign in to comment.