Skip to content

Commit

Permalink
pr feedbacks
Browse files Browse the repository at this point in the history
  • Loading branch information
@HaoranYi
HaoranYi committed Jul 12, 2024
1 parent e115b96 commit c28a32a
Showing 1 changed file with 110 additions and 91 deletions.
201 changes: 110 additions & 91 deletions accounts-db/src/append_vec.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -864,22 +864,20 @@ impl AppendVec {
/// 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.
///
/// 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<AccountOffsets> {
let stored_size_unaligned =
STORE_META_OVERHEAD.checked_add(stored_meta.data_len as usize)?;
/// *Safety* - The caller must ensure that the `stored_meta.data_len` won't overflow the calculation.
fn next_account_offset(start_offset: usize, stored_meta: &StoredMeta) -> AccountOffsets {
let stored_size_unaligned = STORE_META_OVERHEAD
.checked_add(stored_meta.data_len as usize)
.expect("next_account_offset overflow");
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;

Some(AccountOffsets {
AccountOffsets {
next_account_offset,
offset_to_end_of_data,
stored_size_aligned,
})
}
}

/// Iterate over all accounts and call `callback` with `IndexInfo` for each.
Expand All @@ -903,9 +901,7 @@ impl AppendVec {
// we passed the last useful account
return;
}
let Some(next) = Self::next_account_offset(offset, stored_meta) else {
break;
};
let next = Self::next_account_offset(offset, stored_meta);
if next.offset_to_end_of_data > self.len() {
// data doesn't fit, so don't include this account
break;
Expand Down Expand Up @@ -936,7 +932,10 @@ impl AppendVec {
Self::get_type(bytes_subset, next).unwrap();
let (_hash, next): (&AccountHash, _) =
Self::get_type(bytes_subset, next).unwrap();
let stored_size_aligned = u64_align!(next + (meta.data_len as usize));
let stored_size_aligned = next
.checked_add(meta.data_len as usize)
.expect("next_account_offset overflow");
let stored_size_aligned = u64_align!(stored_size_aligned);
callback(IndexInfo {
index_info: {
IndexInfoInner {
Expand Down Expand Up @@ -1916,9 +1915,15 @@ pub mod tests {
assert_eq!(account_sizes, stored_sizes);
}

#[test_case(StorageAccess::Mmap)]
#[test_case(StorageAccess::File)]
fn test_scan_pubkeys(storage_access: StorageAccess) {
/// A helper function for testing different scenario for scan_*.
///
/// `callback` is used to append invalid account's data for testing.
/// `check_fn` performs the check for the scan.
fn test_scan_helper(
storage_access: StorageAccess,
callback: impl Fn(&PathBuf, usize) -> usize,
check_fn: impl Fn(&AppendVec, &[Pubkey]),
) {
const NUM_ACCOUNTS: usize = 37;
let pubkeys: Vec<_> = std::iter::repeat_with(solana_sdk::pubkey::new_rand)
.take(NUM_ACCOUNTS)
Expand Down Expand Up @@ -1958,113 +1963,127 @@ pub mod tests {
.unwrap(),
);

check_fn(&append_vec, &pubkeys);
};

let total_stored_size = callback(&temp_file.path, total_stored_size);
assert_scan(total_stored_size);
}

/// A helper fn to test `scan_pubkeys`.
fn test_scan_pubkeys_helper(
storage_access: StorageAccess,
callback: impl Fn(&PathBuf, usize) -> usize,
) {
test_scan_helper(storage_access, callback, |append_vec, pubkeys| {
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);
/// Test `scan_pubkey` for a valid account storage.
#[test_case(StorageAccess::Mmap)]
#[test_case(StorageAccess::File)]
fn test_scan_pubkeys(storage_access: StorageAccess) {
test_scan_pubkeys_helper(storage_access, |_, size| size);
}

// Append 1 byte of garbage to cover corner cases for incomplete AccountMeta data.
{
/// Test `scan_pubkey` for storage with incomplete account meta data.
#[test_case(StorageAccess::Mmap)]
#[test_case(StorageAccess::File)]
fn test_scan_pubkeys_incomplete_data(storage_access: StorageAccess) {
test_scan_pubkeys_helper(storage_access, |path, size| {
// Append 1 byte of data at the end of the storage file to simulate
// incomplete account's meta data.
let mut f = OpenOptions::new()
.read(true)
.append(true)
.open(&temp_file.path)
.open(&path)
.unwrap();
f.write_all(&[0xFF]).unwrap();
}
assert_scan(total_stored_size + 1);

const SIZE: usize = std::mem::size_of::<StoredMeta>() - 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.write_all(&[0xFF; SIZE]).unwrap();
}
assert_scan(total_stored_size + SIZE + 1);
size + 1
});
}

/// Test `scan_pubkey` for storage with garbage data that cause overflow in next element offset calculation.
#[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();
}
#[should_panic(expected = "next_account_offset overflow")]
fn test_scan_pubkeys_overflow(storage_access: StorageAccess) {
test_scan_pubkeys_helper(storage_access, |path, size| {
// Append a garbage AccountMeta to simulate next_offset overflow.
{
let mut f = OpenOptions::new()
.read(true)
.append(true)
.open(&path)
.unwrap();
f.write_all(&[0xFF; STORE_META_OVERHEAD]).unwrap();
}
size + STORE_META_OVERHEAD
});
}

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(),
);
/// A helper fn to test scan_index
fn test_scan_index_helper(
storage_access: StorageAccess,
callback: impl Fn(&PathBuf, usize) -> usize,
) {
test_scan_helper(storage_access, callback, |append_vec, pubkeys| {
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);
/// Test `scan_index` for a valid account storage.
#[test_case(StorageAccess::Mmap)]
#[test_case(StorageAccess::File)]
fn test_scan_index(storage_access: StorageAccess) {
test_scan_index_helper(storage_access, |_, size| size);
}

// Append 1 byte of garbage to cover corner cases for incomplete AccountMeta data.
{
/// Test `scan_pubkey` for storage with incomplete account meta data.
#[test_case(StorageAccess::Mmap)]
#[test_case(StorageAccess::File)]
fn test_scan_index_incomplete_data(storage_access: StorageAccess) {
test_scan_index_helper(storage_access, |path, size| {
// Append 1 byte of data at the end of the storage file to simulate
// incomplete account's meta data.
let mut f = OpenOptions::new()
.read(true)
.append(true)
.open(&temp_file.path)
.open(&path)
.unwrap();
f.write_all(&[0xFF]).unwrap();
}
assert_scan(total_stored_size + 1);
size + 1
});
}

const SIZE: usize = std::mem::size_of::<StoredMeta>() - 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.write_all(&[0xFF; SIZE]).unwrap();
}
assert_scan(total_stored_size + SIZE + 1);
/// Test `scan_index` for storage with garbage data that cause overflow in next element offset calculation.
#[test_case(StorageAccess::Mmap)]
#[test_case(StorageAccess::File)]
#[should_panic(expected = "next_account_offset overflow")]
fn test_scan_index_overflow(storage_access: StorageAccess) {
test_scan_index_helper(storage_access, |path, size| {
// Append a garbage AccountMeta to simulate next_offset overflow.
{
let mut f = OpenOptions::new()
.read(true)
.append(true)
.open(&path)
.unwrap();
f.write_all(&[0xFF; STORE_META_OVERHEAD]).unwrap();
}
size + STORE_META_OVERHEAD
});
}
}

0 comments on commit c28a32a

Please sign in to comment.