From be1d606dea1bbfe30db413d60ca353f910042b3e Mon Sep 17 00:00:00 2001
From: behzad nouri <behzadnouri@gmail.com>
Date: Wed, 18 May 2022 09:23:04 -0400
Subject: [PATCH] adds sanity checks to Shred::reference_tick_from_data

Shred::reference_tick_from_data should check if payload is indeed a data
shred and has valid size.
---
 ledger/src/blockstore.rs |  6 ++----
 ledger/src/shred.rs      | 18 ++++++++++++------
 ledger/src/shredder.rs   |  4 ++--
 3 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/ledger/src/blockstore.rs b/ledger/src/blockstore.rs
index a60367666..4dca6ba6f 100644
--- a/ledger/src/blockstore.rs
+++ b/ledger/src/blockstore.rs
@@ -1850,10 +1850,8 @@ impl Blockstore {
 
             let upper_index = cmp::min(current_index, end_index);
             // the tick that will be used to figure out the timeout for this hole
-            let reference_tick = u64::from(Shred::reference_tick_from_data(
-                db_iterator.value().expect("couldn't read value"),
-            ));
-
+            let data = db_iterator.value().expect("couldn't read value");
+            let reference_tick = u64::from(Shred::reference_tick_from_data(data).unwrap());
             if ticks_since_first_insert < reference_tick + MAX_TURBINE_DELAY_IN_TICKS {
                 // The higher index holes have not timed out yet
                 break 'outer;
diff --git a/ledger/src/shred.rs b/ledger/src/shred.rs
index 39aa691e9..61b3961f5 100644
--- a/ledger/src/shred.rs
+++ b/ledger/src/shred.rs
@@ -483,10 +483,16 @@ impl Shred {
         limited_deserialize(buffer).ok()
     }
 
-    pub(crate) fn reference_tick_from_data(data: &[u8]) -> u8 {
-        // TODO: should check the type bit as well!
+    pub(crate) fn reference_tick_from_data(data: &[u8]) -> Result<u8, Error> {
         const SHRED_FLAGS_OFFSET: usize = SIZE_OF_COMMON_SHRED_HEADER + std::mem::size_of::<u16>();
-        data[SHRED_FLAGS_OFFSET] & ShredFlags::SHRED_TICK_REFERENCE_MASK.bits()
+        if Self::shred_type_from_payload(data)? != ShredType::Data {
+            return Err(Error::InvalidShredType);
+        }
+        let flags = match data.get(SHRED_FLAGS_OFFSET) {
+            None => return Err(Error::InvalidPayloadSize(data.len())),
+            Some(flags) => flags,
+        };
+        Ok(flags & ShredFlags::SHRED_TICK_REFERENCE_MASK.bits())
     }
 
     pub fn verify(&self, pubkey: &Pubkey) -> bool {
@@ -891,7 +897,7 @@ mod tests {
         assert_eq!(shred, Shred::new_from_serialized_shred(payload).unwrap());
         assert_eq!(
             shred.reference_tick(),
-            Shred::reference_tick_from_data(&packet.data)
+            Shred::reference_tick_from_data(&packet.data).unwrap()
         );
         assert_eq!(Shred::get_slot_from_packet(&packet), Some(shred.slot()));
         assert_eq!(
@@ -932,7 +938,7 @@ mod tests {
         assert_eq!(shred, Shred::new_from_serialized_shred(payload).unwrap());
         assert_eq!(
             shred.reference_tick(),
-            Shred::reference_tick_from_data(&packet.data)
+            Shred::reference_tick_from_data(&packet.data).unwrap()
         );
         assert_eq!(Shred::get_slot_from_packet(&packet), Some(shred.slot()));
         assert_eq!(
@@ -1017,7 +1023,7 @@ mod tests {
             assert_eq!(shred.last_in_slot(), is_last_in_slot);
             assert_eq!(shred.reference_tick(), reference_tick.min(63u8));
             assert_eq!(
-                Shred::reference_tick_from_data(shred.payload()),
+                Shred::reference_tick_from_data(shred.payload()).unwrap(),
                 reference_tick.min(63u8),
             );
         }
diff --git a/ledger/src/shredder.rs b/ledger/src/shredder.rs
index 4d49c1bfe..31dfc7212 100644
--- a/ledger/src/shredder.rs
+++ b/ledger/src/shredder.rs
@@ -519,7 +519,7 @@ mod tests {
         );
         data_shreds.iter().for_each(|s| {
             assert_eq!(s.reference_tick(), 5);
-            assert_eq!(Shred::reference_tick_from_data(s.payload()), 5);
+            assert_eq!(Shred::reference_tick_from_data(s.payload()).unwrap(), 5);
         });
 
         let deserialized_shred =
@@ -555,7 +555,7 @@ mod tests {
                 ShredFlags::SHRED_TICK_REFERENCE_MASK.bits()
             );
             assert_eq!(
-                Shred::reference_tick_from_data(s.payload()),
+                Shred::reference_tick_from_data(s.payload()).unwrap(),
                 ShredFlags::SHRED_TICK_REFERENCE_MASK.bits()
             );
         });