From 20e800230feeee8704a080e151a2871d5bf99d8c Mon Sep 17 00:00:00 2001
From: Pankaj Garg <pankaj@solana.com>
Date: Tue, 15 Oct 2019 15:18:23 -0700
Subject: [PATCH] Don't deserialize coding header for data shreds (#6367)

* Don't deserialize coding hdr for data shreds

* review comments

* fix tests
---
 core/benches/shredder.rs | 14 +++++++++++++-
 core/src/shred.rs        | 36 ++++++++++++++++++++++++++++++++++--
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/core/benches/shredder.rs b/core/benches/shredder.rs
index de9cac6c05..4f18e15f3f 100644
--- a/core/benches/shredder.rs
+++ b/core/benches/shredder.rs
@@ -5,7 +5,7 @@ extern crate test;
 use solana_core::entry::create_ticks;
 use solana_core::entry::Entry;
 use solana_core::shred::{
-    max_entries_per_n_shred, max_ticks_per_n_shreds, Shredder, RECOMMENDED_FEC_RATE,
+    max_entries_per_n_shred, max_ticks_per_n_shreds, Shred, Shredder, RECOMMENDED_FEC_RATE,
     SIZE_OF_DATA_SHRED_HEADER,
 };
 use solana_core::test_tx;
@@ -72,3 +72,15 @@ fn bench_deshredder(bencher: &mut Bencher) {
         assert_ne!(raw.len(), 0);
     })
 }
+
+#[bench]
+fn bench_deserialize_hdr(bencher: &mut Bencher) {
+    let data = vec![0; PACKET_DATA_SIZE - *SIZE_OF_DATA_SHRED_HEADER];
+
+    let shred = Shred::new_from_data(2, 1, 1, Some(&data), true, true);
+
+    bencher.iter(|| {
+        let payload = shred.payload.clone();
+        let _ = Shred::new_from_serialized_shred(payload).unwrap();
+    })
+}
diff --git a/core/src/shred.rs b/core/src/shred.rs
index 39085ab236..53d484fd93 100644
--- a/core/src/shred.rs
+++ b/core/src/shred.rs
@@ -1,4 +1,5 @@
 //! The `shred` module defines data structures and methods to pull MTU sized data frames from the network.
+use crate::blocktree::BlocktreeError;
 use crate::entry::create_ticks;
 use crate::entry::Entry;
 use crate::erasure::Session;
@@ -26,9 +27,14 @@ lazy_static! {
         { serialized_size(&CodingShredHeader::default()).unwrap() as usize };
     pub static ref SIZE_OF_DATA_SHRED_HEADER: usize =
         { serialized_size(&DataShredHeader::default()).unwrap() as usize };
+    pub static ref SIZE_OF_COMMON_SHRED_HEADER: usize =
+        { serialized_size(&ShredCommonHeader::default()).unwrap() as usize };
     static ref SIZE_OF_SIGNATURE: usize =
         { bincode::serialized_size(&Signature::default()).unwrap() as usize };
     pub static ref SIZE_OF_SHRED_TYPE: usize = { bincode::serialized_size(&0u8).unwrap() as usize };
+    pub static ref SIZE_OF_PARENT_OFFSET: usize =
+        { bincode::serialized_size(&0u16).unwrap() as usize };
+    pub static ref SIZE_OF_FLAGS: usize = { bincode::serialized_size(&0u8).unwrap() as usize };
 }
 
 thread_local!(static PAR_THREAD_POOL: RefCell<ThreadPool> = RefCell::new(rayon::ThreadPoolBuilder::new()
@@ -150,6 +156,14 @@ impl Shred {
         Self::new(header, shred_buf)
     }
 
+    fn deserialize_obj<'de, T>(index: &mut usize, size: usize, buf: &'de [u8]) -> result::Result<T>
+    where
+        T: Deserialize<'de>,
+    {
+        let ret = bincode::deserialize(&buf[*index..*index + size])?;
+        *index += size;
+        Ok(ret)
+    }
     pub fn new_from_serialized_shred(shred_buf: Vec<u8>) -> result::Result<Self> {
         let shred_type: u8 = bincode::deserialize(&shred_buf[..*SIZE_OF_SHRED_TYPE])?;
         let header = if shred_type == CODING_SHRED {
@@ -157,9 +171,27 @@ impl Shred {
             let mut header = DataShredHeader::default();
             header.common_header = bincode::deserialize(&shred_buf[..end])?;
             header
+        } else if shred_type == DATA_SHRED {
+            let mut start = *SIZE_OF_CODING_SHRED_HEADER;
+            let common_hdr: ShredCommonHeader =
+                Self::deserialize_obj(&mut start, *SIZE_OF_COMMON_SHRED_HEADER, &shred_buf)?;
+
+            let parent_offset: u16 =
+                Self::deserialize_obj(&mut start, *SIZE_OF_PARENT_OFFSET, &shred_buf)?;
+
+            let flags: u8 = Self::deserialize_obj(&mut start, *SIZE_OF_FLAGS, &shred_buf)?;
+            let mut hdr = DataShredHeader {
+                common_header: CodingShredHeader::default(),
+                data_header: common_hdr,
+                parent_offset,
+                flags,
+            };
+            hdr.common_header.shred_type = shred_type;
+            hdr
         } else {
-            let end = *SIZE_OF_DATA_SHRED_HEADER;
-            bincode::deserialize(&shred_buf[..end])?
+            return Err(Error::BlocktreeError(BlocktreeError::InvalidShredData(
+                Box::new(bincode::ErrorKind::Custom("Invalid shred type".to_string())),
+            )));
         };
 
         Ok(Self::new(header, shred_buf))