TransactionView: Address Table Lookup Iterator (#2639)

transaction-view/src/address_table_lookup_meta.rs

@@ -2,11 +2,12 @@ use {
     crate::{
         bytes::{
             advance_offset_for_array, advance_offset_for_type, check_remaining,
-            optimized_read_compressed_u16, read_byte,
+            optimized_read_compressed_u16, read_byte, read_slice_data, read_type,
         },
         result::Result,
     },
     solana_sdk::{hash::Hash, packet::PACKET_DATA_SIZE, pubkey::Pubkey, signature::Signature},
+    solana_svm_transaction::message_address_table_lookup::SVMMessageAddressTableLookup,
 };
 
 // Each ATL has at least a Pubkey, one byte for the number of write indexes,
@@ -46,7 +47,7 @@ const MAX_ATLS_PER_PACKET: usize = (PACKET_DATA_SIZE - MIN_SIZED_PACKET_WITH_ATL
 /// Contains metadata about the address table lookups in a transaction packet.
 pub struct AddressTableLookupMeta {
     /// The number of address table lookups in the transaction.
-    pub(crate) num_address_table_lookup: u8,
+    pub(crate) num_address_table_lookups: u8,
     /// The offset to the first address table lookup in the transaction.
     pub(crate) offset: u16,
 }
@@ -97,12 +98,84 @@ impl AddressTableLookupMeta {
         }
 
         Ok(Self {
-            num_address_table_lookup: num_address_table_lookups,
+            num_address_table_lookups,
             offset: address_table_lookups_offset,
         })
     }
 }
 
