[lib][uefi] Fix memory leak in load_pe_file

.clang-format

@@ -11,4 +11,4 @@ ColumnLimit: 0
 IndentCaseLabels: true
 IndentGotoLabels: false
 AlignConsecutiveMacros:
-  Enabled: true
+  Enabled: true

lib/uefi/boot_service_provider.cpp

@@ -54,8 +54,11 @@ EfiStatus handle_protocol(EfiHandle handle, const EfiGuid *protocol,
   if (guid_eq(protocol, LOADED_IMAGE_PROTOCOL_GUID)) {
     printf("handle_protocol(%p, LOADED_IMAGE_PROTOCOL_GUID, %p);\n", handle,
            intf);
-    const auto loaded_image = static_cast<EfiLoadedImageProtocol*>(
-        uefi_malloc(sizeof(EfiLoadedImageProtocol)));
+    const auto loaded_image = static_cast<EFI_LOADED_IMAGE_PROTOCOL *>(
+        uefi_malloc(sizeof(EFI_LOADED_IMAGE_PROTOCOL)));
+    if (!loaded_image) {
+      return EFI_STATUS_OUT_OF_RESOURCES;
+    }
     *loaded_image = {};
     loaded_image->revision = EFI_LOADED_IMAGE_PROTOCOL_REVISION;
     loaded_image->parent_handle = nullptr;
@@ -70,7 +73,7 @@ EfiStatus handle_protocol(EfiHandle handle, const EfiGuid *protocol,
   } else if (guid_eq(protocol, LINUX_EFI_LOADED_IMAGE_FIXED_GUID)) {
     printf("handle_protocol(%p, LINUX_EFI_LOADED_IMAGE_FIXED_GUID, %p);\n",
            handle, intf);
-    return EFI_STATUS_SUCCESS;
+    return EFI_STATUS_UNSUPPORTED;
   } else {
     printf("handle_protocol(%p, %p, %p);\n", handle, protocol, intf);
   }

lib/uefi/charset.cpp

@@ -46,14 +46,14 @@ int utf16_strcmp(const char16_t *s1, const char16_t *s2) {
  * This function converts utf-8 string to utf-16 string. However
  * it only supports us-ascii input right now.
  */
-int utf8_to_utf16(char16_t *dest, const char *src, size_t size) {
+size_t utf8_to_utf16(char16_t *dest, const char *src, size_t size) {
   size_t i = 0;
   for (; i < size - 1 && *src; i++, src++) {
     dest[i] = *src;
   }
 
   dest[i] = 0;
-  return 0;
+  return i;
 }
 
 /**
@@ -62,12 +62,12 @@ int utf8_to_utf16(char16_t *dest, const char *src, size_t size) {
  * This function converts utf-16 string to utf-8 string. However
  * it only supports us-ascii output right now.
  */
-int utf16_to_utf8(char *dest, const char16_t *src, size_t size) {
+size_t utf16_to_utf8(char *dest, const char16_t *src, size_t size) {
   size_t i = 0;
   for (; i < size - 1 && *src; i++, src++) {
     dest[i] = *src;
   }
 
   dest[i] = 0;
-  return 0;
+  return i;
 }

lib/uefi/charset.h

@@ -22,7 +22,11 @@
 
 size_t utf16_strlen(const char16_t *str);
 int utf16_strcmp(const char16_t *s1, const char16_t *s2);
-int utf8_to_utf16(char16_t *dest, const char *src, size_t size);
-int utf16_to_utf8(char *dest, const char16_t *src, size_t size);
+static inline int utf16_strcmp(const uint16_t *s1, const char16_t *s2) {
+  return utf16_strcmp(reinterpret_cast<const char16_t *>(s1), s2);
+}
+// Return number of characters converted, including the null terminator
+size_t utf8_to_utf16(char16_t *dest, const char *src, size_t size);
+size_t utf16_to_utf8(char *dest, const char16_t *src, size_t size);
 
 #endif

lib/uefi/include/uefi/system_table.h

@@ -38,7 +38,7 @@ typedef struct {
 
 typedef struct EfiSystemTable {
   EfiTableHeader header;
-  uint16_t* firmware_vendor;
+  const char16_t* firmware_vendor;
   uint32_t firmware_revision;
   EfiHandle console_in_handle;
   void* con_in;

lib/uefi/pe.h

@@ -63,6 +63,7 @@ struct IMAGE_DOS_HEADER { // DOS .EXE header
   u16 e_res2[10];         // Reserved words
   u32 e_lfanew;           // File address of new exe header
 
+  // ASCII "PE\x0\x0"
   constexpr bool CheckMagic() const { return LE32(e_magic) == 0x5A4D; }
   IMAGE_NT_HEADERS64 *GetPEHeader() {
     auto address = reinterpret_cast<char *>(this);
@@ -284,6 +285,7 @@ static constexpr size_t EFI_IMAGE_REL_BASED_HIGHADJ = 4;
 static constexpr size_t EFI_IMAGE_REL_BASED_MIPS_JMPADDR = 5;
 static constexpr size_t EFI_IMAGE_REL_BASED_ARM_MOV32A = 5;
 static constexpr size_t EFI_IMAGE_REL_BASED_ARM_MOV32T = 7;
+static constexpr size_t EFI_IMAGE_REL_BASED_LOONGARCH64_MARK_LA = 8;
 static constexpr size_t EFI_IMAGE_REL_BASED_IA64_IMM64 = 9;
 static constexpr size_t EFI_IMAGE_REL_BASED_MIPS_JMPADDR16 = 9;
 static constexpr size_t EFI_IMAGE_REL_BASED_DIR64 = 10;

lib/uefi/relocation.cpp

@@ -7,102 +7,6 @@
 
 #include "pe.h"
 
-namespace {
-
-constexpr size_t BIT26 = 1 << 26;
-constexpr size_t BIT11 = 1 << 11;
-constexpr size_t BIT10 = 1 << 10;
-
-/**
-  Pass in a pointer to an ARM MOVT or MOVW immediate instruciton and
-  return the immediate data encoded in the instruction.
-
-  @param  Instruction   Pointer to ARM MOVT or MOVW immediate instruction
-
-  @return Immediate address encoded in the instruction
-
-**/
-uint16_t ThumbMovtImmediateAddress(const uint16_t *Instruction) {
-  uint32_t Movt{};
-  uint16_t Address{};
-
-  // Thumb2 is two 16-bit instructions working together. Not a single 32-bit
-  // instruction Example MOVT R0, #0 is 0x0000f2c0 or 0xf2c0 0x0000
-  Movt = (*Instruction << 16) | (*(Instruction + 1));
-
-  // imm16 = imm4:i:imm3:imm8
-  //         imm4 -> Bit19:Bit16
-  //         i    -> Bit26
-  //         imm3 -> Bit14:Bit12
-  //         imm8 -> Bit7:Bit0
-  Address = static_cast<uint16_t>(Movt & 0x000000ff);         // imm8
-  Address |= static_cast<uint16_t>((Movt >> 4) & 0x0000f700); // imm4 imm3
-  Address |= (((Movt & BIT26) != 0) ? BIT11 : 0);             // i
-  return Address;
-}
-
-/**
-  Pass in a pointer to an ARM MOVW/MOVT instruciton pair and
-  return the immediate data encoded in the two` instruction.
-
-  @param  Instructions  Pointer to ARM MOVW/MOVT insturction pair
-
-  @return Immediate address encoded in the instructions
-
-**/
-uint32_t ThumbMovwMovtImmediateAddress(uint16_t *Instructions) {
-  uint16_t *Word{};
-  uint16_t *Top{};
-
-  Word = Instructions; // MOVW
-  Top = Word + 2;      // MOVT
-
-  return (ThumbMovtImmediateAddress(Top) << 16) +
-         ThumbMovtImmediateAddress(Word);
-}
-
-/**
-  Update an ARM MOVT or MOVW immediate instruction immediate data.
-
-  @param  Instruction   Pointer to ARM MOVT or MOVW immediate instruction
-  @param  Address       New addres to patch into the instruction
-**/
-void ThumbMovtImmediatePatch(uint16_t *Instruction, uint16_t Address) {
-  uint16_t Patch{};
-
-  // First 16-bit chunk of instruciton
-  Patch = ((Address >> 12) & 0x000f);              // imm4
-  Patch |= (((Address & BIT11) != 0) ? BIT10 : 0); // i
-  // Mask out instruction bits and or in address
-  *(Instruction) = (*Instruction & ~0x040f) | Patch;
-
-  // Second 16-bit chunk of instruction
-  Patch = Address & 0x000000ff;           // imm8
-  Patch |= ((Address << 4) & 0x00007000); // imm3
-  // Mask out instruction bits and or in address
-  Instruction++;
-  *Instruction = (*Instruction & ~0x70ff) | Patch;
-}
-
-/**
-  Update an ARM MOVW/MOVT immediate instruction instruction pair.
-
-  @param  Instructions  Pointer to ARM MOVW/MOVT instruction pair
-  @param  Address       New addres to patch into the instructions
-**/
-void ThumbMovwMovtImmediatePatch(uint16_t *Instructions, uint32_t Address) {
-  uint16_t *Word{};
-  uint16_t *Top{};
-
-  Word = Instructions; // MOVW
-  Top = Word + 2;      // MOVT
-
-  ThumbMovtImmediatePatch(Word, static_cast<uint16_t>(Address & 0xffff));
-  ThumbMovtImmediatePatch(Top, static_cast<uint16_t>(Address >> 16));
-}
-
-}  // namespace
-
 int relocate_image(char *image) {
   const auto dos_header = reinterpret_cast<IMAGE_DOS_HEADER *>(image);
   const auto pe_header = dos_header->GetPEHeader();
@@ -141,7 +45,6 @@ int relocate_image(char *image) {
       auto Fixup16 = reinterpret_cast<uint16_t *>(Fixup);
       auto Fixup32 = reinterpret_cast<uint32_t *>(Fixup);
       auto Fixup64 = reinterpret_cast<uint64_t *>(Fixup);
-      uint32_t FixupVal = 0;
       switch ((*Reloc) >> 12) {
       case EFI_IMAGE_REL_BASED_ABSOLUTE:
         break;
@@ -167,18 +70,32 @@ int relocate_image(char *image) {
         *Fixup64 = *Fixup64 + static_cast<uint64_t>(Adjust);
         break;
 
-      case EFI_IMAGE_REL_BASED_ARM_MOV32T:
-        FixupVal = ThumbMovwMovtImmediateAddress(Fixup16) +
-                   static_cast<uint32_t>(Adjust);
-        ThumbMovwMovtImmediatePatch(Fixup16, FixupVal);
-
-        break;
-
       case EFI_IMAGE_REL_BASED_ARM_MOV32A:
         printf("Unsupported relocation type: EFI_IMAGE_REL_BASED_ARM_MOV32A\n");
         // break omitted - ARM instruction encoding not implemented
         break;
-
+      case EFI_IMAGE_REL_BASED_LOONGARCH64_MARK_LA: {
+        // The next four instructions are used to load a 64 bit address,
+        // relocate all of them
+        uint64_t Value =
+            (*Fixup32 & 0x1ffffe0) << 7 |           // lu12i.w 20bits from bit5
+            (*(Fixup32 + 1) & 0x3ffc00) >> 10;      // ori     12bits from bit10
+        uint64_t Tmp1 = *(Fixup32 + 2) & 0x1ffffe0; // lu32i.d 20bits from bit5
+        uint64_t Tmp2 = *(Fixup32 + 3) & 0x3ffc00;  // lu52i.d 12bits from bit10
+        Value = Value | (Tmp1 << 27) | (Tmp2 << 42);
+        Value += Adjust;
+
+        *Fixup32 = (*Fixup32 & ~0x1ffffe0) | (((Value >> 12) & 0xfffff) << 5);
+
+        Fixup += sizeof(uint32_t);
+        *Fixup32 = (*Fixup32 & ~0x3ffc00) | ((Value & 0xfff) << 10);
+
+        Fixup += sizeof(uint32_t);
+        *Fixup32 = (*Fixup32 & ~0x1ffffe0) | (((Value >> 32) & 0xfffff) << 5);
+
+        Fixup += sizeof(uint32_t);
+        *Fixup32 = (*Fixup32 & ~0x3ffc00) | (((Value >> 52) & 0xfff) << 10);
+      }
       default:
         printf("Unsupported relocation type: %d\n", (*Reloc) >> 12);
         return -1;

lib/uefi/runtime_service_provider.cpp

@@ -22,29 +22,20 @@
 #include <string.h>
 #include <uefi/types.h>
 
+#include "charset.h"
 #include "variable_mem.h"
 
 namespace {
 
-constexpr auto &&kSecureBoot = "SecureBoot";
+EFI_STATUS GetVariable(const uint16_t *VariableName, const EfiGuid *VendorGuid,
+                       uint32_t *Attributes, size_t *DataSize, void *Data) {
 
-EfiStatus GetVariable(const uint16_t* VariableName, const EfiGuid* VendorGuid,
-                      uint32_t* Attributes, size_t* DataSize, void* Data) {
   if (!VariableName || !VendorGuid || !DataSize) {
     return EFI_STATUS_INVALID_PARAMETER;
   }
 
-  char buffer[512];
-  size_t i = 0;
-  while (VariableName[i] && i < sizeof(buffer)) {
-    size_t j = 0;
-    for (j = 0; j < sizeof(buffer) - 1 && VariableName[i + j]; j++) {
-      buffer[j] = VariableName[i + j];
-    }
-    i += j;
-  }
-  buffer[i] = 0;
-  if (strncmp(buffer, kSecureBoot, sizeof(kSecureBoot)) == 0 || strcmp(buffer, "SetupMode") == 0) {
+  if (utf16_strcmp(VariableName, u"SecureBoot") == 0 &&
+      utf16_strcmp(VariableName, u"SetupMode") == 0) {
     if (DataSize) {
       *DataSize = 1;
     }
@@ -56,7 +47,9 @@ EfiStatus GetVariable(const uint16_t* VariableName, const EfiGuid* VendorGuid,
 
   char *data_in_mem;
   size_t data_in_mem_size;
-  if (efi_get_variable(reinterpret_cast<const char16_t *>(VariableName), VendorGuid, Attributes, &data_in_mem, &data_in_mem_size) == EFI_STATUS_SUCCESS) {
+  if (efi_get_variable(reinterpret_cast<const char16_t *>(VariableName),
+                       VendorGuid, Attributes, &data_in_mem,
+                       &data_in_mem_size) == EFI_STATUS_SUCCESS) {
     if (*DataSize == 0 && !Data) {
       *DataSize = data_in_mem_size;
       return EFI_STATUS_BUFFER_TOO_SMALL;
@@ -68,6 +61,9 @@ EfiStatus GetVariable(const uint16_t* VariableName, const EfiGuid* VendorGuid,
     memcpy(Data, data_in_mem, data_in_mem_size);
     return EFI_STATUS_SUCCESS;
   }
+  char buffer[512];
+  utf16_to_utf8(buffer, reinterpret_cast<const char16_t *>(VariableName),
+                sizeof(buffer));
 
   printf("%s(%s) is unsupported\n", __FUNCTION__, buffer);
   return EFI_STATUS_UNSUPPORTED;
@@ -89,25 +85,22 @@ EfiStatus SetVariable(const uint16_t* VariableName, const EfiGuid* VendorGuid,
   /* Only allow setting non-volatile variables */
   if ((Attributes & EFI_VARIABLE_NON_VOLATILE) == 0) {
     efi_set_variable(reinterpret_cast<const char16_t *>(VariableName),
-                     VendorGuid,
-                     Attributes,
-                     reinterpret_cast<const char *>(Data),
-                     DataSize);
+                     VendorGuid, Attributes,
+                     reinterpret_cast<const char *>(Data), DataSize);
     return EFI_STATUS_SUCCESS;
   }
 
   printf("%s: Only non-volatile is supported. Attributes = 0x%x\n",
-         __FUNCTION__,
-         Attributes);
+         __FUNCTION__, Attributes);
   return EFI_STATUS_UNSUPPORTED;
 }
 
-void ResetSystem(EfiResetType ResetType, EfiStatus ResetStatus, size_t DataSize,
-                 void* ResetData) {
+void ResetSystem(EFI_RESET_TYPE ResetType, EFI_STATUS ResetStatus,
+                 size_t DataSize, void *ResetData) {
   platform_halt(HALT_ACTION_REBOOT, HALT_REASON_SW_RESET);
 }
 
-}  // namespace
+} // namespace
 
 void setup_runtime_service_table(EfiRuntimeService *service) {
   service->get_variable = GetVariable;