Get current return address in UefiMain

1 files changed, 21 insertions, 7 deletions

diff --git a/efi_memtest/MemtestEfi.c b/efi_memtest/MemtestEfi.c
index 87ea80b..876d825 100644
--- a/efi_memtest/MemtestEfi.c
+++ b/efi_memtest/MemtestEfi.c
@@ -34,14 +34,14 @@ UefiMain (
     }
 
   
-    {
-        char log[43] = "Address of UefiMain: ";
-        int length = 23;
+    if (logflag) {
+        char log[41] = "Address of UefiMain: ";
+        int length = 21;
         int_to_charr((unsigned long)UefiMain, log, &length);
         print_log(log, length);
     }
 
-    {
+    if (logflag) {
         char log[43] = "Address of test_start: ";
         int length = 23;
         int_to_charr((unsigned long)test_start, log, &length);
@@ -57,24 +57,38 @@ UefiMain (
         NULL,
         EFI_OPEN_PROTOCOL_GET_PROTOCOL);
     if (status != EFI_SUCCESS) {
-        print_log("Failed to open LoadedImage.", 27);
+        if (logflag) print_log("Failed to open LoadedImage.", 27);
     } else {
         VOID *base = ((EFI_LOADED_IMAGE_PROTOCOL *) LoadedImage)->ImageBase;
         UINT64 size = ((EFI_LOADED_IMAGE_PROTOCOL *) LoadedImage)->ImageSize;
-        {
+        if (logflag) {
             char log[31] = "ImageBase: ";
             int length = 11;
             int_to_charr((unsigned long) base, log, &length);
             print_log(log, length);
         }
 
-        {
+        if (logflag) {
             char log[31] = "ImageSize: ";
             int length = 11;
             int_to_charr(size, log, &length);
             print_log(log, length);
         }
     }
+
+    if (logflag) {
+        char log[45] = "Address of SystemTable = ";
+        int length = 25;
+        int_to_charr((unsigned long) SystemTable, log, &length);
+        print_log(log, length);
+    }
+
+    if (logflag) {
+        char log[45] = "Current return address = ";
+        int length = 25;
+        int_to_charr((unsigned long) ((unsigned long *) __builtin_return_address(0)), log, &length);
+        print_log(log, length);
+    }
   
     test_start();