Use non-zero page size

This avoids potential divide by zero error when calling
`AtomicBitmap::new`.

Signed-off-by: Jonathan Woollett-Light <[email protected]>

Author: Jonathan Woollett-Light

src/bitmap/backend/atomic_bitmap.rs

@@ -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};
@@ -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;
@@ -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)]
-        {
+        let page_size = {
             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 = MaybeUninit::zeroed();
+            // SAFETY: It's safe to call `GetSystemInfo` as `sysinfo` is rightly sized
+            // allocated memory.
+            unsafe { GetSystemInfo(sysinfo.as_mut_ptr()) };
+            // SAFETY: It's safe to call `assume_init` as `GetSystemInfo` initializes `sysinfo`.
+            unsafe { sysinfo.assume_init().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(),
+        )
     }
 }
 
@@ -177,13 +175,16 @@ mod tests {
 
     use crate::bitmap::tests::test_bitmap;
 
+    #[allow(clippy::undocumented_unsafe_blocks)]
+    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));
@@ -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));
@@ -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);
     }
 }

src/bitmap/backend/atomic_bitmap_arc.rs

@@ -77,10 +77,14 @@ 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));
+        // SAFETY: `128` is non-zero.
+        let b = AtomicBitmapArc::new(AtomicBitmap::new(0x2000, unsafe {
+            NonZeroUsize::new_unchecked(128)
+        }));
         test_bitmap(&b);
     }
 }

src/bitmap/backend/slice.rs

@@ -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() {
@@ -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);
 
@@ -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);
         }

src/mmap_unix.rs

@@ -446,6 +446,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;
@@ -599,7 +600,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);

src/volatile_memory.rs

@@ -1642,13 +1642,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!(
@@ -2298,14 +2301,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(
@@ -2321,18 +2323,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()));
@@ -2388,9 +2390,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)
         };
@@ -2400,8 +2401,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);
@@ -2428,14 +2432,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,
@@ -2452,7 +2455,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,
@@ -2469,8 +2472,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);
     }
 }