feat: pre-release improvements and harderning.

Signed-off-by: Amlal El Mahrouss <amlal@nekernel.org>

8 files changed, 35 insertions, 59 deletions

diff --git a/src/boot/modules/BootNet/amd64.json b/src/boot/modules/BootNet/amd64.json
index c0037580..5cbefe3d 100644
--- a/src/boot/modules/BootNet/amd64.json
+++ b/src/boot/modules/BootNet/amd64.json
@@ -27,6 +27,7 @@
     "__BOOTZ__",
     "__BOOTZ_STANDALONE__",
     "__NE_AMD64__",
+    "__nekernel_max_cores=8 ",
     "kBootNetVersionHighest=0x0100",
     "kBootNetVersionLowest=0x0100",
     "kBootNetEFIVersion=0x0100"
diff --git a/src/boot/modules/SysChk/SysChk.cc b/src/boot/modules/SysChk/SysChk.cc
index baedf3e8..ebc2be68 100644
--- a/src/boot/modules/SysChk/SysChk.cc
+++ b/src/boot/modules/SysChk/SysChk.cc
@@ -25,9 +25,9 @@ EXTERN_C Int32 SysChkModuleMain(Kernel::HEL::BootInfoHeader* handover) {
   fw_init_efi((EfiSystemTable*) handover->f_FirmwareCustomTables[Kernel::HEL::kHandoverTableST]);
 
 #if defined(__ATA_PIO__)
-  Boot::BDiskFormatFactory<BootDeviceATA> partition_factory;
+  Boot::BDiskFormatFactory<Boot::BootDeviceATA> partition_factory;
 #elif defined(__AHCI__)
-  Boot::BDiskFormatFactory<BootDeviceSATA> partition_factory;
+  Boot::BDiskFormatFactory<Boot::BootDeviceSATA> partition_factory;
 #endif
 
   if (partition_factory.IsPartitionValid()) return kEfiOk;
diff --git a/src/boot/modules/SysChk/amd64-ahci-epm.json b/src/boot/modules/SysChk/amd64-ahci-epm.json
index 8ce9bfd8..48da0f86 100644
--- a/src/boot/modules/SysChk/amd64-ahci-epm.json
+++ b/src/boot/modules/SysChk/amd64-ahci-epm.json
@@ -33,6 +33,7 @@
     "__AHCI__",
     "__SYSCHK__",
     "BOOTZ_EPM_SUPPORT",
+    "__nekernel_max_cores=8 ",
     "__NE_VEPM__",
     "__NE_MODULAR_KERNEL_COMPONENTS__",
     "kChkVersionHighest=0x0100",
diff --git a/src/boot/modules/SysChk/amd64-ahci-gpt.json b/src/boot/modules/SysChk/amd64-ahci-gpt.json
index 80bb433e..91ab6291 100644
--- a/src/boot/modules/SysChk/amd64-ahci-gpt.json
+++ b/src/boot/modules/SysChk/amd64-ahci-gpt.json
@@ -30,6 +30,7 @@
     "__NEOSKRNL__",
     "__NE_AMD64__",
     "__AHCI__",
+    "__nekernel_max_cores=8 ",
     "__SYSCHK__",
     "BOOTZ_GPT_SUPPORT",
     "__NE_MODULAR_KERNEL_COMPONENTS__",
diff --git a/src/boot/modules/SysChk/amd64-pio-epm.json b/src/boot/modules/SysChk/amd64-pio-epm.json
index b1b95d8d..c7276d4e 100644
--- a/src/boot/modules/SysChk/amd64-pio-epm.json
+++ b/src/boot/modules/SysChk/amd64-pio-epm.json
@@ -31,6 +31,7 @@
     "__BOOTZ__",
     "__BOOTZ_STANDALONE__",
     "__NE_AMD64__",
+    "__nekernel_max_cores=8 ",
     "__ATA_PIO__",
     "BOOTZ_EPM_SUPPORT",
     "__NE_VEPM__",
diff --git a/src/boot/modules/SysChk/amd64-pio-gpt.json b/src/boot/modules/SysChk/amd64-pio-gpt.json
index b1a4d38b..88dae9be 100644
--- a/src/boot/modules/SysChk/amd64-pio-gpt.json
+++ b/src/boot/modules/SysChk/amd64-pio-gpt.json
@@ -33,6 +33,7 @@
     "__NE_AMD64__",
     "__ATA_PIO__",
     "__NE_VEPM__",
+    "__nekernel_max_cores=8 ",
     "BOOTZ_GPT_SUPPORT",
     "kChkVersionHighest=0x0100",
     "kChkVersionLowest=0x0100",
diff --git a/src/boot/src/HEL/AMD64/BootEFI.cc b/src/boot/src/HEL/AMD64/BootEFI.cc
index 59b0d239..d275a13d 100644
--- a/src/boot/src/HEL/AMD64/BootEFI.cc
+++ b/src/boot/src/HEL/AMD64/BootEFI.cc
@@ -36,7 +36,7 @@ STATIC Bool boot_init_fb() {
 
   if (BS->LocateProtocol(&kGopGuid, nullptr, (VoidPtr*) &kGop) != kEfiOk) return No;
 
-  // AMLALE: Ok that ain't great, open to fixes.
+  // TODO: Ok that ain't great, open to fixes.
   kGopStride = 4;
 
   return Yes;
@@ -124,38 +124,41 @@ EFI_EXTERN_C EFI_API Int32 BootloaderMain(EfiHandlePtr image_handle, EfiSystemTa
 
   // Fill handover header now.
 
-  handover_hdr->f_BitMapStart = nullptr;           /* Start of bitmap. */
-  handover_hdr->f_BitMapSize  = kHandoverBitMapSz; /* Size of bitmap in bytes. */
+  handover_hdr->f_BitMapStart = nullptr; /* Start of bitmap. */
+  handover_hdr->f_BitMapSize  = 0UL;     /* Size of bitmap in bytes. */
 
-  kHandoverHeader->f_BitMapStart = nullptr;           /* Start of bitmap. */
-  kHandoverHeader->f_BitMapSize  = kHandoverBitMapSz; /* Size of bitmap in bytes. */
+  kHandoverHeader->f_BitMapStart = nullptr; /* Start of bitmap. */
+  kHandoverHeader->f_BitMapSize  = 0UL;     /* Size of bitmap in bytes. */
 
-  // open to patches.
-  UInt16 trials = 15;
-
-  while (BS->AllocatePool(EfiLoaderData, kHandoverHeader->f_BitMapSize,
-                          &kHandoverHeader->f_BitMapStart) != kEfiOk) {
-    --trials;
-
-    if (trials) {
-      writer.Write("BootZ: Unable to allocate sufficient memory, trying again...\r");
+  // Get memory map to determine available memory for bitmap allocation.
+  BS->GetMemoryMap(&size_struct_ptr, struct_ptr, &map_key, &sz_desc, &rev_desc);
 
-      if (kHandoverHeader->f_BitMapSize > 0)
-        kHandoverHeader->f_BitMapSize =
-            kHandoverHeader->f_BitMapSize / 2; /* Size of bitmap in bytes. */
+  // Allocate space for the memory descriptors.
+  BS->AllocatePool(EfiLoaderData, size_struct_ptr, (VoidPtr*) &struct_ptr);
+  BS->GetMemoryMap(&size_struct_ptr, struct_ptr, &map_key, &sz_desc, &rev_desc);
 
-      while (BS->AllocatePool(EfiLoaderData, kHandoverHeader->f_BitMapSize,
-                              &kHandoverHeader->f_BitMapStart) != kEfiOk) {
-        --trials;
+  // Calculate initial bitmap size by summing all free memory pages.
+  UInt64  free_pages      = 0;
+  VoidPtr first_free_page = nullptr;
 
-        if (!trials) {
-          writer.Write("BootZ: Unable to allocate sufficient memory, aborting...\r");
-          Boot::Stop();
-        }
+  for (UInt32 i = 0; i < size_struct_ptr / sz_desc; ++i) {
+    EfiMemoryDescriptor* desc = (EfiMemoryDescriptor*) ((UInt8*) struct_ptr + (i * sz_desc));
+    if (desc->Kind == EfiConventionalMemory) {
+      if (first_free_page == nullptr) {
+        first_free_page = (VoidPtr) desc->PhysicalStart;
       }
+      free_pages += desc->NumberOfPages;
     }
   }
 
+  // Set bitmap to use the first free page region found.
+  kHandoverHeader->f_BitMapStart = first_free_page;
+  handover_hdr->f_BitMapStart    = first_free_page;
+
+  // Convert pages to bytes (assuming 4K pages) for bitmap size.
+  kHandoverHeader->f_BitMapSize = free_pages * 4096;
+  handover_hdr->f_BitMapSize    = free_pages * 4096;
+
   handover_hdr->f_FirmwareCustomTables[Kernel::HEL::kHandoverTableBS] = (VoidPtr) BS;
   handover_hdr->f_FirmwareCustomTables[Kernel::HEL::kHandoverTableST] = (VoidPtr) ST;
 
@@ -175,8 +178,6 @@ EFI_EXTERN_C EFI_API Int32 BootloaderMain(EfiHandlePtr image_handle, EfiSystemTa
     syschk_thread->Start(handover_hdr, NO);
   }
 
-  BS->GetMemoryMap(&size_struct_ptr, struct_ptr, &map_key, &sz_desc, &rev_desc);
-
   handover_hdr->f_FirmwareVendorLen = Boot::BStrLen(sys_table->FirmwareVendor);
 
   handover_hdr->f_Magic   = kHandoverMagic;
diff --git a/src/boot/src/HEL/ARM64/BootEFI.cc b/src/boot/src/HEL/ARM64/BootEFI.cc
index a393e841..6322d311 100644
--- a/src/boot/src/HEL/ARM64/BootEFI.cc
+++ b/src/boot/src/HEL/ARM64/BootEFI.cc
@@ -120,39 +120,9 @@ EFI_EXTERN_C EFI_API Int32 BootloaderMain(EfiHandlePtr image_handle, EfiSystemTa
   }
 
   //-------------------------------------------------------------//
-  // Allocate heap.
+  // TODO: Allocate heap.
   //-------------------------------------------------------------//
 
-  Boot::BootTextWriter writer;
-
-  kHandoverHeader->f_BitMapStart = nullptr;           /* Start of bitmap. */
-  kHandoverHeader->f_BitMapSize  = kHandoverBitMapSz; /* Size of bitmap in bytes. */
-
-  UInt16 trials = 15;
-
-  while (BS->AllocatePool(EfiLoaderData, kHandoverHeader->f_BitMapSize,
-                          &kHandoverHeader->f_BitMapStart) != kEfiOk) {
-    --trials;
-
-    if (trials) {
-      writer.Write("BootZ: Unable to allocate sufficient memory, trying again...\r");
-
-      if (kHandoverHeader->f_BitMapSize > 0)
-        kHandoverHeader->f_BitMapSize =
-            kHandoverHeader->f_BitMapSize / 2; /* Size of bitmap in bytes. */
-
-      while (BS->AllocatePool(EfiLoaderData, kHandoverHeader->f_BitMapSize,
-                              &kHandoverHeader->f_BitMapStart) != kEfiOk) {
-        --trials;
-
-        if (!trials) {
-          writer.Write("BootZ: Unable to allocate sufficient memory, aborting...\r");
-          Boot::Stop();
-        }
-      }
-    }
-  }
-
   // ------------------------------------------ //
   // null these fields, to avoid being reused later.
   // ------------------------------------------ //