+pub struct AddressTableLookupIterator<'a> {
+    pub(crate) bytes: &'a [u8],
+    pub(crate) offset: usize,
+    pub(crate) num_address_table_lookups: u8,
+    pub(crate) index: u8,
+}
+
+impl<'a> Iterator for AddressTableLookupIterator<'a> {
+    type Item = SVMMessageAddressTableLookup<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index < self.num_address_table_lookups {
+            self.index = self.index.wrapping_add(1);
+
+            // Each ATL has 3 pieces:
+            // 1. Address (Pubkey)
+            // 2. write indexes ([u8])
+            // 3. read indexes ([u8])
+
+            // Advance offset for address of the lookup table.
+            const _: () = assert!(core::mem::align_of::<Pubkey>() == 1, "Pubkey alignment");
+            // SAFETY:
+            // - The offset is checked to be valid in the slice.
+            // - The alignment of Pubkey is 1.
+            // - `Pubkey` is a byte array, it cannot be improperly initialized.
+            let account_key = unsafe { read_type::<Pubkey>(self.bytes, &mut self.offset) }.ok()?;
+
+            // Read the number of write indexes, and then update the offset.
+            let num_write_accounts =
+                optimized_read_compressed_u16(self.bytes, &mut self.offset).ok()?;
+
+            const _: () = assert!(core::mem::align_of::<u8>() == 1, "u8 alignment");
+            // SAFETY:
+            // - The offset is checked to be valid in the byte slice.
+            // - The alignment of u8 is 1.
+            // - The slice length is checked to be valid.
+            // - `u8` cannot be improperly initialized.
+            let writable_indexes =
+                unsafe { read_slice_data::<u8>(self.bytes, &mut self.offset, num_write_accounts) }
+                    .ok()?;
+
+            // Read the number of read indexes, and then update the offset.
+            let num_read_accounts =
+                optimized_read_compressed_u16(self.bytes, &mut self.offset).ok()?;
+
+            const _: () = assert!(core::mem::align_of::<u8>() == 1, "u8 alignment");
+            // SAFETY:
+            // - The offset is checked to be valid in the byte slice.
+            // - The alignment of u8 is 1.
+            // - The slice length is checked to be valid.
+            // - `u8` cannot be improperly initialized.
+            let readonly_indexes =
+                unsafe { read_slice_data::<u8>(self.bytes, &mut self.offset, num_read_accounts) }
+                    .ok()?;
+
+            Some(SVMMessageAddressTableLookup {
+                account_key,
+                writable_indexes,
+                readonly_indexes,
+            })
+        } else {
+            None
+        }
+    }
+}
+
+impl ExactSizeIterator for AddressTableLookupIterator<'_> {
+    fn len(&self) -> usize {
+        usize::from(self.num_address_table_lookups.wrapping_sub(self.index))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use {
@@ -115,7 +188,7 @@ mod tests {
         let bytes = bincode::serialize(&ShortVec::<MessageAddressTableLookup>(vec![])).unwrap();
         let mut offset = 0;
         let meta = AddressTableLookupMeta::try_new(&bytes, &mut offset).unwrap();
-        assert_eq!(meta.num_address_table_lookup, 0);
+        assert_eq!(meta.num_address_table_lookups, 0);
         assert_eq!(meta.offset, 1);
         assert_eq!(offset, bytes.len());
     }
@@ -141,7 +214,7 @@ mod tests {
         .unwrap();
         let mut offset = 0;
         let meta = AddressTableLookupMeta::try_new(&bytes, &mut offset).unwrap();
-        assert_eq!(meta.num_address_table_lookup, 1);
+        assert_eq!(meta.num_address_table_lookups, 1);
         assert_eq!(meta.offset, 1);
         assert_eq!(offset, bytes.len());
     }
@@ -163,7 +236,7 @@ mod tests {
         .unwrap();
         let mut offset = 0;
         let meta = AddressTableLookupMeta::try_new(&bytes, &mut offset).unwrap();
-        assert_eq!(meta.num_address_table_lookup, 2);
+        assert_eq!(meta.num_address_table_lookups, 2);
         assert_eq!(meta.offset, 1);
         assert_eq!(offset, bytes.len());
     }

transaction-view/src/bytes.rs

@@ -3,11 +3,15 @@ use crate::result::{Result, TransactionParsingError};
 /// Check that the buffer has at least `len` bytes remaining starting at
 /// `offset`. Returns Err if the buffer is too short.
 ///
+/// * `bytes` - Slice of bytes to read from.
+/// * `offset` - Current offset into `bytes`.
+/// * `num_bytes` - Number of bytes that must be remaining.
+///
 /// Assumptions:
 /// - The current offset is not greater than `bytes.len()`.
 #[inline(always)]
-pub fn check_remaining(bytes: &[u8], offset: usize, len: usize) -> Result<()> {
-    if len > bytes.len().wrapping_sub(offset) {
+pub fn check_remaining(bytes: &[u8], offset: usize, num_bytes: usize) -> Result<()> {
+    if num_bytes > bytes.len().wrapping_sub(offset) {
         Err(TransactionParsingError)
     } else {
         Ok(())
@@ -29,6 +33,9 @@ pub fn read_byte(bytes: &[u8], offset: &mut usize) -> Result<u8> {
 /// If the buffer is too short or the encoding is invalid, return Err.
 /// `offset` is updated to point to the byte after the compressed u16.
 ///
+/// * `bytes` - Slice of bytes to read from.
+/// * `offset` - Current offset into `bytes`.
+///
 /// Assumptions:
 /// - The current offset is not greater than `bytes.len()`.
 #[allow(dead_code)]
@@ -61,6 +68,7 @@ pub fn read_compressed_u16(bytes: &[u8], offset: &mut usize) -> Result<u16> {
 }
 
 /// Domain-specific optimization for reading a compressed u16.
+///
 /// The compressed u16's are only used for array-lengths in our transaction
 /// format. The transaction packet has a maximum size of 1232 bytes.
 /// This means that the maximum array length within a **valid** transaction is
@@ -70,6 +78,9 @@ pub fn read_compressed_u16(bytes: &[u8], offset: &mut usize) -> Result<u16> {
 /// case, and reads a maximum of 2 bytes.
 /// If the buffer is too short or the encoding is invalid, return Err.
 /// `offset` is updated to point to the byte after the compressed u16.
+///
+/// * `bytes` - Slice of bytes to read from.
+/// * `offset` - Current offset into `bytes`.
 #[inline(always)]
 pub fn optimized_read_compressed_u16(bytes: &[u8], offset: &mut usize) -> Result<u16> {
     let mut result = 0u16;
@@ -98,6 +109,10 @@ pub fn optimized_read_compressed_u16(bytes: &[u8], offset: &mut usize) -> Result
 /// Update the `offset` to point to the byte after an array of length `len` and
 /// of type `T`. If the buffer is too short, return Err.
 ///
+/// * `bytes` - Slice of bytes to read from.
+/// * `offset` - Curernt offset into `bytes`.
+/// * `num_elements` - Number of `T` elements in the array.
+///
 /// Assumptions:
 /// 1. The current offset is not greater than `bytes.len()`.
 /// 2. The size of `T` is small enough such that a usize will not overflow if
@@ -106,9 +121,9 @@ pub fn optimized_read_compressed_u16(bytes: &[u8], offset: &mut usize) -> Result
 pub fn advance_offset_for_array<T: Sized>(
     bytes: &[u8],
     offset: &mut usize,
-    len: u16,
+    num_elements: u16,
 ) -> Result<()> {
-    let array_len_bytes = usize::from(len).wrapping_mul(core::mem::size_of::<T>());
+    let array_len_bytes = usize::from(num_elements).wrapping_mul(core::mem::size_of::<T>());
     check_remaining(bytes, *offset, array_len_bytes)?;
     *offset = offset.wrapping_add(array_len_bytes);
     Ok(())
@@ -117,6 +132,9 @@ pub fn advance_offset_for_array<T: Sized>(
 /// Update the `offset` to point t the byte after the `T`.
 /// If the buffer is too short, return Err.
 ///
+/// * `bytes` - Slice of bytes to read from.
+/// * `offset` - Curernt offset into `bytes`.
+///
 /// Assumptions:
 /// 1. The current offset is not greater than `bytes.len()`.
 /// 2. The size of `T` is small enough such that a usize will not overflow.
@@ -128,6 +146,51 @@ pub fn advance_offset_for_type<T: Sized>(bytes: &[u8], offset: &mut usize) -> Re
     Ok(())
 }
 
+/// Return a reference to the next slice of `T` in the buffer, checking bounds
+/// and advancing the offset.
+/// If the buffer is too short, return Err.
+///
+/// * `bytes` - Slice of bytes to read from.
+/// * `offset` - Curernt offset into `bytes`.
+/// * `num_elements` - Number of `T` elements in the slice.
+///
+/// # Safety
+/// 1. `bytes` must be a valid slice of bytes.
+/// 2. `offset` must be a valid offset into `bytes`.
+/// 3. `bytes + offset` must be properly aligned for `T`.
+/// 4. `T` slice must be validly initialized.
+/// 5. The size of `T` is small enough such that a usize will not overflow if
+///    given the maximum slice size (u16::MAX).
+#[inline(always)]
+pub unsafe fn read_slice_data<'a, T: Sized>(
+    bytes: &'a [u8],
+    offset: &mut usize,
+    num_elements: u16,
+) -> Result<&'a [T]> {
+    let current_ptr = bytes.as_ptr().wrapping_add(*offset);
+    advance_offset_for_array::<T>(bytes, offset, num_elements)?;
+    Ok(unsafe { core::slice::from_raw_parts(current_ptr as *const T, usize::from(num_elements)) })
+}
+
+/// Return a reference to the next `T` in the buffer, checking bounds and
+/// advancing the offset.
+/// If the buffer is too short, return Err.
+///
+/// * `bytes` - Slice of bytes to read from.
+/// * `offset` - Curernt offset into `bytes`.
+///
+/// # Safety
+/// 1. `bytes` must be a valid slice of bytes.
+/// 2. `offset` must be a valid offset into `bytes`.
+/// 3. `bytes + offset` must be properly aligned for `T`.
+/// 4. `T` must be validly initialized.
+#[inline(always)]
+pub unsafe fn read_type<'a, T: Sized>(bytes: &'a [u8], offset: &mut usize) -> Result<&'a T> {
+    let current_ptr = bytes.as_ptr().wrapping_add(*offset);
+    advance_offset_for_type::<T>(bytes, offset)?;
+    Ok(unsafe { &*(current_ptr as *const T) })
+}
+
 #[cfg(test)]
 mod tests {
     use {

transaction-view/src/instructions_meta.rs

@@ -2,6 +2,7 @@ use {
     crate::{
         bytes::{
             advance_offset_for_array, check_remaining, optimized_read_compressed_u16, read_byte,
+            read_slice_data,
         },
         result::Result,
     },
@@ -82,6 +83,8 @@ impl<'a> Iterator for InstructionsIterator<'a> {
 
     fn next(&mut self) -> Option<Self::Item> {
         if self.index < self.num_instructions {
+            self.index = self.index.wrapping_add(1);
+
             // Each instruction has 3 pieces:
             // 1. Program ID index (u8)
             // 2. Accounts indexes ([u8])
@@ -93,27 +96,29 @@ impl<'a> Iterator for InstructionsIterator<'a> {
             // Read the number of account indexes, and then update the offset
             // to skip over the account indexes.
             let num_accounts = optimized_read_compressed_u16(self.bytes, &mut self.offset).ok()?;
-            // SAFETY: Only returned after we check that there are enough bytes.
-            let accounts = unsafe {
-                core::slice::from_raw_parts(
-                    self.bytes.as_ptr().add(self.offset),
-                    usize::from(num_accounts),
-                )
-            };
-            advance_offset_for_array::<u8>(self.bytes, &mut self.offset, num_accounts).ok()?;
+
+            const _: () = assert!(core::mem::align_of::<u8>() == 1, "u8 alignment");
+            // SAFETY:
+            // - The offset is checked to be valid in the byte slice.
+            // - The alignment of u8 is 1.
+            // - The slice length is checked to be valid.
+            // - `u8` cannot be improperly initialized.
+            let accounts =
+                unsafe { read_slice_data::<u8>(self.bytes, &mut self.offset, num_accounts) }
+                    .ok()?;
 
             // Read the length of the data, and then update the offset to skip
             // over the data.
             let data_len = optimized_read_compressed_u16(self.bytes, &mut self.offset).ok()?;
-            // SAFETY: Only returned after we check that there are enough bytes.
-            let data = unsafe {
-                core::slice::from_raw_parts(
-                    self.bytes.as_ptr().add(self.offset),
-                    usize::from(data_len),
-                )
-            };
-            advance_offset_for_array::<u8>(self.bytes, &mut self.offset, data_len).ok()?;
-            self.index = self.index.wrapping_add(1);
+
+            const _: () = assert!(core::mem::align_of::<u8>() == 1, "u8 alignment");
+            // SAFETY:
+            // - The offset is checked to be valid in the byte slice.
+            // - The alignment of u8 is 1.
+            // - The slice length is checked to be valid.
+            // - `u8` cannot be improperly initialized.
+            let data =
+                unsafe { read_slice_data::<u8>(self.bytes, &mut self.offset, data_len) }.ok()?;
 
             Some(SVMInstruction {
                 program_id_index,