Kernel/aarch64: Bootstrap the MMU in pre_init similar to RISC-V

The old code had two bugs.
The first bug was caused by us jumping into high memory inside of a C++
function (pre_init). Since the compiler (rightfully so) assumes that
all global variables stay at the same address, it decided to load the
address of g_boot_info (which we pass to init) before we enable the MMU.
This caused the boot_info info argument to point to physical memory.

Normally, this should cause a crash when we try to access this argument
in the kernel itself, since unmap_identity_map() unmaps the
identity map, causing all accesses to low memory to fail.
However, unmap_identity_map() is also broken, since it assumes that all
physical memory is below 0x3f000000 (likely a leftover from when we
only supported the Rasberry Pi 3). This causes it to not unmap the
identity mapping if physical memory is above that address, like on QEMU.

We can easily fix both of these problems if we bootstrap the MMU like
we do on RISC-V. Instead of identity mapping the entire kernel range,
only identity map a single "bootstrap" page that contains the code to
enable the MMU and afterwards jump into virtual memory.
This means that we don't jump into high virtual memory in the middle of
a C++ function, so the compiler can't incorrectly load the g_boot_info
address before the MMU is enabled. Additionally, we only have to unmap
a single page, which is significantly easier than having to unmap some
arbitrary range of physical memory.

This fixes booting on AArch64 systems which have physical memory below
0x3f000000, like the Pi 5.
This regressed in d7f65d06d6.

Kernel/Arch/aarch64/CPU.h

@@ -20,8 +20,7 @@ void dbgln_without_mmu(StringView);
 
 namespace Memory {
 
-void init_page_tables(PhysicalPtr flattened_devicetree_paddr);
-void unmap_identity_map();
+[[noreturn]] void init_page_tables_and_jump_to_init(PhysicalPtr flattened_devicetree_paddr);
 
 }
 

Kernel/Arch/aarch64/MMU.cpp

@@ -29,10 +29,6 @@ extern u8 end_of_kernel_image[];
 
 namespace Kernel::Memory {
 
-// physical memory
-constexpr u32 START_OF_NORMAL_MEMORY = 0x00000000;
-constexpr u32 END_OF_NORMAL_MEMORY = 0x3EFFFFFF;
-
 ALWAYS_INLINE static u64* descriptor_to_pointer(FlatPtr descriptor)
 {
     return (u64*)(descriptor & DESCRIPTOR_MASK);
@@ -159,9 +155,6 @@ static void build_mappings(PageBumpAllocator& allocator, u64* root_table)
 
     auto start_of_physical_kernel_range = PhysicalAddress { bit_cast<PhysicalPtr>(+start_of_kernel_image) };
 
-    // Insert identity mapping
-    insert_entries_for_memory_range(allocator, root_table, start_of_kernel_range.offset(-calculate_physical_to_link_time_address_offset()), end_of_kernel_range.offset(-calculate_physical_to_link_time_address_offset()), start_of_physical_kernel_range, normal_memory_flags);
-
     // Map kernel into high virtual memory
     insert_entries_for_memory_range(allocator, root_table, start_of_kernel_range, end_of_kernel_range, start_of_physical_kernel_range, normal_memory_flags);
 }
@@ -172,7 +165,7 @@ static void switch_to_page_table(u8* page_table)
     Aarch64::Asm::set_ttbr1_el1((FlatPtr)page_table);
 }
 
-static void activate_mmu()
+static void configure_mmu()
 {
     Aarch64::MAIR_EL1 mair_el1 = {};
     mair_el1.Attr[0] = 0xFF;       // Normal memory
@@ -201,15 +194,6 @@ static void activate_mmu()
     tcr_el1.IPS = feature_register.PARange;
 
     Aarch64::TCR_EL1::write(tcr_el1);
-
-    // Enable MMU in the system control register
-    Aarch64::SCTLR_EL1 sctlr_el1 = Aarch64::SCTLR_EL1::read();
-    sctlr_el1.M = 1; // Enable MMU
-    sctlr_el1.C = 1; // Enable data cache
-    sctlr_el1.I = 1; // Enable instruction cache
-    Aarch64::SCTLR_EL1::write(sctlr_el1);
-
-    Aarch64::Asm::flush();
 }
 
 static u64* get_page_directory(u64* root_table, VirtualAddress virtual_addr)
@@ -256,7 +240,52 @@ static void setup_kernel_page_directory(u64* root_table)
     g_boot_info.boot_pdpt = PhysicalAddress((PhysicalPtr)get_page_directory_table(root_table, VirtualAddress { g_boot_info.kernel_mapping_base }));
 }
 
