Skip to content

Commit

Permalink
Use non-zero page size
Browse files Browse the repository at this point in the history 
This avoids potential divide by zero error.

Signed-off-by: Jonathan Woollett-Light <[email protected]>
  • Loading branch information
Jonathan Woollett-Light committed Oct 31, 2023
1 parent 07ab853 commit 607daa5
Show file tree
Hide file tree
Showing 5 changed files with 53 additions and 44 deletions.
51 changes: 26 additions & 25 deletions src/bitmap/backend/atomic_bitmap.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

//! Bitmap backend implementation based on atomic integers.

use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicU64, Ordering};

use crate::bitmap::{Bitmap, RefSlice, WithBitmapSlice};
Expand All @@ -17,14 +18,14 @@ use crate::mmap::NewBitmap;
pub struct AtomicBitmap {
map: Vec<AtomicU64>,
size: usize,
page_size: usize,
page_size: NonZeroUsize,
}

#[allow(clippy::len_without_is_empty)]
impl AtomicBitmap {
/// Create a new bitmap of `byte_size`, with one bit per page. This is effectively
/// rounded up, and we get a new vector of the next multiple of 64 bigger than `bit_size`.
pub fn new(byte_size: usize, page_size: usize) -> Self {
pub fn new(byte_size: usize, page_size: NonZeroUsize) -> Self {
let mut num_pages = byte_size / page_size;
if byte_size % page_size > 0 {
num_pages += 1;
Expand Down Expand Up @@ -136,38 +137,35 @@ impl Bitmap for AtomicBitmap {

impl Default for AtomicBitmap {
fn default() -> Self {
AtomicBitmap::new(0, 0x1000)
// SAFETY: Safe as `0x1000` is non-zero.
AtomicBitmap::new(0, unsafe { NonZeroUsize::new_unchecked(0x1000) })
}
}

#[cfg(feature = "backend-mmap")]
impl NewBitmap for AtomicBitmap {
fn with_len(len: usize) -> Self {
let page_size;

#[cfg(unix)]
{
// SAFETY: There's no unsafe potential in calling this function.
page_size = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) };
}
// SAFETY: There's no unsafe potential in calling this function.
let page_size = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) };

#[cfg(windows)]
{
// SAFETY: It's safe to initialize this object from a zeroed memory region.
// It's safe to call `GetSystemInfo` as the pointer is based on the address
// of the previously initialized `sysinfo` object.
let page_size = unsafe {
use winapi::um::sysinfoapi::{GetSystemInfo, LPSYSTEM_INFO, SYSTEM_INFO};

// It's safe to initialize this object from a zeroed memory region.
let mut sysinfo: SYSTEM_INFO = unsafe { std::mem::zeroed() };

// It's safe to call this method as the pointer is based on the address
// of the previously initialized `sysinfo` object.
unsafe { GetSystemInfo(&mut sysinfo as LPSYSTEM_INFO) };

page_size = sysinfo.dwPageSize;
}
let mut sysinfo: SYSTEM_INFO = std::mem::zeroed();
GetSystemInfo(&mut sysinfo as LPSYSTEM_INFO);
sysinfo.dwPageSize
};

// The `unwrap` is safe to use because the above call should always succeed on the
// supported platforms, and the size of a page will always fit within a `usize`.
AtomicBitmap::new(len, usize::try_from(page_size).unwrap())
AtomicBitmap::new(
len,
NonZeroUsize::try_from(usize::try_from(page_size).unwrap()).unwrap(),
)
}
}

Expand All @@ -177,13 +175,16 @@ mod tests {

use crate::bitmap::tests::test_bitmap;

// SAFETY: `128` is non-zero.
const DEFAULT_PAGE_SIZE: NonZeroUsize = unsafe { NonZeroUsize::new_unchecked(128) };

#[test]
fn test_bitmap_basic() {
// Test that bitmap size is properly rounded up.
let a = AtomicBitmap::new(1025, 128);
let a = AtomicBitmap::new(1025, DEFAULT_PAGE_SIZE);
assert_eq!(a.len(), 9);

let b = AtomicBitmap::new(1024, 128);
let b = AtomicBitmap::new(1024, DEFAULT_PAGE_SIZE);
assert_eq!(b.len(), 8);
b.set_addr_range(128, 129);
assert!(!b.is_addr_set(0));
Expand Down Expand Up @@ -214,7 +215,7 @@ mod tests {

#[test]
fn test_bitmap_out_of_range() {
let b = AtomicBitmap::new(1024, 1);
let b = AtomicBitmap::new(1024, NonZeroUsize::MIN);
// Set a partial range that goes beyond the end of the bitmap
b.set_addr_range(768, 512);
assert!(b.is_addr_set(768));
Expand All @@ -225,7 +226,7 @@ mod tests {

#[test]
fn test_bitmap_impl() {
let b = AtomicBitmap::new(0x2000, 128);
let b = AtomicBitmap::new(0x2000, DEFAULT_PAGE_SIZE);
test_bitmap(&b);
}
}
5 changes: 4 additions & 1 deletion src/bitmap/backend/atomic_bitmap_arc.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,10 +77,13 @@ mod tests {
use super::*;

use crate::bitmap::tests::test_bitmap;
use std::num::NonZeroUsize;

#[test]
fn test_bitmap_impl() {
let b = AtomicBitmapArc::new(AtomicBitmap::new(0x2000, 128));
let b = AtomicBitmapArc::new(AtomicBitmap::new(0x2000, unsafe {
NonZeroUsize::new_unchecked(128)
}));
test_bitmap(&b);
}
}
5 changes: 3 additions & 2 deletions src/bitmap/backend/slice.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ mod tests {

use crate::bitmap::tests::{range_is_clean, range_is_dirty, test_bitmap};
use crate::bitmap::AtomicBitmap;
use std::num::NonZeroUsize;

#[test]
fn test_slice() {
Expand All @@ -107,7 +108,7 @@ mod tests {
let dirty_len = 0x100;

{
let bitmap = AtomicBitmap::new(bitmap_size, 1);
let bitmap = AtomicBitmap::new(bitmap_size, NonZeroUsize::MIN);
let slice1 = bitmap.slice_at(0);
let slice2 = bitmap.slice_at(dirty_offset);

Expand All @@ -121,7 +122,7 @@ mod tests {
}

{
let bitmap = AtomicBitmap::new(bitmap_size, 1);
let bitmap = AtomicBitmap::new(bitmap_size, NonZeroUsize::MIN);
let slice = bitmap.slice_at(0);
test_bitmap(&slice);
}
Expand Down
3 changes: 2 additions & 1 deletion src/mmap_unix.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -445,6 +445,7 @@ mod tests {
use super::*;

use std::io::Write;
use std::num::NonZeroUsize;
use std::slice;
use std::sync::Arc;
use vmm_sys_util::tempfile::TempFile;
Expand Down Expand Up @@ -598,7 +599,7 @@ mod tests {
assert!(r.owned());

let region_size = 0x10_0000;
let bitmap = AtomicBitmap::new(region_size, 0x1000);
let bitmap = AtomicBitmap::new(region_size, unsafe { NonZeroUsize::new_unchecked(0x1000) });
let builder = MmapRegionBuilder::new_with_bitmap(region_size, bitmap)
.with_hugetlbfs(true)
.with_mmap_prot(libc::PROT_READ | libc::PROT_WRITE);
Expand Down
33 changes: 18 additions & 15 deletions src/volatile_memory.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1640,13 +1640,16 @@ mod tests {
use std::thread::spawn;

use matches::assert_matches;
use std::num::NonZeroUsize;
use vmm_sys_util::tempfile::TempFile;

use crate::bitmap::tests::{
check_range, range_is_clean, range_is_dirty, test_bytes, test_volatile_memory,
};
use crate::bitmap::{AtomicBitmap, RefSlice};

const DEFAULT_PAGE_SIZE: NonZeroUsize = unsafe { NonZeroUsize::new_unchecked(0x1000) };

#[test]
fn test_display_error() {
assert_eq!(
Expand Down Expand Up @@ -2296,14 +2299,13 @@ mod tests {
let val = 123u64;
let dirty_offset = 0x1000;
let dirty_len = size_of_val(&val);
let page_size = 0x1000;

let len = 0x10000;
let buf = unsafe { std::alloc::alloc_zeroed(Layout::from_size_align(len, 8).unwrap()) };

// Invoke the `Bytes` test helper function.
{
let bitmap = AtomicBitmap::new(len, page_size);
let bitmap = AtomicBitmap::new(len, DEFAULT_PAGE_SIZE);
let slice = unsafe { VolatileSlice::with_bitmap(buf, len, bitmap.slice_at(0), None) };

test_bytes(
Expand All @@ -2319,18 +2321,18 @@ mod tests {

// Invoke the `VolatileMemory` test helper function.
{
let bitmap = AtomicBitmap::new(len, page_size);
let bitmap = AtomicBitmap::new(len, DEFAULT_PAGE_SIZE);
let slice = unsafe { VolatileSlice::with_bitmap(buf, len, bitmap.slice_at(0), None) };
test_volatile_memory(&slice);
}

let bitmap = AtomicBitmap::new(len, page_size);
let bitmap = AtomicBitmap::new(len, DEFAULT_PAGE_SIZE);
let slice = unsafe { VolatileSlice::with_bitmap(buf, len, bitmap.slice_at(0), None) };

let bitmap2 = AtomicBitmap::new(len, page_size);
let bitmap2 = AtomicBitmap::new(len, DEFAULT_PAGE_SIZE);
let slice2 = unsafe { VolatileSlice::with_bitmap(buf, len, bitmap2.slice_at(0), None) };

let bitmap3 = AtomicBitmap::new(len, page_size);
let bitmap3 = AtomicBitmap::new(len, DEFAULT_PAGE_SIZE);
let slice3 = unsafe { VolatileSlice::with_bitmap(buf, len, bitmap3.slice_at(0), None) };

assert!(range_is_clean(slice.bitmap(), 0, slice.len()));
Expand Down Expand Up @@ -2386,9 +2388,8 @@ mod tests {
fn test_volatile_ref_dirty_tracking() {
let val = 123u64;
let mut buf = vec![val];
let page_size = 0x1000;

let bitmap = AtomicBitmap::new(size_of_val(&val), page_size);
let bitmap = AtomicBitmap::new(size_of_val(&val), DEFAULT_PAGE_SIZE);
let vref = unsafe {
VolatileRef::with_bitmap(buf.as_mut_ptr() as *mut u8, bitmap.slice_at(0), None)
};
Expand All @@ -2398,8 +2399,11 @@ mod tests {
assert!(range_is_dirty(vref.bitmap(), 0, vref.len()));
}

fn test_volatile_array_ref_copy_from_tracking<T>(buf: &mut [T], index: usize, page_size: usize)
where
fn test_volatile_array_ref_copy_from_tracking<T>(
buf: &mut [T],
index: usize,
page_size: NonZeroUsize,
) where
T: ByteValued + From<u8>,
{
let bitmap = AtomicBitmap::new(size_of_val(buf), page_size);
Expand All @@ -2426,14 +2430,13 @@ mod tests {
let dirty_len = size_of_val(&val);
let index = 0x1000;
let dirty_offset = dirty_len * index;
let page_size = 0x1000;

let mut buf = vec![0u64; index + 1];
let mut byte_buf = vec![0u8; index + 1];

// Test `ref_at`.
{
let bitmap = AtomicBitmap::new(buf.len() * size_of_val(&val), page_size);
let bitmap = AtomicBitmap::new(buf.len() * size_of_val(&val), DEFAULT_PAGE_SIZE);
let arr = unsafe {
VolatileArrayRef::with_bitmap(
buf.as_mut_ptr() as *mut u8,
Expand All @@ -2450,7 +2453,7 @@ mod tests {

// Test `store`.
{
let bitmap = AtomicBitmap::new(buf.len() * size_of_val(&val), page_size);
let bitmap = AtomicBitmap::new(buf.len() * size_of_val(&val), DEFAULT_PAGE_SIZE);
let arr = unsafe {
VolatileArrayRef::with_bitmap(
buf.as_mut_ptr() as *mut u8,
Expand All @@ -2467,8 +2470,8 @@ mod tests {
}

// Test `copy_from` when size_of::<T>() == 1.
test_volatile_array_ref_copy_from_tracking(&mut byte_buf, index, page_size);
test_volatile_array_ref_copy_from_tracking(&mut byte_buf, index, DEFAULT_PAGE_SIZE);
// Test `copy_from` when size_of::<T>() > 1.
test_volatile_array_ref_copy_from_tracking(&mut buf, index, page_size);
test_volatile_array_ref_copy_from_tracking(&mut buf, index, DEFAULT_PAGE_SIZE);
}
}

0 comments on commit 607daa5

Please sign in to comment.