1 files changed, 49 insertions, 18 deletions

diff --git a/arch/x86_64/kapi/memory.cpp b/arch/x86_64/kapi/memory.cpp
index a9e1216..5b870d5 100644
--- a/arch/x86_64/kapi/memory.cpp
+++ b/arch/x86_64/kapi/memory.cpp
@@ -1,18 +1,19 @@
-#include "kapi/memory.hpp"
+#include <kapi/memory.hpp>
 
-#include "kapi/boot.hpp"
-#include "kapi/system.hpp"
+#include <arch/boot/boot.hpp>
+#include <arch/boot/ld.hpp>
+#include <arch/cpu/registers.hpp>
+#include <arch/memory/higher_half_mapper.hpp>
+#include <arch/memory/kernel_mapper.hpp>
+#include <arch/memory/page_table.hpp>
+#include <arch/memory/page_utilities.hpp>
+#include <arch/memory/region_allocator.hpp>
 
-#include "arch/boot/boot.hpp"
-#include "arch/boot/ld.hpp"
-#include "arch/cpu/registers.hpp"
-#include "arch/memory/higher_half_mapper.hpp"
-#include "arch/memory/kernel_mapper.hpp"
-#include "arch/memory/page_table.hpp"
-#include "arch/memory/page_utilities.hpp"
-#include "arch/memory/region_allocator.hpp"
+#include <kapi/boot.hpp>
+#include <kapi/system.hpp>
 
 #include <kstd/print>
+#include <kstd/units>
 
 #include <multiboot2/constants.hpp>
 #include <multiboot2/information.hpp>
@@ -28,6 +29,8 @@
 #include <span>
 #include <utility>
 
+using namespace kstd::units_literals;
+
 namespace kapi::memory
 {
 
@@ -46,13 +49,14 @@ namespace kapi::memory
       }
 
       auto const & mbi = boot::bootstrap_information.mbi;
-      auto mbi_span = std::span{std::bit_cast<std::byte *>(mbi), mbi->size_bytes()};
+      auto mbi_span = std::span{std::bit_cast<std::byte *>(mbi), static_cast<std::size_t>(mbi->size())};
       auto image_span = std::span{&arch::boot::_start_physical, &arch::boot::_end_physical};
 
       return arch::memory::region_allocator::memory_information{
         .image_range = std::make_pair(physical_address{&image_span.front()}, physical_address{&image_span.back()}),
         .mbi_range = std::make_pair(physical_address{&mbi_span.front()}, physical_address{&mbi_span.back()}),
         .memory_map = *memory_map,
+        .mbi = mbi,
       };
     }
 
@@ -106,10 +110,10 @@ namespace kapi::memory
     [[maybe_unused]] auto remap_multiboot_information(page_mapper & mapper) -> void
     {
       auto mbi_base = std::bit_cast<std::uintptr_t>(boot::bootstrap_information.mbi);
-      auto mbi_size = boot::bootstrap_information.mbi->size_bytes();
+      auto mbi_size = boot::bootstrap_information.mbi->size();
       auto mbi_physical_start = physical_address{mbi_base & ~std::bit_cast<std::uintptr_t>(&arch::boot::TEACHOS_VMA)};
       auto mbi_virtual_start = linear_address{mbi_base};
-      auto mbi_block_count = (mbi_size + PLATFORM_FRAME_SIZE - 1) / PLATFORM_FRAME_SIZE;
+      auto mbi_block_count = (mbi_size + frame::size - 1_B) / frame::size;
 
       for (auto i = 0uz; i < mbi_block_count; ++i)
       {
@@ -119,6 +123,24 @@ namespace kapi::memory
       }
     }
 
+    [[maybe_unused]] auto remap_bootloader_modules(page_mapper & mapper) -> void
+    {
+      std::ranges::for_each(boot::bootstrap_information.mbi->modules(), [&mapper](auto const & module) {
+        auto module_physical_start = physical_address{module.start_address};
+        auto module_virtual_start =
+            linear_address{module.start_address + std::bit_cast<std::uintptr_t>(&arch::boot::TEACHOS_VMA)};
+        auto module_size = static_cast<kstd::units::bytes>(module.end_address - module.start_address);
+        auto module_block_count = (module_size + frame::size - 1_B) / frame::size;
+
+        for (auto i = 0uz; i < module_block_count; ++i)
+        {
+          auto page = page::containing(module_virtual_start) + i;
+          auto frame = frame::containing(module_physical_start) + i;
+          mapper.map(page, frame, page_mapper::flags::writable | page_mapper::flags::supervisor_only);
+        }
+      });
+    }
+
     [[maybe_unused]] auto handoff_to_kernel_pmm(frame_allocator & new_allocator) -> void
     {
       auto memory_map = boot::bootstrap_information.mbi->memory_map();
@@ -128,7 +150,7 @@ namespace kapi::memory
                                  }))
       {
         auto start = frame::containing(physical_address{region.base});
-        auto count = region.size_in_B / page::size;
+        auto count = kstd::units::bytes{region.size_in_B} / page::size;
         new_allocator.release_many({start, count});
       }
 
@@ -148,14 +170,22 @@ namespace kapi::memory
                             [&](auto frame) { new_allocator.mark_used(frame); });
 
       auto mbi_base = std::bit_cast<std::uintptr_t>(boot::bootstrap_information.mbi);
-      auto mbi_size = boot::bootstrap_information.mbi->size_bytes();
+      auto mbi_size = boot::bootstrap_information.mbi->size();
       auto mbi_address = physical_address{mbi_base & ~std::bit_cast<std::uintptr_t>(&arch::boot::TEACHOS_VMA)};
       auto mbi_start = frame::containing(mbi_address);
       auto mbi_end = frame::containing(mbi_address + mbi_size) + 1;
 
-      // TODO BA-FS26: Protect MB2 boot modules
-
       std::ranges::for_each(std::views::iota(mbi_start, mbi_end), [&](auto frame) { new_allocator.mark_used(frame); });
+
+      std::ranges::for_each(boot::bootstrap_information.mbi->modules(), [&](auto const & module) {
+        auto module_physical_start = physical_address{module.start_address};
+        auto module_size = module.end_address - module.start_address;
+        auto module_start = frame::containing(module_physical_start);
+        auto module_end = frame::containing(module_physical_start + module_size) + 1;
+
+        std::ranges::for_each(std::views::iota(module_start, module_end),
+                              [&](auto frame) { new_allocator.mark_used(frame); });
+      });
     }
 
   }  // namespace
@@ -196,6 +226,7 @@ namespace kapi::memory
     remap_kernel(*higher_half_mapper);
     remap_vga_text_mode_buffer(*higher_half_mapper);
     remap_multiboot_information(*higher_half_mapper);
+    remap_bootloader_modules(*higher_half_mapper);
 
     auto current_cr3 = arch::cpu::cr3::read();
     auto old_pml4 = static_cast<arch::memory::page_table *>(current_cr3.address());