Skip to content

Commit

Permalink
[TieredStorage] Boundary check for get_account_offset() (#34531)
Browse files Browse the repository at this point in the history
#### Problem
TieredStorage doesn't perform boundary check in get_account_offset
when the input IndexOffset isn't valid.

#### Summary of Changes
This PR adds two checks.  First, it checks whether the IndexOffset exceeds
the boundary of the index block.  Second, when an index format that has the
same index entries as account entries is used, it also checks whether IndexOffset
is smaller than account_entry_count.

#### Test Plan
Two new tests are added to this PR.
  • Loading branch information
yhchiang-sol authored Dec 21, 2023
1 parent 03386cc commit a583035
Showing 1 changed file with 92 additions and 7 deletions.
99 changes: 92 additions & 7 deletions accounts-db/src/tiered_storage/index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,16 +110,26 @@ impl IndexBlockFormat {
footer: &TieredStorageFooter,
index_offset: IndexOffset,
) -> TieredStorageResult<Offset> {
match self {
let offset = match self {
Self::AddressAndBlockOffsetOnly => {
let offset = footer.index_block_offset as usize
debug_assert!(index_offset.0 < footer.account_entry_count);
footer.index_block_offset as usize
+ std::mem::size_of::<Pubkey>() * footer.account_entry_count as usize
+ std::mem::size_of::<Offset>() * index_offset.0 as usize;
let (account_offset, _) = get_pod::<Offset>(mmap, offset)?;

Ok(*account_offset)
+ std::mem::size_of::<Offset>() * index_offset.0 as usize
}
}
};

debug_assert!(
offset.saturating_add(std::mem::size_of::<Offset>())
<= footer.owners_block_offset as usize,
"reading IndexOffset ({}) would exceed index block boundary ({}).",
offset,
footer.owners_block_offset,
);

let (account_offset, _) = get_pod::<Offset>(mmap, offset)?;

Ok(*account_offset)
}

/// Returns the size of one index entry.
Expand Down Expand Up @@ -266,4 +276,79 @@ mod tests {
.get_account_address(&mmap, &footer, IndexOffset(2))
.unwrap();
}

#[test]
#[should_panic(expected = "index_offset.0 < footer.account_entry_count")]
fn test_get_account_offset_out_of_bounds() {
let temp_dir = TempDir::new().unwrap();
let path = temp_dir
.path()
.join("test_get_account_offset_out_of_bounds");

let footer = TieredStorageFooter {
account_entry_count: 100,
index_block_format: IndexBlockFormat::AddressAndBlockOffsetOnly,
..TieredStorageFooter::default()
};

{
// we only writes a footer here as the test should hit an assert
// failure before it actually reads the file.
let file = TieredStorageFile::new_writable(&path).unwrap();
footer.write_footer_block(&file).unwrap();
}

let file = OpenOptions::new()
.read(true)
.create(false)
.open(&path)
.unwrap();
let mmap = unsafe { MmapOptions::new().map(&file).unwrap() };
footer
.index_block_format
.get_account_offset::<HotAccountOffset>(
&mmap,
&footer,
IndexOffset(footer.account_entry_count),
)
.unwrap();
}

#[test]
#[should_panic(expected = "would exceed index block boundary")]
fn test_get_account_offset_exceeds_index_block_boundary() {
let temp_dir = TempDir::new().unwrap();
let path = temp_dir
.path()
.join("test_get_account_offset_exceeds_index_block_boundary");

let footer = TieredStorageFooter {
account_entry_count: 100,
index_block_format: IndexBlockFormat::AddressAndBlockOffsetOnly,
index_block_offset: 1024,
// only holds one index entry
owners_block_offset: 1024 + std::mem::size_of::<HotAccountOffset>() as u64,
..TieredStorageFooter::default()
};

{
// we only write a footer here as the test should hit an assert
// failure before we actually read the file.
let file = TieredStorageFile::new_writable(&path).unwrap();
footer.write_footer_block(&file).unwrap();
}

let file = OpenOptions::new()
.read(true)
.create(false)
.open(&path)
.unwrap();
let mmap = unsafe { MmapOptions::new().map(&file).unwrap() };
// IndexOffset does not exceeds the account_entry_count but exceeds
// the index block boundary.
footer
.index_block_format
.get_account_offset::<HotAccountOffset>(&mmap, &footer, IndexOffset(2))
.unwrap();
}
}

0 comments on commit a583035

Please sign in to comment.