From a81da361615ad7048a4f5a3d8732a136903fe535 Mon Sep 17 00:00:00 2001 From: HaoranYi Date: Thu, 11 Jul 2024 16:56:42 +0000 Subject: [PATCH] handle overflow for next_account_offset calculation --- accounts-db/src/append_vec.rs | 157 ++++++++++++++++++++++++++++++---- 1 file changed, 140 insertions(+), 17 deletions(-) diff --git a/accounts-db/src/append_vec.rs b/accounts-db/src/append_vec.rs index 2b5241aaf5ad59..ec584fa511e9bb 100644 --- a/accounts-db/src/append_vec.rs +++ b/accounts-db/src/append_vec.rs @@ -483,6 +483,8 @@ impl AppendVec { } /// Creates an appendvec from file without performing sanitize checks or counting the number of accounts + /// + /// This function should only be used for tests, benchmarks and debug tools. For production code, use `new_from_file` instead. #[cfg_attr(not(unix), allow(unused_variables))] pub fn new_from_file_unchecked( path: impl Into, @@ -840,17 +842,23 @@ impl AppendVec { /// data is at the end of each account and is variable sized /// the next account is then aligned on a 64 bit boundary. /// With these helpers, we can skip over reading some of the data depending on what the caller wants. - fn next_account_offset(start_offset: usize, stored_meta: &StoredMeta) -> AccountOffsets { - let stored_size_unaligned = STORE_META_OVERHEAD + stored_meta.data_len as usize; + /// + /// For safety, when `stored_meta.data_len` cause the calculation to overflow, Return None. + fn next_account_offset( + start_offset: usize, + stored_meta: &StoredMeta, + ) -> Option { + let stored_size_unaligned = + STORE_META_OVERHEAD.checked_add(stored_meta.data_len as usize)?; let stored_size_aligned = u64_align!(stored_size_unaligned); let offset_to_end_of_data = start_offset + stored_size_unaligned; let next_account_offset = start_offset + stored_size_aligned; - AccountOffsets { + Some(AccountOffsets { next_account_offset, offset_to_end_of_data, stored_size_aligned, - } + }) } /// Iterate over all accounts and call `callback` with `IndexInfo` for each. @@ -874,7 +882,9 @@ impl AppendVec { // we passed the last useful account return; } - let next = Self::next_account_offset(offset, stored_meta); + let Some(next) = Self::next_account_offset(offset, stored_meta) else { + break; + }; if next.offset_to_end_of_data > self.len() { // data doesn't fit, so don't include this account break; @@ -991,7 +1001,9 @@ impl AppendVec { let Some((stored_meta, _)) = Self::get_type::(slice, offset) else { break; }; - let next = Self::next_account_offset(offset, stored_meta); + let Some(next) = Self::next_account_offset(offset, stored_meta) else { + break; + }; if next.offset_to_end_of_data > self.len() { // data doesn't fit, so don't include break; @@ -1025,7 +1037,9 @@ impl AppendVec { // eof break; }; - let next = Self::next_account_offset(offset, stored_meta); + let Some(next) = Self::next_account_offset(offset, stored_meta) else { + break; + }; if next.offset_to_end_of_data > self.len() { // data doesn't fit, so don't include this pubkey break; @@ -1890,16 +1904,125 @@ pub mod tests { append_vec.flush().unwrap(); } - // now open the append vec with the given storage access method - // then scan the pubkeys to ensure they are correct - let (append_vec, _) = - AppendVec::new_from_file(&temp_file.path, total_stored_size, storage_access).unwrap(); + let assert_scan = |stored_size| { + // now open the append vec with the given storage access method + // then scan the pubkeys to ensure they are correct + let append_vec = ManuallyDrop::new( + AppendVec::new_from_file_unchecked(&temp_file.path, stored_size, storage_access) + .unwrap(), + ); - let mut i = 0; - append_vec.scan_pubkeys(|pubkey| { - assert_eq!(pubkey, pubkeys.get(i).unwrap()); - i += 1; - }); - assert_eq!(i, pubkeys.len()); + let mut i = 0; + append_vec.scan_pubkeys(|pubkey| { + assert_eq!(pubkey, pubkeys.get(i).unwrap()); + i += 1; + }); + assert_eq!(i, pubkeys.len()); + }; + + assert_scan(total_stored_size); + + // Append 1 byte of garbage to cover corner cases for incomplete AccountMeta data. + { + let mut f = OpenOptions::new() + .read(true) + .append(true) + .open(&temp_file.path) + .unwrap(); + f.seek(SeekFrom::End(0)).unwrap(); + f.write_all(&[0xFF]).unwrap(); + } + assert_scan(total_stored_size + 1); + + const SIZE: usize = std::mem::size_of::() - 1; + // Append a garbage AccountMeta to simulate next_offset overflow. + { + let mut f = OpenOptions::new() + .read(true) + .append(true) + .open(&temp_file.path) + .unwrap(); + f.seek(SeekFrom::End(0)).unwrap(); + f.write_all(&[0xFF; SIZE]).unwrap(); + } + assert_scan(total_stored_size + SIZE + 1); + } + + #[test_case(StorageAccess::Mmap)] + #[test_case(StorageAccess::File)] + fn test_scan_index(storage_access: StorageAccess) { + const NUM_ACCOUNTS: usize = 37; + let pubkeys: Vec<_> = std::iter::repeat_with(solana_sdk::pubkey::new_rand) + .take(NUM_ACCOUNTS) + .collect(); + + let mut rng = thread_rng(); + let mut accounts = Vec::with_capacity(pubkeys.len()); + let mut total_stored_size = 0; + for _ in &pubkeys { + let lamports = rng.gen(); + let data_len = rng.gen_range(0..MAX_PERMITTED_DATA_LENGTH) as usize; + let account = AccountSharedData::new(lamports, data_len, &Pubkey::default()); + accounts.push(account); + total_stored_size += aligned_stored_size(data_len); + } + let accounts = accounts; + let total_stored_size = total_stored_size; + + let temp_file = get_append_vec_path("test_scan_index"); + { + // wrap AppendVec in ManuallyDrop to ensure we do not remove the backing file when dropped + let append_vec = + ManuallyDrop::new(AppendVec::new(&temp_file.path, true, total_stored_size)); + let slot = 42; // the specific slot does not matter + let storable_accounts: Vec<_> = std::iter::zip(&pubkeys, &accounts).collect(); + append_vec + .append_accounts(&(slot, storable_accounts.as_slice()), 0) + .unwrap(); + append_vec.flush().unwrap(); + } + + let assert_scan = |stored_size| { + // now open the append vec with the given storage access method + // then scan the pubkeys to ensure they are correct + let append_vec = ManuallyDrop::new( + AppendVec::new_from_file_unchecked(&temp_file.path, stored_size, storage_access) + .unwrap(), + ); + let mut i = 0; + append_vec.scan_index(|info| { + let pubkey = info.index_info.pubkey; + assert_eq!(&pubkey, pubkeys.get(i).unwrap()); + i += 1; + }); + assert_eq!(i, pubkeys.len()); + }; + + assert_scan(total_stored_size); + + // Append 1 byte of garbage to cover corner cases for incomplete AccountMeta data. + { + let mut f = OpenOptions::new() + .read(true) + .append(true) + .open(&temp_file.path) + .unwrap(); + f.seek(SeekFrom::End(0)).unwrap(); + f.write_all(&[0xFF]).unwrap(); + } + assert_scan(total_stored_size + 1); + + const SIZE: usize = std::mem::size_of::() - 1; + // Append a garbage AccountMeta to simulate next_offset overflow. + { + let mut f = OpenOptions::new() + .read(true) + .append(true) + .open(&temp_file.path) + .unwrap(); + f.seek(SeekFrom::End(0)).unwrap(); + f.write_all(&[0xFF; SIZE]).unwrap(); + } + assert_scan(total_stored_size + SIZE + 1); } }