Retrieve position (address) where the efi image is loaded

1 files changed, 63 insertions, 36 deletions

diff --git a/efi_memtest/MemtestEfi.c b/efi_memtest/MemtestEfi.c
index 3f14289..87ea80b 100644
--- a/efi_memtest/MemtestEfi.c
+++ b/efi_memtest/MemtestEfi.c
@@ -3,6 +3,7 @@
 #include "Library/UefiLib.h"
 #include "Library/UefiApplicationEntryPoint.h"
 #include "Guid/FileSystemInfo.h"
+#include "Protocol/LoadedImage.h"
 
 #include "main.h" // TODO move into main dir
 #include "memtest86+/logger.h"
@@ -15,51 +16,77 @@ extern EFI_SYSTEM_TABLE *gST;
 EFI_STATUS
 EFIAPI
 UefiMain (
-  IN EFI_HANDLE        ImageHandle,
-  IN EFI_SYSTEM_TABLE  *SystemTable
-  )
+    IN EFI_HANDLE        ImageHandle,
+    IN EFI_SYSTEM_TABLE  *SystemTable
+    )
 {
 
-  SystemTable->ConOut->ClearScreen(SystemTable->ConOut);
-  SystemTable->ConOut->SetCursorPosition(SystemTable->ConOut, 0, 0);
+    SystemTable->ConOut->ClearScreen(SystemTable->ConOut);
+    SystemTable->ConOut->SetCursorPosition(SystemTable->ConOut, 0, 0);
 
-  trace_init("Starting MemtestEfi");
+    trace_init("Starting MemtestEfi");
 
-  init_logger();
+    init_logger();
+    
+    if (logflag) {
+        char msg[] = "MemtestEfi started";
+        print_log(msg, sizeof(msg) - 1);
+    }
 
-  //while(1);
-
-
-  if (logflag) {
-  char msg[] = "MemtestEfi started";
-  print_log(msg, sizeof(msg) - 1);
-  }
-
-  {
-    char log[43] = "Address of UefiMain: ";
-    int length = 23;
-    int_to_charr((unsigned long)UefiMain, log, &length);
-    print_log(log, length);
-  }
-
-  {
-    char log[43] = "Address of test_start: ";
-    int length = 23;
-    int_to_charr((unsigned long)test_start, log, &length);
-    print_log(log, length);
-  }
   
-  test_start();
+    {
+        char log[43] = "Address of UefiMain: ";
+        int length = 23;
+        int_to_charr((unsigned long)UefiMain, log, &length);
+        print_log(log, length);
+    }
+
+    {
+        char log[43] = "Address of test_start: ";
+        int length = 23;
+        int_to_charr((unsigned long)test_start, log, &length);
+        print_log(log, length);
+    }
+
+    VOID *LoadedImage;
+    EFI_STATUS status = SystemTable->BootServices->OpenProtocol(
+        ImageHandle,
+        &gEfiLoadedImageProtocolGuid,
+        &LoadedImage,
+        ImageHandle,
+        NULL,
+        EFI_OPEN_PROTOCOL_GET_PROTOCOL);
+    if (status != EFI_SUCCESS) {
+        print_log("Failed to open LoadedImage.", 27);
+    } else {
+        VOID *base = ((EFI_LOADED_IMAGE_PROTOCOL *) LoadedImage)->ImageBase;
+        UINT64 size = ((EFI_LOADED_IMAGE_PROTOCOL *) LoadedImage)->ImageSize;
+        {
+            char log[31] = "ImageBase: ";
+            int length = 11;
+            int_to_charr((unsigned long) base, log, &length);
+            print_log(log, length);
+        }
+
+        {
+            char log[31] = "ImageSize: ";
+            int length = 11;
+            int_to_charr(size, log, &length);
+            print_log(log, length);
+        }
+    }
+  
+    test_start();
 
-  if (logflag) {
-  char msg[] = "MemtestEfi finished!\n";
-  print_log(msg, sizeof(msg) - 1);
-  }
+    if (logflag) {
+        char msg[] = "MemtestEfi finished!\n";
+        print_log(msg, sizeof(msg) - 1);
+    }
 
-  //while(1) {}
+    //while(1) {}
 
-  SystemTable->ConOut->SetCursorPosition(SystemTable->ConOut, 0, 25);
+    SystemTable->ConOut->SetCursorPosition(SystemTable->ConOut, 0, 25);
 
-  return EFI_SUCCESS;
+    return EFI_SUCCESS;
 }