From 4777319220cf8cc5c847cb4052fbfa4f4514a003 Mon Sep 17 00:00:00 2001
From: Jonathan Klimt <jonathan.klimt@eonerc.rwth-aachen.de>
Date: Thu, 19 Sep 2024 10:42:42 +0200
Subject: [PATCH] Enhanced uhyve memory detection

---
 src/arch/x86_64/mm/paging.rs      | 11 +++++-----
 src/arch/x86_64/mm/physicalmem.rs | 35 +++++++++++++++++++------------
 2 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/src/arch/x86_64/mm/paging.rs b/src/arch/x86_64/mm/paging.rs
index d8947df827..fe4d144196 100644
--- a/src/arch/x86_64/mm/paging.rs
+++ b/src/arch/x86_64/mm/paging.rs
@@ -10,11 +10,12 @@ use x86_64::structures::idt::PageFaultErrorCode;
 use x86_64::structures::paging::mapper::{TranslateResult, UnmapError};
 pub use x86_64::structures::paging::PageTableFlags as PageTableEntryFlags;
 use x86_64::structures::paging::{
-	Mapper, Page, PageTableIndex, PhysFrame, RecursivePageTable, Size2MiB, Translate,
+	Mapper, Page, PhysFrame, RecursivePageTable, Size2MiB, Translate,
 };
 
 use crate::arch::x86_64::kernel::processor;
 use crate::arch::x86_64::mm::{physicalmem, PhysAddr, VirtAddr};
+use crate::kernel::get_limit;
 use crate::{env, mm, scheduler};
 
 pub trait PageTableEntryFlagsExt {
@@ -305,11 +306,9 @@ pub fn init_page_tables() {
 		// See https://github.com/hermit-os/uhyve/issues/426
 		let kernel_end_addr = x86_64::VirtAddr::new(mm::kernel_end_address().as_u64());
 		let start_page = Page::<Size2MiB>::from_start_address(kernel_end_addr).unwrap();
-		let end_page = Page::from_page_table_indices_2mib(
-			start_page.p4_index(),
-			start_page.p3_index(),
-			PageTableIndex::new(511),
-		);
+		let end_page =
+			Page::<Size2MiB>::from_start_address(x86_64::VirtAddr::new(get_limit() as u64))
+				.unwrap();
 		let page_range = Page::range_inclusive(start_page, end_page);
 
 		let mut page_table = unsafe { recursive_page_table() };
diff --git a/src/arch/x86_64/mm/physicalmem.rs b/src/arch/x86_64/mm/physicalmem.rs
index cb25c5c4c7..c18943a445 100644
--- a/src/arch/x86_64/mm/physicalmem.rs
+++ b/src/arch/x86_64/mm/physicalmem.rs
@@ -4,7 +4,7 @@ use free_list::{AllocError, FreeList, PageLayout, PageRange};
 use hermit_sync::InterruptTicketMutex;
 use multiboot::information::{MemoryType, Multiboot};
 
-use crate::arch::x86_64::kernel::{get_fdt, get_limit, get_mbinfo};
+use crate::arch::x86_64::kernel::{boot_info, get_fdt, get_limit, get_mbinfo};
 use crate::arch::x86_64::mm::paging::{BasePageSize, PageSize};
 use crate::arch::x86_64::mm::{MultibootMemory, PhysAddr, VirtAddr};
 use crate::{env, mm};
@@ -110,36 +110,45 @@ fn detect_from_uhyve() -> Result<(), ()> {
 		return Err(());
 	}
 
-	let limit = get_limit();
-	assert_ne!(limit, 0);
+	let physmem_end = get_limit();
+	assert_ne!(physmem_end, 0);
 	let mut free_list = PHYSICAL_FREE_LIST.lock();
 	let total_memory;
 
+	let kernel_end = mm::kernel_end_address().as_usize();
 	// add gap for the APIC
-	if limit > KVM_32BIT_GAP_START {
-		let range =
-			PageRange::new(mm::kernel_end_address().as_usize(), KVM_32BIT_GAP_START).unwrap();
+	assert!(
+		!(KVM_32BIT_GAP_START..=KVM_32BIT_GAP_START + KVM_32BIT_GAP_SIZE).contains(&kernel_end),
+		"Kernel was loaded into the KVM 32BIT GAP"
+	);
+	if physmem_end > KVM_32BIT_GAP_START && kernel_end < KVM_32BIT_GAP_START {
+		let range = PageRange::new(kernel_end, KVM_32BIT_GAP_START).unwrap();
 		unsafe {
 			free_list.deallocate(range).unwrap();
 		}
-		if limit > KVM_32BIT_GAP_START + KVM_32BIT_GAP_SIZE {
-			let range = PageRange::new(KVM_32BIT_GAP_START + KVM_32BIT_GAP_SIZE, limit).unwrap();
+		if physmem_end > KVM_32BIT_GAP_START + KVM_32BIT_GAP_SIZE {
+			let range =
+				PageRange::new(KVM_32BIT_GAP_START + KVM_32BIT_GAP_SIZE, physmem_end).unwrap();
 			unsafe {
 				free_list.deallocate(range).unwrap();
 			}
-			total_memory = limit - KVM_32BIT_GAP_SIZE;
+			total_memory = boot_info().hardware_info.phys_addr_range.end
+				- boot_info().hardware_info.phys_addr_range.start
+				- KVM_32BIT_GAP_SIZE as u64;
 		} else {
-			total_memory = KVM_32BIT_GAP_START;
+			total_memory =
+				KVM_32BIT_GAP_START as u64 - boot_info().hardware_info.phys_addr_range.start;
 		}
 	} else {
-		let range = PageRange::new(mm::kernel_end_address().as_usize(), limit).unwrap();
+		let range = PageRange::new(kernel_end, physmem_end).unwrap();
 		unsafe {
 			free_list.deallocate(range).unwrap();
 		}
-		total_memory = limit;
+		total_memory = boot_info().hardware_info.phys_addr_range.end
+			- boot_info().hardware_info.phys_addr_range.start;
 	}
 
-	TOTAL_MEMORY.store(total_memory, Ordering::Relaxed);
+	TOTAL_MEMORY.store(total_memory as usize, Ordering::Relaxed);
 
 	Ok(())
 }