From 14ee9c2dc7ed8b91ec18511f26d769026b9a6c28 Mon Sep 17 00:00:00 2001
From: Ruihan Li <lrh2000@pku.edu.cn>
Date: Tue, 26 Dec 2023 02:20:03 +0800
Subject: [PATCH] Fix missing kernel/initramfs memory regions

---
 .../src/arch/x86/boot/linux_boot/mod.rs         | 16 ++++++++++++++--
 .../src/arch/x86/boot/multiboot/mod.rs          | 11 +----------
 .../src/arch/x86/boot/multiboot2/mod.rs         | 11 +----------
 framework/jinux-frame/src/boot/memory_region.rs | 17 +++++++++++++++++
 4 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/framework/jinux-frame/src/arch/x86/boot/linux_boot/mod.rs b/framework/jinux-frame/src/arch/x86/boot/linux_boot/mod.rs
index 0716da81..fb932305 100644
--- a/framework/jinux-frame/src/arch/x86/boot/linux_boot/mod.rs
+++ b/framework/jinux-frame/src/arch/x86/boot/linux_boot/mod.rs
@@ -6,7 +6,7 @@ use boot_params::E820Type;
 
 use crate::boot::{
     kcmdline::KCmdlineArg,
-    memory_region::{MemoryRegion, MemoryRegionType},
+    memory_region::{non_overlapping_regions_from, MemoryRegion, MemoryRegionType},
     BootloaderAcpiArg, BootloaderFramebufferArg,
 };
 use crate::{config::PHYS_OFFSET, vm::paddr_to_vaddr};
@@ -114,6 +114,8 @@ fn init_memory_regions(memory_regions: &'static Once<Vec<MemoryRegion>>) {
     let mut regions = Vec::<MemoryRegion>::new();
 
     let boot_params = BOOT_PARAMS.get().unwrap();
+
+    // Add regions from E820.
     let num_entries = boot_params.e820_entries as usize;
     for e820_entry in &boot_params.e820_table[0..num_entries] {
         regions.push(MemoryRegion::new(
@@ -123,7 +125,17 @@ fn init_memory_regions(memory_regions: &'static Once<Vec<MemoryRegion>>) {
         ));
     }
 
-    memory_regions.call_once(|| regions);
+    // Add the kernel region.
+    regions.push(MemoryRegion::kernel());
+
+    // Add the initramfs region.
+    regions.push(MemoryRegion::new(
+        boot_params.hdr.ramdisk_image as usize,
+        boot_params.hdr.ramdisk_size as usize,
+        MemoryRegionType::Module,
+    ));
+
+    memory_regions.call_once(|| non_overlapping_regions_from(regions.as_ref()));
 }
 
 /// The entry point of Rust code called by the Linux 64-bit boot compatible bootloader.
diff --git a/framework/jinux-frame/src/arch/x86/boot/multiboot/mod.rs b/framework/jinux-frame/src/arch/x86/boot/multiboot/mod.rs
index 4247b3e5..45a1582b 100644
--- a/framework/jinux-frame/src/arch/x86/boot/multiboot/mod.rs
+++ b/framework/jinux-frame/src/arch/x86/boot/multiboot/mod.rs
@@ -136,16 +136,7 @@ fn init_memory_regions(memory_regions: &'static Once<Vec<MemoryRegion>>) {
     ));
 
     // Add the kernel region.
-    // These are physical addresses provided by the linker script.
-    extern "C" {
-        fn __kernel_start();
-        fn __kernel_end();
-    }
-    regions.push(MemoryRegion::new(
-        __kernel_start as usize,
-        __kernel_end as usize - __kernel_start as usize,
-        MemoryRegionType::Kernel,
-    ));
+    regions.push(MemoryRegion::kernel());
 
     // Add the initramfs area.
     if info.mods_count != 0 {
diff --git a/framework/jinux-frame/src/arch/x86/boot/multiboot2/mod.rs b/framework/jinux-frame/src/arch/x86/boot/multiboot2/mod.rs
index 6f437e5e..dda8fa31 100644
--- a/framework/jinux-frame/src/arch/x86/boot/multiboot2/mod.rs
+++ b/framework/jinux-frame/src/arch/x86/boot/multiboot2/mod.rs
@@ -141,16 +141,7 @@ fn init_memory_regions(memory_regions: &'static Once<Vec<MemoryRegion>>) {
     }
 
     // Add the kernel region since Grub does not specify it.
-    // These are physical addresses provided by the linker script.
-    extern "C" {
-        fn __kernel_start();
-        fn __kernel_end();
-    }
-    regions.push(MemoryRegion::new(
-        __kernel_start as usize,
-        __kernel_end as usize - __kernel_start as usize,
-        MemoryRegionType::Kernel,
-    ));
+    regions.push(MemoryRegion::kernel());
 
     // Add the boot module region since Grub does not specify it.
     let mb2_module_tag = mb2_info.module_tags();
diff --git a/framework/jinux-frame/src/boot/memory_region.rs b/framework/jinux-frame/src/boot/memory_region.rs
index d7b7aa8c..ac2fcdd2 100644
--- a/framework/jinux-frame/src/boot/memory_region.rs
+++ b/framework/jinux-frame/src/boot/memory_region.rs
@@ -40,6 +40,23 @@ impl MemoryRegion {
         MemoryRegion { base, len, typ }
     }
 
+    /// Construct a memory region where kernel sections are loaded.
+    ///
+    /// Most boot protocols do not mark the place where the kernel loads as unusable. In this case,
+    /// we need to explicitly construct and append this memory region.
+    pub fn kernel() -> Self {
+        // These are physical addresses provided by the linker script.
+        extern "C" {
+            fn __kernel_start();
+            fn __kernel_end();
+        }
+        MemoryRegion {
+            base: __kernel_start as usize,
+            len: __kernel_end as usize - __kernel_start as usize,
+            typ: MemoryRegionType::Kernel,
+        }
+    }
+
     /// The physical address of the base of the region.
     pub fn base(&self) -> usize {
         self.base