Skip to content

Commit 2d5afa0

Browse files
Jonathan Woollett-LightJonathanWoollett-Light
Jonathan Woollett-Light
authored and
JonathanWoollett-Light
committed
Use non-zero page size
This avoids potential divide by zero error when calling `AtomicBitmap::new`. Signed-off-by: Jonathan Woollett-Light <[email protected]>
1 parent 910d296 commit 2d5afa0

File tree

5 files changed

+54
-44
lines changed

5 files changed

+54
-44
lines changed

‎src/bitmap/backend/atomic_bitmap.rs

Lines changed: 26 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
//! Bitmap backend implementation based on atomic integers.
55
6+
use std::num::NonZeroUsize;
67
use std::sync::atomic::{AtomicU64, Ordering};
78

89
use crate::bitmap::{Bitmap, RefSlice, WithBitmapSlice};
@@ -17,14 +18,14 @@ use crate::mmap::NewBitmap;
1718
pub struct AtomicBitmap {
1819
map: Vec<AtomicU64>,
1920
size: usize,
20-
page_size: usize,
21+
page_size: NonZeroUsize,
2122
}
2223

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

137138
impl Default for AtomicBitmap {
138139
fn default() -> Self {
139-
AtomicBitmap::new(0, 0x1000)
140+
// SAFETY: Safe as `0x1000` is non-zero.
141+
AtomicBitmap::new(0, unsafe { NonZeroUsize::new_unchecked(0x1000) })
140142
}
141143
}
142144

143145
#[cfg(feature = "backend-mmap")]
144146
impl NewBitmap for AtomicBitmap {
145147
fn with_len(len: usize) -> Self {
146-
let page_size;
147-
148148
#[cfg(unix)]
149-
{
150-
// SAFETY: There's no unsafe potential in calling this function.
151-
page_size = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) };
152-
}
149+
// SAFETY: There's no unsafe potential in calling this function.
150+
let page_size = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) };
153151

154152
#[cfg(windows)]
155-
{
153+
let page_size = {
156154
use winapi::um::sysinfoapi::{GetSystemInfo, LPSYSTEM_INFO, SYSTEM_INFO};
157-
158-
// It's safe to initialize this object from a zeroed memory region.
159-
let mut sysinfo: SYSTEM_INFO = unsafe { std::mem::zeroed() };
160-
161-
// It's safe to call this method as the pointer is based on the address
162-
// of the previously initialized `sysinfo` object.
163-
unsafe { GetSystemInfo(&mut sysinfo as LPSYSTEM_INFO) };
164-
165-
page_size = sysinfo.dwPageSize;
166-
}
155+
let mut sysinfo = MaybeUninit::zeroed();
156+
// SAFETY: It's safe to call `GetSystemInfo` as `sysinfo` is rightly sized
157+
// allocated memory.
158+
unsafe { GetSystemInfo(sysinfo.as_mut_ptr()) };
159+
// SAFETY: It's safe to call `assume_init` as `GetSystemInfo` initializes `sysinfo`.
160+
unsafe { sysinfo.assume_init().dwPageSize }
161+
};
167162

168163
// The `unwrap` is safe to use because the above call should always succeed on the
169164
// supported platforms, and the size of a page will always fit within a `usize`.
170-
AtomicBitmap::new(len, usize::try_from(page_size).unwrap())
165+
AtomicBitmap::new(
166+
len,
167+
NonZeroUsize::try_from(usize::try_from(page_size).unwrap()).unwrap(),
168+
)
171169
}
172170
}
173171

