Virtualizing memory after exiting boot services

src/boot/loader.c

@@ -52,8 +52,8 @@ EFI_STATUS loader_load_kernel(void **kernel_image, uint64_t *length) {
             CHECK_EFI_STATUS_OR_RETURN(status, "Freeing kernel file info memory");
 
             UINTN page_count = ((buffer_size - 1) / 0x1000) + 1;
-            EFI_PHYSICAL_ADDRESS addr = 0x100000; // Try to load the kernel in at 1MiB
-            EFI_MEMORY_TYPE mem_type = 0xFFFFFFFF;  // Special value to tell the kernel it's here
+            EFI_PHYSICAL_ADDRESS addr = KERNEL_PHYS_ADDRESS; // Try to load the kernel in at 1MiB
+            EFI_MEMORY_TYPE mem_type = KERNEL_MEMTYPE;  // Special value to tell the kernel it's here
             status = ST->BootServices->AllocatePages(AllocateAddress, mem_type, page_count, &addr);
             if (status == EFI_NOT_FOUND) {
                 // couldn't get the address we wanted, try loading the kernel anywhere

src/boot/loader.h

@@ -1,6 +1,18 @@
 #pragma once
 #include <efi.h>
 
+#ifndef KERNEL_PHYS_ADDRESS
+#define KERNEL_PHYS_ADDRESS 0x0000000000100000
+#endif
+
+#ifndef KERNEL_VIRT_ADDRESS
+#define KERNEL_VIRT_ADDRESS 0x7fff000000000000
+#endif
+
+#ifndef KERNEL_MEMTYPE
+#define KERNEL_MEMTYPE 0xffffffff
+#endif
+
 #ifndef KERNEL_FILENAME
 #define KERNEL_FILENAME L"kernel.bin"
 #endif

src/boot/main.c

@@ -33,8 +33,6 @@ efi_main (EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable)
     CHECK_EFI_STATUS_OR_FAIL(status);
 	con_status_ok();
 
-	// memory_dump_map();
-
     con_status_begin(L"Loading kernel into memory...");
     void *kernel_image = NULL;
 	uint64_t kernel_length = 0;
@@ -43,28 +41,46 @@ efi_main (EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable)
 	Print(L" %u bytes at 0x%x", kernel_length, kernel_image);
     con_status_ok();
 
+	//memory_dump_map();
+
+	con_status_begin(L"Exiting boot services...");
+	UINTN memmap_size = 0, memmap_key = 0;
+	UINTN desc_size = 0;
+    UINT32 desc_version = 0;
+    EFI_MEMORY_DESCRIPTOR *memory_map;
+
+    status = memory_mark_address_for_update((void**)&ST);
+	CHECK_EFI_STATUS_OR_FAIL(status);
+
+    status = memory_mark_address_for_update((void**)&kernel_image);
+	CHECK_EFI_STATUS_OR_FAIL(status);
+
+    status = memory_get_map(&memory_map,
+            &memmap_size, &memmap_key,
+            &desc_size, &desc_version);
+	CHECK_EFI_STATUS_OR_FAIL(status);
+
+	status = ST->BootServices->ExitBootServices(ImageHandle, memmap_key);
+	CHECK_EFI_STATUS_OR_FAIL(status);
+
+    status = memory_virtualize(
+            &kernel_image,
+            memory_map, memmap_size,
+            desc_size, desc_version);
+
+    CHECK_EFI_STATUS_OR_FAIL(status);
+
     void (*kernel_main)() = kernel_image;
     kernel_main();
 
     /*
+	con_status_ok();
+
 	con_status_begin(L"Virtualizing memory...");
-    status = memory_virtualize();
-    CHECK_EFI_STATUS_OR_FAIL(status);
 	con_status_ok();
 
-	UINTN memmap_size = 0;
-	EFI_MEMORY_DESCRIPTOR *memmap;
-	UINTN memmap_key;
-	UINTN desc_size;
-	UINT32 desc_version;
-
-	con_status_begin(L"Exiting boot services");
-	memmap = LibMemoryMap(&memmap_size, &memmap_key,
-		&desc_size, &desc_version);
-
-	status = ST->BootServices->ExitBootServices(ImageHandle, memmap_key);
-	CHECK_EFI_STATUS_OR_FAIL(status);
-	con_status_ok();
+    void (*kernel_main)() = kernel_image;
+    kernel_main();
 
 	con_status_begin(L"Setting virtual address map");
 	status = ST->RuntimeServices->SetVirtualAddressMap(

src/boot/memory.c

@@ -2,6 +2,7 @@
 #include <efilib.h>
 
 #include "memory.h"
+#include "loader.h"
 #include "utility.h"
 
 const UINTN PAGE_SIZE = 4096;
@@ -30,20 +31,35 @@ static const CHAR16 *memory_type_name(UINT32 value) {
     return memory_type_names[value];
 }
 
-void EFIAPI memory_update_addresses(EFI_EVENT UNUSED *event, void UNUSED *context) {
-    //ST->RuntimeServices->ConvertPointer(0, (void **)&BS);
-    ST->RuntimeServices->ConvertPointer(0, (void **)&ST);
+void EFIAPI memory_update_addresses(EFI_EVENT UNUSED *event, void *context) {
+    ST->RuntimeServices->ConvertPointer(0, (void **)context);
 }
 
-EFI_STATUS memory_virtualize() {
-    EFI_STATUS status;
+EFI_STATUS memory_mark_address_for_update(void **pointer) {
     EFI_EVENT event;
+    EFI_STATUS status;
+
+    status = ST->BootServices->CreateEvent(
+            EVT_SIGNAL_VIRTUAL_ADDRESS_CHANGE,
+            TPL_CALLBACK,
+            (EFI_EVENT_NOTIFY)&memory_update_addresses,
+            (void*)pointer,
+            &event);
 
-    status = ST->BootServices->CreateEvent(EVT_SIGNAL_VIRTUAL_ADDRESS_CHANGE,
-            TPL_CALLBACK, (EFI_EVENT_NOTIFY)&memory_update_addresses, NULL, &event);
     CHECK_EFI_STATUS_OR_RETURN(status, "Failed to create memory update event");
+}
 
-    return status;
+EFI_STATUS memory_virtualize(void **kernel_addr, EFI_MEMORY_DESCRIPTOR *memory_map,
+                             UINTN memmap_size, UINTN desc_size, UINT32 desc_version) {
+    unsigned count = memmap_size / desc_size;
+    for (unsigned i=0; i<count; ++i) {
+        if (memory_map[i].Type == KERNEL_MEMTYPE)
+            memory_map[i].VirtualStart = (EFI_VIRTUAL_ADDRESS)KERNEL_VIRT_ADDRESS;
+        else
+            memory_map[i].VirtualStart = (EFI_VIRTUAL_ADDRESS)memory_map[i].PhysicalStart;
+    }
+
+    return ST->RuntimeServices->SetVirtualAddressMap(memmap_size, desc_size, desc_version, memory_map);
 }
 
 EFI_STATUS memory_get_map(EFI_MEMORY_DESCRIPTOR **buffer, UINTN *buffer_size,
@@ -58,7 +74,6 @@ EFI_STATUS memory_get_map(EFI_MEMORY_DESCRIPTOR **buffer, UINTN *buffer_size,
 
     // Give some extra buffer to account for changes.
     *buffer_size = needs_size + 256;
-    Print(L"Trying to load memory map with size %d.\n", *buffer_size);
     status = ST->BootServices->AllocatePool(EfiLoaderData, *buffer_size, (void**)buffer);
     CHECK_EFI_STATUS_OR_RETURN(status, "Failed to allocate space for memory map");
 
@@ -88,10 +103,8 @@ EFI_STATUS memory_dump_map() {
     while (d < end) {
         UINTN size_bytes = d->NumberOfPages * PAGE_SIZE;
 
-        Print(L"Type: %s  Attr: 0x%x\n", memory_type_name(d->Type), d->Attribute);
-        Print(L"\t Physical  %016llx - %016llx\n", d->PhysicalStart, d->PhysicalStart + size_bytes);
-        if (d->VirtualStart != 0)
-            Print(L"\t  Virtual  %016llx - %016llx\n\n", d->VirtualStart, d->VirtualStart + size_bytes);
+        Print(L"%23s ", memory_type_name(d->Type));
+        Print(L"%016llx (%3d pages)\n", d->PhysicalStart, size_bytes / 0x1000);
 
         d = (EFI_MEMORY_DESCRIPTOR *)((uint8_t *)d + desc_size);
     }

src/boot/memory.h

@@ -1,7 +1,14 @@
 #pragma once
 #include <efi.h>
 
-EFI_STATUS memory_virtualize();
+EFI_STATUS memory_mark_address_for_update(void **pointer);
+
+EFI_STATUS memory_virtualize(
+    void **kernel_image,
+    EFI_MEMORY_DESCRIPTOR *memory_map,
+    UINTN memmap_size,
+    UINTN desc_size,
+    UINT32 desc_version);
 
 EFI_STATUS memory_get_map(
     EFI_MEMORY_DESCRIPTOR **buffer,

src/modules/main/main.c

@@ -1,5 +1,8 @@
 __attribute__((section(".text.entry")))
 void kernel_main() {
-    volatile int foo = 13;
-    return;
+    volatile register int foo = 0x1a1b1c10;
+    volatile register int bar = 0;
+
+    while(1)
+        foo = foo | 0xfffffff0 + bar++ | 0xf;
 }