Fix replay stage test (#6406)

Author: pgarg66

core/src/blocktree.rs

@@ -429,7 +429,7 @@ impl Blocktree {
                     &mut write_batch,
                     &mut just_inserted_data_shreds,
                 );
-            } else {
+            } else if shred.is_code() {
                 self.check_insert_coding_shred(
                     shred,
                     &mut erasure_metas,
@@ -1033,7 +1033,14 @@ impl Blocktree {
                 if is_complete {
                     if let Ok(deshred_payload) = Shredder::deshred(&shred_chunk) {
                         debug!("{:?} shreds in last FEC set", shred_chunk.len(),);
-                        let entries: Vec<Entry> = bincode::deserialize(&deshred_payload)?;
+                        let entries: Vec<Entry> =
+                            bincode::deserialize(&deshred_payload).map_err(|_| {
+                                Error::BlocktreeError(BlocktreeError::InvalidShredData(Box::new(
+                                    bincode::ErrorKind::Custom(
+                                        "could not construct entries".to_string(),
+                                    ),
+                                )))
+                            })?;
                         return Ok((entries, shred_chunk.len()));
                     } else {
                         debug!("Failed in deshredding shred payloads");

core/src/replay_stage.rs

@@ -878,14 +878,15 @@ impl Service for ReplayStage {
 mod test {
     use super::*;
     use crate::blocktree::tests::make_slot_entries;
-    use crate::blocktree::{entries_to_test_shreds, get_tmp_ledger_path};
+    use crate::blocktree::{entries_to_test_shreds, get_tmp_ledger_path, BlocktreeError};
     use crate::confidence::BankConfidence;
     use crate::entry;
     use crate::genesis_utils::{create_genesis_block, create_genesis_block_with_leader};
     use crate::replay_stage::ReplayStage;
-    use crate::shred::Shred;
+    use crate::shred::{Shred, ShredHeader, DATA_COMPLETE_SHRED, SIZE_OF_SHRED_HEADER};
     use solana_runtime::genesis_utils::GenesisBlockInfo;
     use solana_sdk::hash::{hash, Hash};
+    use solana_sdk::packet::PACKET_DATA_SIZE;
     use solana_sdk::signature::{Keypair, KeypairUtil};
     use solana_sdk::system_transaction;
     use solana_sdk::transaction::TransactionError;
@@ -1004,28 +1005,24 @@ mod test {
     }
 
     #[test]
-    fn test_dead_fork_blob_deserialize_failure() {
-        let keypair1 = Keypair::new();
-        let keypair2 = Keypair::new();
-        // Insert entry that causes blob deserialization failure
-
-        let res = check_dead_fork(|blockhash, slot| {
-            let entry = entry::next_entry(
-                &blockhash,
-                1,
-                vec![system_transaction::create_user_account(
-                    &keypair1,
-                    &keypair2.pubkey(),
-                    2,
-                    *blockhash,
-                )],
+    fn test_dead_fork_entry_deserialize_failure() {
+        // Insert entry that causes deserialization failure
+        let res = check_dead_fork(|_, _| {
+            let payload_len = PACKET_DATA_SIZE - *SIZE_OF_SHRED_HEADER;
+            let gibberish = [0xa5u8; PACKET_DATA_SIZE];
+            let mut header = ShredHeader::default();
+            header.data_header.flags = DATA_COMPLETE_SHRED;
+            let mut shred = Shred::new_empty_from_header(header);
+            let _ = bincode::serialize_into(
+                &mut shred.payload[*SIZE_OF_SHRED_HEADER..],
+                &gibberish[..payload_len],
             );
-            entries_to_test_shreds(vec![entry], slot, slot.saturating_sub(1), false)
+            vec![shred]
         });
 
         assert_matches!(
             res,
-            Err(Error::TransactionError(TransactionError::AccountNotFound))
+            Err(Error::BlocktreeError(BlocktreeError::InvalidShredData(_)))
         );
     }
 

core/src/shred.rs

@@ -52,10 +52,10 @@ pub const MAX_DATA_SHREDS_PER_FEC_BLOCK: u32 = 16;
 pub const RECOMMENDED_FEC_RATE: f32 = 0.25;
 
 const LAST_SHRED_IN_SLOT: u8 = 0b0000_0001;
-const DATA_COMPLETE_SHRED: u8 = 0b0000_0010;
+pub const DATA_COMPLETE_SHRED: u8 = 0b0000_0010;
 
 #[derive(Serialize, Clone, Deserialize, PartialEq, Debug)]
-pub struct ShredType(u8);
+pub struct ShredType(pub u8);
 
 /// A common header that is present in data and code shred headers
 #[derive(Serialize, Clone, Deserialize, Default, PartialEq, Debug)]
@@ -245,6 +245,9 @@ impl Shred {
     pub fn is_data(&self) -> bool {
         self.headers.shred_type == ShredType(DATA_SHRED)
     }
+    pub fn is_code(&self) -> bool {
+        self.headers.shred_type == ShredType(CODING_SHRED)
+    }
 
     pub fn last_in_slot(&self) -> bool {
         if self.is_data() {
@@ -271,7 +274,7 @@ impl Shred {
     }
 
     pub fn coding_params(&self) -> Option<(u16, u16, u16)> {
-        if !self.is_data() {
+        if self.is_code() {
             let header = &self.headers.coding_header;
             Some((
                 header.num_data_shreds,
@@ -286,8 +289,10 @@ impl Shred {
     pub fn verify(&self, pubkey: &Pubkey) -> bool {
         let signed_payload_offset = if self.is_data() {
             *SIZE_OF_CODING_SHRED_HEADER + *SIZE_OF_SHRED_TYPE
-        } else {
+        } else if self.is_code() {
             *SIZE_OF_SHRED_TYPE
+        } else {
+            return false;
         } + *SIZE_OF_SIGNATURE;
         self.signature()
             .verify(pubkey.as_ref(), &self.payload[signed_payload_offset..])