From 7ba599ad0f8100448613bbf0738895cd96243737 Mon Sep 17 00:00:00 2001
From: bzt <bztemail@gmail.com>
Date: Mon, 22 Nov 2021 17:50:32 +0100
Subject: [PATCH] Added more failsafes, fixed issue #19

---
 uefi/crt_aarch64.c | 13 +++++++++----
 uefi/crt_x86_64.c  | 14 +++++++++-----
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/uefi/crt_aarch64.c b/uefi/crt_aarch64.c
index f7ed6bf..23b31a8 100644
--- a/uefi/crt_aarch64.c
+++ b/uefi/crt_aarch64.c
@@ -144,7 +144,7 @@ void bootstrap()
 /**
  * Initialize POSIX-UEFI and call the application's main() function
  */
-int uefi_init (
+efi_status_t uefi_init (
     efi_handle_t image, efi_system_table_t *systab
 #ifndef __clang__
     , uintptr_t ldbase, Elf64_Dyn *dyn
@@ -156,10 +156,10 @@ int uefi_init (
     efi_shell_interface_protocol_t *shi = NULL;
     efi_guid_t lipGuid = EFI_LOADED_IMAGE_PROTOCOL_GUID;
     efi_status_t status;
-    int argc = 0, i;
+    int argc = 0, i, ret;
     wchar_t **argv = NULL;
 #if USE_UTF8
-    int ret, j;
+    int j;
     char *s;
 #endif
 #ifndef __clang__
@@ -186,6 +186,10 @@ int uefi_init (
 #else
     (void)i;
 #endif
+    /* failsafes, should never happen */
+    if(!image || !systab || !systab->BootServices || !systab->BootServices->HandleProtocol ||
+        !systab->BootServices->OpenProtocol || !systab->BootServices->AllocatePool || !systab->BootServices->FreePool)
+            return EFI_UNSUPPORTED;
     /* save EFI pointers and loaded image into globals */
     IM = image;
     ST = systab;
@@ -227,6 +231,7 @@ int uefi_init (
     if(__argvutf8) BS->FreePool(__argvutf8);
     return ret;
 #else
-    return main(argc, argv);
+    ret = main(argc, argv);
 #endif
+    return ret ? EFIERR(ret) : EFI_SUCCESS;
 }
diff --git a/uefi/crt_x86_64.c b/uefi/crt_x86_64.c
index 2597adf..a084103 100644
--- a/uefi/crt_x86_64.c
+++ b/uefi/crt_x86_64.c
@@ -137,7 +137,7 @@ void bootstrap()
 /**
  * Initialize POSIX-UEFI and call the application's main() function
  */
-int uefi_init (
+efi_status_t uefi_init (
 #ifndef __clang__
     uintptr_t ldbase, Elf64_Dyn *dyn, efi_system_table_t *systab, efi_handle_t image
 #else
@@ -150,10 +150,10 @@ int uefi_init (
     efi_shell_interface_protocol_t *shi = NULL;
     efi_guid_t lipGuid = EFI_LOADED_IMAGE_PROTOCOL_GUID;
     efi_status_t status;
-    int argc = 0, i;
+    int argc = 0, i, ret;
     wchar_t **argv = NULL;
 #if USE_UTF8
-    int ret, j;
+    int j;
     char *s;
 #endif
 #ifndef __clang__
@@ -189,6 +189,10 @@ int uefi_init (
     "	orw $3 << 9, %ax\n"
     "	mov %rax, %cr4\n"
     );
+    /* failsafes, should never happen */
+    if(!image || !systab || !systab->BootServices || !systab->BootServices->HandleProtocol ||
+        !systab->BootServices->OpenProtocol || !systab->BootServices->AllocatePool || !systab->BootServices->FreePool)
+            return EFI_UNSUPPORTED;
     /* save EFI pointers and loaded image into globals */
     IM = image;
     ST = systab;
@@ -228,8 +232,8 @@ int uefi_init (
     }
     ret = main(argc, (char**)__argvutf8);
     if(__argvutf8) BS->FreePool(__argvutf8);
-    return ret;
 #else
-    return main(argc, argv);
+    ret = main(argc, argv);
 #endif
+    return ret ? EFIERR(ret) : EFI_SUCCESS;
 }