-void init_page_tables(PhysicalPtr flattened_devicetree_paddr)
+// This function has to fit into one page as it will be identity mapped.
+[[gnu::aligned(PAGE_SIZE)]] [[noreturn]] UNMAP_AFTER_INIT static void activate_mmu(BootInfo const& info, u64* activate_mmu_pte)
+{
+    // Enable the MMU. This will immediately take effect, but we won't crash as this function is identity mapped.
+    auto offset = calculate_physical_to_link_time_address_offset();
+    register FlatPtr x0 asm("x0") = bit_cast<FlatPtr>(&info) + offset;
+    asm volatile(
+        R"(
+            // Enable the MMU, data cache, and instruction cache.
+            mrs x1, sctlr_el1
+            mov w2, #(1 << 0) | (1 << 2) | (1 << 12)
+            orr x1, x1, x2
+            msr sctlr_el1, x1
+            isb
+
+            // Continue execution at high virtual address.
+            adrp x1, 1f
+            add x1, x1, :lo12:1f
+            add x1, x1, %[offset]
+            br x1
+        1:
+
+            // Add offset to the stack pointer, such that it is also using the mapping in high virtual memory.
+            add sp, sp, %[offset]
+
+            // Zero the PTE which identity maps this function.
+            add x1, %[activate_mmu_pte], %[offset]
+            str xzr, [x1]
+            dsb ishst
+            tlbi vmalle1
+            dsb ish
+            isb
+
+            mov lr, xzr
+            mov fp, xzr
+
+            b init
+        )"
+        :
+        : "r"(x0), [offset] "r"(offset), [activate_mmu_pte] "r"(activate_mmu_pte)
+        : "x1", "x2", "memory");
+
+    VERIFY_NOT_REACHED();
+}
+
+[[noreturn]] void init_page_tables_and_jump_to_init(PhysicalPtr flattened_devicetree_paddr)
 {
     ::DeviceTree::FlattenedDeviceTreeHeader* fdt_header = bit_cast<::DeviceTree::FlattenedDeviceTreeHeader*>(flattened_devicetree_paddr);
     if (fdt_header->magic != 0xd00dfeed)
@@ -285,29 +314,19 @@ void init_page_tables(PhysicalPtr flattened_devicetree_paddr)
     setup_quickmap_page_table(allocator, root_table);
     setup_kernel_page_directory(root_table);
 
+    // Identity map the `activate_mmu` function and save the level 3 table address in order to remove the identity mapping in `activate_mmu` again.
+    auto const activate_mmu_vaddr = VirtualAddress { bit_cast<FlatPtr>(&activate_mmu) };
+    auto const activate_mmu_paddr = PhysicalAddress { bit_cast<PhysicalPtr>(&activate_mmu) };
+
+    u64* activate_mmu_level3_table = insert_page_table(allocator, root_table, activate_mmu_vaddr);
+
+    size_t activate_mmu_level3_index = (activate_mmu_vaddr.get() >> 12) & 0x1ff;
+    activate_mmu_level3_table[activate_mmu_level3_index] = activate_mmu_paddr.get();
+    activate_mmu_level3_table[activate_mmu_level3_index] |= ACCESS_FLAG | PAGE_DESCRIPTOR | INNER_SHAREABLE | NORMAL_MEMORY;
+
     switch_to_page_table(page_tables_phys_start);
-    activate_mmu();
-}
-
-void unmap_identity_map()
-{
-    auto start_of_physical_memory = FlatPtr(START_OF_NORMAL_MEMORY);
-
-    u64 level0_idx = (start_of_physical_memory >> 39) & 0x1FF;
-    u64 level1_idx = (start_of_physical_memory >> 30) & 0x1FF;
-
-    u64* level1_table = (u64*)page_tables_phys_start;
-
-    auto level2_table = FlatPtr(descriptor_to_pointer(level1_table[level0_idx]));
-    if (!level2_table)
-        panic_without_mmu("Could not find table!"sv);
-
-    // NOTE: The function descriptor_to_pointer returns a physical address, but we want to unmap that range
-    //       so, the pointer must be converted to a virtual address by adding calculate_physical_to_link_time_address_offset().
-    level2_table += calculate_physical_to_link_time_address_offset();
-
-    // Unmap the complete identity map
-    ((u64*)level2_table)[level1_idx] = 0;
+    configure_mmu();
+    activate_mmu(g_boot_info, &activate_mmu_level3_table[activate_mmu_level3_index]);
 }
 
 }

Kernel/Arch/aarch64/pre_init.cpp

@@ -6,6 +6,7 @@
 
 #include <AK/StringView.h>
 #include <AK/Types.h>
+#include <Kernel/Arch/Processor.h>
 #include <Kernel/Arch/aarch64/CPU.h>
 #include <Kernel/Sections.h>
 
@@ -39,42 +40,8 @@ extern "C" [[noreturn]] void pre_init(PhysicalPtr flattened_devicetree_paddr)
     // exception level the kernel should run at.
     initialize_exceptions();
 
-    // Next step is to set up page tables and enable the MMU.
-    Memory::init_page_tables(flattened_devicetree_paddr);
-
-    // At this point the MMU is enabled, physical memory is identity mapped,
-    // and the kernel is also mapped into higher virtual memory. However we are still executing
-    // from the physical memory address, so we have to jump to the kernel in high memory. We also need to
-    // switch the stack pointer to high memory, such that we can unmap the identity mapping.
-
-    // Continue execution at high virtual address.
-    asm volatile(
-        "adrp x0, 1f \n"
-        "add x0, x0, :lo12:1f \n"
-        "add x0, x0, %[base] \n"
-        "br x0 \n"
-        "1: \n" ::[base] "r"(g_boot_info.physical_to_virtual_offset)
-        : "x0");
-
-    // Add kernel_mapping_base to the stack pointer, such that it is also using the mapping
-    // in high virtual memory.
-    asm volatile(
-        "mov x0, %[base] \n"
-        "add sp, sp, x0 \n" ::[base] "r"(g_boot_info.physical_to_virtual_offset)
-        : "x0");
-
-    // We can now unmap the identity map as everything is running in high virtual memory at this point.
-    Memory::unmap_identity_map();
-
-    // Clear the frame pointer (x29) and link register (x30) to make sure the kernel cannot backtrace
-    // into this code, and jump to actual init function in the kernel.
-    register FlatPtr const x0 asm("x0") = bit_cast<FlatPtr>(&g_boot_info);
-    asm volatile(
-        "mov x29, xzr \n"
-        "mov x30, xzr \n"
-        "b init \n" ::"r"(x0));
-
-    VERIFY_NOT_REACHED();
+    // Set up page tables, enable the MMU, and jump to init.
+    Memory::init_page_tables_and_jump_to_init(flattened_devicetree_paddr);
 }
 
 }