From 4cfdb374ae72eb28525c07e7ee54edda02c10323 Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <93241502+yhchiang-sol@users.noreply.github.com>
Date: Wed, 28 Jun 2023 12:33:29 +0800
Subject: [PATCH] [TieredStorage] byte_block::read_type (#32295)

#### Summary of Changes
This PR adds byte_block::read_type(), an util function that
reads the raw part of the input byte block at the specified offset
as type T.  This function will be later used by the hot storage.

#### Test Plan
Modified existing unit tests to use byte_block::read_type()
when the input byte block is properly aligned.
---
 runtime/src/tiered_storage/byte_block.rs | 53 ++++++++++++++++--------
 1 file changed, 36 insertions(+), 17 deletions(-)

diff --git a/runtime/src/tiered_storage/byte_block.rs b/runtime/src/tiered_storage/byte_block.rs
index dca07afc85d60e..8795aa7b48cf63 100644
--- a/runtime/src/tiered_storage/byte_block.rs
+++ b/runtime/src/tiered_storage/byte_block.rs
@@ -115,6 +115,21 @@ impl ByteBlockWriter {
 /// The util struct for reading byte blocks.
 pub struct ByteBlockReader;
 
+/// Reads the raw part of the input byte_block at the specified offset
+/// as type T.
+///
+/// If `offset` + size_of::<T>() exceeds the size of the input byte_block,
+/// then None will be returned.
+pub fn read_type<T>(byte_block: &[u8], offset: usize) -> Option<&T> {
+    let (next, overflow) = offset.overflowing_add(std::mem::size_of::<T>());
+    if overflow || next > byte_block.len() {
+        return None;
+    }
+    let ptr = byte_block[offset..].as_ptr() as *const T;
+    debug_assert!(ptr as usize % std::mem::align_of::<T>() == 0);
+    Some(unsafe { &*ptr })
+}
+
 impl ByteBlockReader {
     /// Decode the input byte array using the specified format.
     ///
@@ -143,7 +158,7 @@ mod tests {
         solana_sdk::{hash::Hash, stake_history::Epoch},
     };
 
-    fn read_type<T>(buffer: &[u8], offset: usize) -> (T, usize) {
+    fn read_type_unaligned<T>(buffer: &[u8], offset: usize) -> (T, usize) {
         let size = std::mem::size_of::<T>();
         let (next, overflow) = offset.overflowing_add(size);
         assert!(!overflow && next <= buffer.len());
@@ -170,7 +185,7 @@ mod tests {
 
         assert_eq!(decoded_buffer.len(), mem::size_of::<u32>());
 
-        let (value_from_buffer, next) = read_type::<u32>(&decoded_buffer, 0);
+        let (value_from_buffer, next) = read_type_unaligned::<u32>(&decoded_buffer, 0);
         assert_eq!(value, value_from_buffer);
 
         if format != AccountBlockFormat::AlignedRaw {
@@ -250,7 +265,7 @@ mod tests {
         );
 
         // verify meta1 and its data
-        let (meta1_from_buffer, next1) = read_type::<TestMetaStruct>(&decoded_buffer, 0);
+        let (meta1_from_buffer, next1) = read_type_unaligned::<TestMetaStruct>(&decoded_buffer, 0);
         assert_eq!(test_metas[0], meta1_from_buffer);
         assert_eq!(
             test_data1,
@@ -258,8 +273,10 @@ mod tests {
         );
 
         // verify meta2 and its data
-        let (meta2_from_buffer, next2) =
-            read_type::<TestMetaStruct>(&decoded_buffer, next1 + meta1_from_buffer.data_len);
+        let (meta2_from_buffer, next2) = read_type_unaligned::<TestMetaStruct>(
+            &decoded_buffer,
+            next1 + meta1_from_buffer.data_len,
+        );
         assert_eq!(test_metas[1], meta2_from_buffer);
         assert_eq!(
             test_data2,
@@ -267,8 +284,10 @@ mod tests {
         );
 
         // verify meta3 and its data
-        let (meta3_from_buffer, next3) =
-            read_type::<TestMetaStruct>(&decoded_buffer, next2 + meta2_from_buffer.data_len);
+        let (meta3_from_buffer, next3) = read_type_unaligned::<TestMetaStruct>(
+            &decoded_buffer,
+            next2 + meta2_from_buffer.data_len,
+        );
         assert_eq!(test_metas[2], meta3_from_buffer);
         assert_eq!(
             test_data3,
@@ -337,23 +356,23 @@ mod tests {
         let mut offset = 0;
         for opt_fields in &opt_fields_vec {
             if let Some(expected_rent_epoch) = opt_fields.rent_epoch {
-                let (rent_epoch, next) = read_type::<Epoch>(&decoded_buffer, offset);
-                assert_eq!(rent_epoch, expected_rent_epoch);
+                let rent_epoch = read_type::<Epoch>(&decoded_buffer, offset).unwrap();
+                assert_eq!(*rent_epoch, expected_rent_epoch);
                 verified_count += 1;
-                offset = next;
+                offset += std::mem::size_of::<Epoch>();
             }
             if let Some(expected_hash) = opt_fields.account_hash {
-                let (hash, next) = read_type::<Hash>(&decoded_buffer, offset);
-                assert_eq!(hash, expected_hash);
+                let hash = read_type::<Hash>(&decoded_buffer, offset).unwrap();
+                assert_eq!(hash, &expected_hash);
                 verified_count += 1;
-                offset = next;
+                offset += std::mem::size_of::<Hash>();
             }
             if let Some(expected_write_version) = opt_fields.write_version {
-                let (write_version, next) =
-                    read_type::<StoredMetaWriteVersion>(&decoded_buffer, offset);
-                assert_eq!(write_version, expected_write_version);
+                let write_version =
+                    read_type::<StoredMetaWriteVersion>(&decoded_buffer, offset).unwrap();
+                assert_eq!(*write_version, expected_write_version);
                 verified_count += 1;
-                offset = next;
+                offset += std::mem::size_of::<StoredMetaWriteVersion>();
             }
         }