@@ -177,13 +175,16 @@ mod tests {
177175

178176
use crate::bitmap::tests::test_bitmap;
179177

178+
#[allow(clippy::undocumented_unsafe_blocks)]
179+
const DEFAULT_PAGE_SIZE: NonZeroUsize = unsafe { NonZeroUsize::new_unchecked(128) };
180+
180181
#[test]
181182
fn test_bitmap_basic() {
182183
// Test that bitmap size is properly rounded up.
183-
let a = AtomicBitmap::new(1025, 128);
184+
let a = AtomicBitmap::new(1025, DEFAULT_PAGE_SIZE);
184185
assert_eq!(a.len(), 9);
185186

186-
let b = AtomicBitmap::new(1024, 128);
187+
let b = AtomicBitmap::new(1024, DEFAULT_PAGE_SIZE);
187188
assert_eq!(b.len(), 8);
188189
b.set_addr_range(128, 129);
189190
assert!(!b.is_addr_set(0));
@@ -214,7 +215,7 @@ mod tests {
214215

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

226227
#[test]
227228
fn test_bitmap_impl() {
228-
let b = AtomicBitmap::new(0x2000, 128);
229+
let b = AtomicBitmap::new(0x2000, DEFAULT_PAGE_SIZE);
229230
test_bitmap(&b);
230231
}
231232
}

‎src/bitmap/backend/atomic_bitmap_arc.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,10 +77,14 @@ mod tests {
7777
use super::*;
7878

7979
use crate::bitmap::tests::test_bitmap;
80+
use std::num::NonZeroUsize;
8081

8182
#[test]
8283
fn test_bitmap_impl() {
83-
let b = AtomicBitmapArc::new(AtomicBitmap::new(0x2000, 128));
84+
// SAFETY: `128` is non-zero.
85+
let b = AtomicBitmapArc::new(AtomicBitmap::new(0x2000, unsafe {
86+
NonZeroUsize::new_unchecked(128)
87+
}));
8488
test_bitmap(&b);
8589
}
8690
}

‎src/bitmap/backend/slice.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ mod tests {
9999

100100
use crate::bitmap::tests::{range_is_clean, range_is_dirty, test_bitmap};
101101
use crate::bitmap::AtomicBitmap;
102+
use std::num::NonZeroUsize;
102103

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

109110
{
110-
let bitmap = AtomicBitmap::new(bitmap_size, 1);
111+
let bitmap = AtomicBitmap::new(bitmap_size, NonZeroUsize::MIN);
111112
let slice1 = bitmap.slice_at(0);
112113
let slice2 = bitmap.slice_at(dirty_offset);
113114

@@ -121,7 +122,7 @@ mod tests {
121122
}
122123

123124
{
124-
let bitmap = AtomicBitmap::new(bitmap_size, 1);
125+
let bitmap = AtomicBitmap::new(bitmap_size, NonZeroUsize::MIN);
125126
let slice = bitmap.slice_at(0);
126127
test_bitmap(&slice);
127128
}

‎src/mmap_unix.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -446,6 +446,7 @@ mod tests {
446446
use super::*;
447447

448448
use std::io::Write;
449+
use std::num::NonZeroUsize;
449450
use std::slice;
450451
use std::sync::Arc;
451452
use vmm_sys_util::tempfile::TempFile;
@@ -599,7 +600,7 @@ mod tests {
599600
assert!(r.owned());
600601

601602
let region_size = 0x10_0000;
602-
let bitmap = AtomicBitmap::new(region_size, 0x1000);
603+
let bitmap = AtomicBitmap::new(region_size, unsafe { NonZeroUsize::new_unchecked(0x1000) });
603604
let builder = MmapRegionBuilder::new_with_bitmap(region_size, bitmap)
604605
.with_hugetlbfs(true)
605606
.with_mmap_prot(libc::PROT_READ | libc::PROT_WRITE);

‎src/volatile_memory.rs

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1642,13 +1642,16 @@ mod tests {
16421642
use std::thread::spawn;
16431643

16441644
use matches::assert_matches;
1645+
use std::num::NonZeroUsize;
16451646
use vmm_sys_util::tempfile::TempFile;
16461647

16471648
use crate::bitmap::tests::{
16481649
check_range, range_is_clean, range_is_dirty, test_bytes, test_volatile_memory,
16491650
};
16501651
use crate::bitmap::{AtomicBitmap, RefSlice};
16511652

1653+
const DEFAULT_PAGE_SIZE: NonZeroUsize = unsafe { NonZeroUsize::new_unchecked(0x1000) };
1654+
16521655
#[test]
16531656
fn test_display_error() {
16541657
assert_eq!(
@@ -2298,14 +2301,13 @@ mod tests {
22982301
let val = 123u64;
22992302
let dirty_offset = 0x1000;
23002303
let dirty_len = size_of_val(&val);
2301-
let page_size = 0x1000;
23022304

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

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

23112313
test_bytes(
@@ -2321,18 +2323,18 @@ mod tests {
23212323

23222324
// Invoke the `VolatileMemory` test helper function.
23232325
{
2324-
let bitmap = AtomicBitmap::new(len, page_size);
2326+
let bitmap = AtomicBitmap::new(len, DEFAULT_PAGE_SIZE);
23252327
let slice = unsafe { VolatileSlice::with_bitmap(buf, len, bitmap.slice_at(0), None) };
23262328
test_volatile_memory(&slice);
23272329
}
23282330

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

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

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

23382340
assert!(range_is_clean(slice.bitmap(), 0, slice.len()));
@@ -2388,9 +2390,8 @@ mod tests {
23882390
fn test_volatile_ref_dirty_tracking() {
23892391
let val = 123u64;
23902392
let mut buf = vec![val];
2391-
let page_size = 0x1000;
23922393

2393-
let bitmap = AtomicBitmap::new(size_of_val(&val), page_size);
2394+
let bitmap = AtomicBitmap::new(size_of_val(&val), DEFAULT_PAGE_SIZE);
23942395
let vref = unsafe {
23952396
VolatileRef::with_bitmap(buf.as_mut_ptr() as *mut u8, bitmap.slice_at(0), None)
23962397
};
@@ -2400,8 +2401,11 @@ mod tests {
24002401
assert!(range_is_dirty(vref.bitmap(), 0, vref.len()));
24012402
}
24022403

2403-
fn test_volatile_array_ref_copy_from_tracking<T>(buf: &mut [T], index: usize, page_size: usize)
2404-
where
2404+
fn test_volatile_array_ref_copy_from_tracking<T>(
2405+
buf: &mut [T],
2406+
index: usize,
2407+
page_size: NonZeroUsize,
2408+
) where
24052409
T: ByteValued + From<u8>,
24062410
{
24072411
let bitmap = AtomicBitmap::new(size_of_val(buf), page_size);
@@ -2428,14 +2432,13 @@ mod tests {
24282432
let dirty_len = size_of_val(&val);
24292433
let index = 0x1000;
24302434
let dirty_offset = dirty_len * index;
2431-
let page_size = 0x1000;
24322435

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

24362439
// Test `ref_at`.
24372440
{
2438-
let bitmap = AtomicBitmap::new(buf.len() * size_of_val(&val), page_size);
2441+
let bitmap = AtomicBitmap::new(buf.len() * size_of_val(&val), DEFAULT_PAGE_SIZE);
24392442
let arr = unsafe {
24402443
VolatileArrayRef::with_bitmap(
24412444
buf.as_mut_ptr() as *mut u8,
@@ -2452,7 +2455,7 @@ mod tests {
24522455

24532456
// Test `store`.
24542457
{
2455-
let bitmap = AtomicBitmap::new(buf.len() * size_of_val(&val), page_size);
2458+
let bitmap = AtomicBitmap::new(buf.len() * size_of_val(&val), DEFAULT_PAGE_SIZE);
24562459
let arr = unsafe {
24572460
VolatileArrayRef::with_bitmap(
24582461
buf.as_mut_ptr() as *mut u8,
@@ -2469,8 +2472,8 @@ mod tests {
24692472
}
24702473

24712474
// Test `copy_from` when size_of::<T>() == 1.
2472-
test_volatile_array_ref_copy_from_tracking(&mut byte_buf, index, page_size);
2475+
test_volatile_array_ref_copy_from_tracking(&mut byte_buf, index, DEFAULT_PAGE_SIZE);
24732476
// Test `copy_from` when size_of::<T>() > 1.
2474-
test_volatile_array_ref_copy_from_tracking(&mut buf, index, page_size);
2477+
test_volatile_array_ref_copy_from_tracking(&mut buf, index, DEFAULT_PAGE_SIZE);
24752478
}
24762479
}

0 commit comments

Comments
 (0)