Added HugePageAlignedMem struct for tests

src/arch/x86_64/mod.rs

@@ -252,7 +252,14 @@ mod tests {
 			.try_init();
 
 		let guest_address = GuestPhysAddr::new(0x11111000);
-		let mem = MmapMemory::new(0, MIN_PHYSMEM_SIZE * 2, guest_address, true, true);
+		let mem = MmapMemory::new(
+			0,
+			align_up!(MIN_PHYSMEM_SIZE * 2, 0x20_0000),
+			guest_address,
+			true,
+			true,
+		);
+		println!("mmap memory created {mem:?}");
 		init_guest_mem(
 			unsafe { mem.as_slice_mut() }.try_into().unwrap(),
 			guest_address,

src/arch/x86_64/paging/mod.rs

@@ -113,7 +113,10 @@ pub fn create_gdt_entry(flags: u64, base: u64, limit: u64) -> u64 {
 #[cfg(test)]
 mod tests {
 	use super::*;
-	use crate::consts::{GDT_OFFSET, PDE_OFFSET, PDPTE_OFFSET, PML4_OFFSET};
+	use crate::{
+		consts::{GDT_OFFSET, PDE_OFFSET, PDPTE_OFFSET, PML4_OFFSET},
+		mem::HugePageAlignedMem,
+	};
 
 	#[test]
 	fn test_pagetable_initialization() {
@@ -124,16 +127,13 @@ mod tests {
 
 		let guest_address = GuestPhysAddr::new(0x20_0000);
 
-		let mut mem: Vec<u8> = vec![0; MIN_PHYSMEM_SIZE];
+		let aligned_mem = HugePageAlignedMem::<MIN_PHYSMEM_SIZE>::new();
 		// This will return a pagetable setup that we will check.
-		initialize_pagetables(
-			(&mut mem[0..MIN_PHYSMEM_SIZE]).try_into().unwrap(),
-			guest_address,
-		);
+		initialize_pagetables((aligned_mem.mem).try_into().unwrap(), guest_address);
 
 		// Check PDPTE address
 		let addr_pdpte = GuestPhysAddr::new(u64::from_le_bytes(
-			mem[(PML4_OFFSET as usize)..(PML4_OFFSET as usize + 8)]
+			aligned_mem.mem[(PML4_OFFSET as usize)..(PML4_OFFSET as usize + 8)]
 				.try_into()
 				.unwrap(),
 		));
@@ -144,7 +144,7 @@ mod tests {
 
 		// Check PDE
 		let addr_pde = GuestPhysAddr::new(u64::from_le_bytes(
-			mem[(PDPTE_OFFSET as usize)..(PDPTE_OFFSET as usize + 8)]
+			aligned_mem.mem[(PDPTE_OFFSET as usize)..(PDPTE_OFFSET as usize + 8)]
 				.try_into()
 				.unwrap(),
 		));
@@ -156,7 +156,11 @@ mod tests {
 		// Check PDE's pagetable bits
 		for i in (0..4096).step_by(8) {
 			let pde_addr = (PDE_OFFSET) as usize + i;
-			let entry = u64::from_le_bytes(mem[pde_addr..(pde_addr + 8)].try_into().unwrap());
+			let entry = u64::from_le_bytes(
+				aligned_mem.mem[pde_addr..(pde_addr + 8)]
+					.try_into()
+					.unwrap(),
+			);
 			assert!(
 				PageTableFlags::from_bits_truncate(entry)
 					.difference(
@@ -173,7 +177,8 @@ mod tests {
 		let gdt_results = [0x0, 0xAF9B000000FFFF, 0xCF93000000FFFF];
 		for (i, res) in gdt_results.iter().enumerate() {
 			let gdt_addr = GDT_OFFSET as usize + i * 8;
-			let gdt_entry = u64::from_le_bytes(mem[gdt_addr..gdt_addr + 8].try_into().unwrap());
+			let gdt_entry =
+				u64::from_le_bytes(aligned_mem.mem[gdt_addr..gdt_addr + 8].try_into().unwrap());
 			assert_eq!(*res, gdt_entry);
 		}
 	}

src/mem.rs

@@ -180,6 +180,42 @@ impl Index<usize> for MmapMemory {
 	}
 }
 
+#[cfg(test)]
+/// Wrapper aroud a memory allocation that is aligned to x86 HugePages
+/// (`0x20_0000`). Intended for testing purposes only
+pub(crate) struct HugePageAlignedMem<'a, const SIZE: usize> {
+	ptr: *mut u8,
+	pub mem: &'a mut [u8],
+}
+#[cfg(test)]
+impl<'a, const SIZE: usize> HugePageAlignedMem<'a, SIZE> {
+	pub fn new() -> Self {
+		use std::alloc::{alloc_zeroed, Layout};
+		// TODO: Make this generic to arbitrary alignments.
+		let layout = Layout::from_size_align(SIZE, 0x20_0000).unwrap();
+		unsafe {
+			let ptr = alloc_zeroed(layout);
+			if ptr.is_null() {
+				panic!("Allocation failed");
+			}
+			Self {
+				ptr,
+				mem: std::slice::from_raw_parts_mut(ptr, SIZE),
+			}
+		}
+	}
+}
+#[cfg(test)]
+impl<'a, const SIZE: usize> Drop for HugePageAlignedMem<'a, SIZE> {
+	fn drop(&mut self) {
+		use std::alloc::{dealloc, Layout};
+		let layout = Layout::from_size_align(SIZE, 0x20_0000).unwrap();
+		unsafe {
+			dealloc(self.ptr, layout);
+		}
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;