[lib][uefi] Move several memory APIs to uefi_platform.cpp

This allows OEMs to customize exactly how their memories are allocated.
For examples, OEMs might want to allocate certain parts of the UEFI
from high address kernel space, for ease of access in LK world.
Since UEFI spec requires everything to be identity mapped, the default
implementation of memory allocation is identity mapped.

lib/uefi/blockio2_protocols.cpp

@@ -32,6 +32,7 @@
 #include "memory_protocols.h"
 #include "switch_stack.h"
 #include "thread_utils.h"
+#include "uefi_platform.h"
 
 #define LOCAL_TRACE 0
 

lib/uefi/blockio_protocols.cpp

@@ -25,6 +25,7 @@
 #include "io_stack.h"
 #include "memory_protocols.h"
 #include "switch_stack.h"
+#include "uefi_platform.h"
 
 namespace {
 

lib/uefi/boot_service_provider.cpp

@@ -30,7 +30,6 @@
 #include <uefi/protocols/loaded_image_protocol.h>
 #include <uefi/types.h>
 
-#include "arch/defines.h"
 #include "blockio2_protocols.h"
 #include "blockio_protocols.h"
 #include "events.h"
@@ -112,23 +111,6 @@ EfiStatus locate_protocol(const EfiGuid *protocol, void *registration,
   return UNSUPPORTED;
 }
 
-EfiStatus allocate_pages(EfiAllocatorType type, EfiMemoryType memory_type,
-                         size_t pages, EfiPhysicalAddr *memory) {
-  if (memory == nullptr) {
-    return INVALID_PARAMETER;
-  }
-  if (type == ALLOCATE_MAX_ADDRESS && *memory < 0xFFFFFFFF) {
-    printf("allocate_pages(%d, %d, %zu, 0x%llx) unsupported\n", type,
-           memory_type, pages, *memory);
-    return UNSUPPORTED;
-  }
-  *memory = reinterpret_cast<EfiPhysicalAddr>(alloc_page(pages * PAGE_SIZE));
-  if (*memory == 0) {
-    return OUT_OF_RESOURCES;
-  }
-  return SUCCESS;
-}
-
 EfiStatus uninstall_multiple_protocol_interfaces(EfiHandle handle, ...) {
   printf("%s is unsupported\n", __FUNCTION__);
   return UNSUPPORTED;
@@ -176,48 +158,6 @@ EfiTpl raise_tpl(EfiTpl new_tpl) {
   return APPLICATION;
 }
 
-const char *GetImageType(const char16_t *ImageType) {
-  if (memcmp(ImageType, GBL_IMAGE_TYPE_OS_LOAD,
-             sizeof(GBL_IMAGE_TYPE_OS_LOAD)) == 0) {
-    return "os_load";
-  } else if (memcmp(ImageType, GBL_IMAGE_TYPE_FASTBOOT,
-                    sizeof(GBL_IMAGE_TYPE_FASTBOOT)) == 0) {
-    return "fastboot";
-  } else if (memcmp(ImageType, GBL_IMAGE_TYPE_PVMFW_DATA,
-                    sizeof(GBL_IMAGE_TYPE_PVMFW_DATA)) == 0) {
-    return "pvmfw_data";
-  }
-  return "unknown";
-}
-template <typename T>
-T clamp(T n, T lower, T upper) {
-  if (n < lower) {
-    return lower;
-  }
-  if (n > upper) {
-    return upper;
-  }
-  return n;
-}
-
-EfiStatus get_buffer(struct GblEfiImageLoadingProtocol *self,
-                     const GblEfiImageInfo *ImageInfo,
-                     GblEfiImageBuffer *Buffer) {
-  printf("%s(%s, %lu)\n", __FUNCTION__, GetImageType(ImageInfo->ImageType),
-         ImageInfo->SizeBytes);
-
-  // Allow maximum of 128MB buffer
-  const size_t buffer_size =
-      clamp(Buffer->SizeBytes, PAGE_SIZE, 128ul * 1024 * 1024);
-  Buffer->Memory = alloc_page(buffer_size);
-  if (Buffer->Memory == nullptr) {
-    return OUT_OF_RESOURCES;
-  }
-
-  Buffer->SizeBytes = buffer_size;
-  return SUCCESS;
-}
-
 EfiStatus get_verify_partitions(struct GblEfiImageLoadingProtocol *self,
                                 size_t *NumberOfPartitions,
                                 GblEfiPartitionName *Partitions) {
@@ -409,6 +349,10 @@ EfiStatus stall(size_t microseconds) {
   return SUCCESS;
 }
 
+EfiStatus free_pages(EfiPhysicalAddr memory, size_t pages) {
+  return ::free_pages(reinterpret_cast<void *>(memory), pages);
+}
+
 } // namespace
 
 void setup_boot_service_table(EfiBootService *service) {

lib/uefi/configuration_table.cpp

@@ -21,9 +21,9 @@
 #include <string.h>
 #include <uefi/system_table.h>
 
-#include "memory_protocols.h"
-#include "platform.h"
 #include "debug_support.h"
+#include "platform.h"
+#include "uefi_platform.h"
 
 void setup_configuration_table(EfiSystemTable *table) {
   auto &rng = table->configuration_table[table->number_of_table_entries++];
@@ -42,7 +42,8 @@ void setup_configuration_table(EfiSystemTable *table) {
     memcpy(dtb.vendor_table, fdt, fdt_size);
   }
 
-  auto &debug_image_info_table = table->configuration_table[table->number_of_table_entries++];
+  auto &debug_image_info_table =
+      table->configuration_table[table->number_of_table_entries++];
   debug_image_info_table.vendor_guid = EFI_DEBUG_IMAGE_INFO_TABLE_GUID;
   debug_image_info_table.vendor_table = &efi_m_debug_info_table_header;
 }

lib/uefi/debug_support.cpp

@@ -25,6 +25,7 @@
 
 #include "boot_service_provider.h"
 #include "memory_protocols.h"
+#include "uefi_platform.h"
 
 struct EFI_DEVICE_PATH_FILE_PATH_PROTOCOL {
   struct EFI_DEVICE_PATH_PROTOCOL dp;

lib/uefi/memory_protocols.cpp

@@ -16,17 +16,21 @@
  */
 
 #include "memory_protocols.h"
+#include "uefi_platform.h"
 
 #include <arch/defines.h>
 #include <arch/mmu.h>
 #include <kernel/vm.h>
 #include <lib/dlmalloc.h>
+#include <lk/trace.h>
 #include <stdio.h>
 #include <string.h>
 #include <sys/types.h>
 #include <uefi/boot_service.h>
 #include <uefi/types.h>
 
+#define LOCAL_TRACE 0
+
 // MACRO list_for_every_entry cast between int/ptr
 // NOLINTBEGIN(performance-no-int-to-ptr)
 
@@ -60,27 +64,29 @@ mspace get_mspace() {
 vmm_aspace_t *uefi_aspace = nullptr;
 
 void restore_aspace() {
-  vmm_set_active_aspace(old_aspace);
-  vmm_free_aspace(uefi_aspace);
-  uefi_aspace = nullptr;
+  if (uefi_aspace != nullptr) {
+    vmm_set_active_aspace(old_aspace);
+    vmm_free_aspace(uefi_aspace);
+    uefi_aspace = nullptr;
+  }
 }
 
 }  // namespace
 
-EfiStatus free_pages(EfiPhysicalAddr memory, size_t pages) {
-  return free_pages(reinterpret_cast<void *>(memory), pages);
-}
-
-void setup_heap() {
+__WEAK void setup_heap() {
   set_boot_aspace();
   get_mspace();
 }
 
-void reset_heap() {
-  destroy_mspace(space);
-  space = nullptr;
-  free_pages(heap, kHeapSize / PAGE_SIZE);
-  heap = nullptr;
+__WEAK void reset_heap() {
+  if (space != nullptr) {
+    destroy_mspace(space);
+    space = nullptr;
+  }
+  if (heap != nullptr) {
+    free_pages(heap, kHeapSize / PAGE_SIZE);
+    heap = nullptr;
+  }
   restore_aspace();
 }
 
@@ -96,7 +102,7 @@ vmm_aspace_t *set_boot_aspace() {
   return uefi_aspace;
 }
 
-void *alloc_page(size_t size, size_t align_log2) {
+__WEAK void *alloc_page(size_t size, size_t align_log2) {
   size = ROUNDUP(size, PAGE_SIZE);
   auto aspace = set_boot_aspace();
   paddr_t pa{};
@@ -120,7 +126,7 @@ void *alloc_page(size_t size, size_t align_log2) {
   return reinterpret_cast<void *>(pa);
 }
 
-void *alloc_page(void *addr, size_t size, size_t align_log2) {
+__WEAK void *alloc_page(void *addr, size_t size, size_t align_log2) {
   if (addr == nullptr) {
     return alloc_page(size, align_log2);
   }
@@ -146,9 +152,29 @@ void *alloc_page(void *addr, size_t size, size_t align_log2) {
   return addr;
 }
 
-void *uefi_malloc(size_t size) { return mspace_malloc(get_mspace(), size); }
+EfiStatus allocate_pages(EfiAllocatorType type, EfiMemoryType memory_type,
+                         size_t pages, EfiPhysicalAddr *memory) {
+  if (memory == nullptr) {
+    return INVALID_PARAMETER;
+  }
+  if (type == ALLOCATE_MAX_ADDRESS && *memory < 0xFFFFFFFF) {
+    LTRACEF("%s(%d, %d, %zu, 0x%llx) unsupported\n", __FUNCTION__, type,
+            memory_type, pages, *memory);
+    return UNSUPPORTED;
+  }
+  *memory = reinterpret_cast<EfiPhysicalAddr>(alloc_page(pages * PAGE_SIZE));
+  if (*memory == 0) {
+    return OUT_OF_RESOURCES;
+  }
+  return SUCCESS;
+}
 
-EfiStatus allocate_pool(EfiMemoryType pool_type, size_t size, void **buf) {
+__WEAK void *uefi_malloc(size_t size) {
+  return mspace_malloc(get_mspace(), size);
+}
+
+__WEAK EfiStatus allocate_pool(EfiMemoryType pool_type, size_t size,
+                               void **buf) {
   if (buf == nullptr) {
     return INVALID_PARAMETER;
   }
@@ -163,12 +189,12 @@ EfiStatus allocate_pool(EfiMemoryType pool_type, size_t size, void **buf) {
   return OUT_OF_RESOURCES;
 }
 
-EfiStatus free_pool(void *mem) {
+__WEAK EfiStatus free_pool(void *mem) {
   mspace_free(get_mspace(), mem);
   return SUCCESS;
 }
 
-EfiStatus free_pages(void *memory, size_t pages) {
+__WEAK EfiStatus free_pages(void *memory, size_t pages) {
   auto pa = reinterpret_cast<void *>(vaddr_to_paddr(memory));
   if (pa != memory) {
     printf("WARN: virtual address %p is not identity mapped, physical addr: "

lib/uefi/memory_protocols.h

@@ -1,3 +1,20 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
 #ifndef __LIB_UEFI_MEMORY_PROTOCOLS_H
 #define __LIB_UEFI_MEMORY_PROTOCOLS_H
 
@@ -10,14 +27,15 @@ vmm_aspace_t *set_boot_aspace();
 void setup_heap();
 void reset_heap();
 
-void *alloc_page(void *addr, size_t size, size_t align_log2 = PAGE_SIZE_SHIFT);
-void *alloc_page(size_t size, size_t align_log2 = PAGE_SIZE_SHIFT);
-EfiStatus free_pages(EfiPhysicalAddr memory, size_t pages);
-EfiStatus free_pages(void *memory, size_t pages);
-void *identity_map(void *addr, size_t size);
-void *uefi_malloc(size_t size);
-EfiStatus allocate_pool(EfiMemoryType pool_type, size_t size, void **buf);
-EfiStatus free_pool(void *mem);
+EfiStatus allocate_pages(EfiAllocatorType type, EfiMemoryType memory_type,
+                         size_t pages, EfiPhysicalAddr *memory);
+
+// Declaration moved to uefi_platform.h to keep all weak functions declared in
+// one place.
+// void *uefi_malloc(size_t size);
+// EfiStatus allocate_pool(EfiMemoryType pool_type, size_t size, void **buf);
+// EfiStatus free_pool(void *mem);
+// EfiStatus free_pages(void *memory, size_t pages);
 
 EfiStatus get_physical_memory_map(size_t *memory_map_size,
                                   EfiMemoryDescriptor *memory_map,

lib/uefi/uefi_platform.cpp

@@ -21,6 +21,7 @@
 #include <libfdt.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <uefi/protocols/gbl_efi_image_loading_protocol.h>
 #include <uefi/protocols/gbl_efi_os_configuration_protocol.h>
 #include <uefi/types.h>
 
@@ -109,3 +110,52 @@ __WEAK EfiStatus get_timestamp_properties(EfiTimestampProperties *properties) {
   properties->end_value = UINT64_MAX;
   return SUCCESS;
 }
+
+namespace {
+
+const char *GetImageType(const char16_t *ImageType) {
+  if (memcmp(ImageType, GBL_IMAGE_TYPE_OS_LOAD,
+             sizeof(GBL_IMAGE_TYPE_OS_LOAD)) == 0) {
+    return "os_load";
+  } else if (memcmp(ImageType, GBL_IMAGE_TYPE_FASTBOOT,
+                    sizeof(GBL_IMAGE_TYPE_FASTBOOT)) == 0) {
+    return "fastboot";
+  } else if (memcmp(ImageType, GBL_IMAGE_TYPE_PVMFW_DATA,
+                    sizeof(GBL_IMAGE_TYPE_PVMFW_DATA)) == 0) {
+    return "pvmfw_data";
+  }
+  return "unknown";
+}
+template <typename T> T clamp(T n, T lower, T upper) {
+  if (n < lower) {
+    return lower;
+  }
+  if (n > upper) {
+    return upper;
+  }
+  return n;
+}
+
+} // namespace
+
+__WEAK EfiStatus get_buffer(struct GblEfiImageLoadingProtocol *self,
+                            const GblEfiImageInfo *ImageInfo,
+                            GblEfiImageBuffer *Buffer) {
+  printf("%s(%s, %lu)\n", __FUNCTION__, GetImageType(ImageInfo->ImageType),
+         ImageInfo->SizeBytes);
+
+  // Allow maximum of 128MB buffer
+  const size_t buffer_size =
+      clamp(Buffer->SizeBytes, PAGE_SIZE, 128ul * 1024 * 1024);
+  // Bottom line, kernel, ramdisk, device tree must be identity mapped.
+  // Otherwise linux kernel would crash immediately after entering.
+  // Other buffers can be allocated from kernel address space, up to
+  // OEM for customization.
+  Buffer->Memory = alloc_page(buffer_size);
+  if (Buffer->Memory == nullptr) {
+    return OUT_OF_RESOURCES;
+  }
+
+  Buffer->SizeBytes = buffer_size;
+  return SUCCESS;
+}

lib/uefi/uefi_platform.h

@@ -18,7 +18,9 @@
 #ifndef __GBL_OS_CONFIGURATION_
 #define __GBL_OS_CONFIGURATION_
 
+#include <arch/defines.h>
 #include <uefi/protocols/efi_timestamp_protocol.h>
+#include <uefi/protocols/gbl_efi_image_loading_protocol.h>
 #include <uefi/protocols/gbl_efi_os_configuration_protocol.h>
 #include <uefi/system_table.h>
 #include <uefi/types.h>
@@ -45,4 +47,28 @@ uint64_t get_timestamp();
 
 EfiStatus get_timestamp_properties(EfiTimestampProperties *properties);
 
+// alloc_page/free_pages is implemented in memory_protocols.cpp
+void *alloc_page(void *addr, size_t size, size_t align_log2 = PAGE_SIZE_SHIFT);
+
+void *alloc_page(size_t size, size_t align_log2 = PAGE_SIZE_SHIFT);
+
+EfiStatus free_pages(void *memory, size_t pages);
+
+EfiStatus get_buffer(struct GblEfiImageLoadingProtocol *self,
+                     const GblEfiImageInfo *ImageInfo,
+                     GblEfiImageBuffer *Buffer);
+
+// uefi_malloc is used by LK to allocate memory that would be used by UEFI
+// applications
+void *uefi_malloc(size_t size);
+
+// Used by UEFI application to allocate heap memory.
+EfiStatus allocate_pool(EfiMemoryType pool_type, size_t size, void **buf);
+EfiStatus free_pool(void *mem);
+
+// Called by LK once before executing UEFI application to setup the heap
+void setup_heap();
+// Caled by LK once after executing UEFI application to tear down the heap
+void reset_heap();
+
 #endif