From aa03fc9cb230cf37194babea21d59b25292800b2 Mon Sep 17 00:00:00 2001
From: Rafal Kupiec <belliash@codingworkshop.eu.org>
Date: Sat, 30 Dec 2023 10:54:38 +0100
Subject: [PATCH] Add more checks to BlLoadModule() and make it more verbose

---
 xtldr2/xtldr.c | 62 ++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 52 insertions(+), 10 deletions(-)

diff --git a/xtldr2/xtldr.c b/xtldr2/xtldr.c
index 48c1255..138a5cd 100644
--- a/xtldr2/xtldr.c
+++ b/xtldr2/xtldr.c
@@ -169,7 +169,10 @@ EFI_STATUS
 BlLoadModules(IN PWCHAR ModulesList)
 {
     PWCHAR LastModule, Module;
-    EFI_STATUS Status;
+    EFI_STATUS ReturnStatus, Status;
+
+    /* Set default return value */
+    ReturnStatus = STATUS_EFI_SUCCESS;
 
     if(ModulesList != NULL)
     {
@@ -182,9 +185,9 @@ BlLoadModules(IN PWCHAR ModulesList)
             Status = BlLoadModule(Module);
             if(Status != STATUS_EFI_SUCCESS)
             {
-                /* Failed to load module, print error message and return status code */
+                /* Failed to load module, print error message and set new return value */
                 BlDebugPrint(L"Failed to load module '%S', status = 0x%lx\n", Module, Status);
-                // return Status;
+                ReturnStatus = STATUS_EFI_LOAD_ERROR;
             }
 
             /* Take next module from the list */
@@ -193,14 +196,16 @@ BlLoadModules(IN PWCHAR ModulesList)
     }
 
     /* Return success */
-    return STATUS_EFI_SUCCESS;
+    return ReturnStatus;
 }
 
 XTCDECL
 EFI_STATUS
 BlLoadModule(IN PWCHAR ModuleName)
 {
+    EFI_GUID LIPGuid = EFI_LOADED_IMAGE_PROTOCOL_GUID;
     EFI_MEMMAP_DEVICE_PATH ModuleDevicePath[2];
+    PEFI_LOADED_IMAGE_PROTOCOL LoadedImage;
     PEFI_FILE_HANDLE DirHandle, FsHandle;
     EFI_HANDLE DiskHandle, ModuleHandle;
     PPECOFF_IMAGE_SECTION_HEADER SectionHeader;
@@ -224,7 +229,7 @@ BlLoadModule(IN PWCHAR ModuleName)
         if(RtlCompareWideStringInsensitive(ModuleInfo->ModuleName, ModuleName, 0) == 0)
         {
             /* Module already loaded */
-            BlDebugPrint(L"Module '%S' already loaded\n", ModuleName);
+            BlDebugPrint(L"WARNING: Module '%S' already loaded!\n", ModuleName);
             return STATUS_EFI_SUCCESS;
         }
 
@@ -270,7 +275,7 @@ BlLoadModule(IN PWCHAR ModuleName)
         return Status;
     }
 
-    /* Allocate memory for new option */
+    /* Allocate memory for module information block */
     Status = BlMemoryAllocatePool(sizeof(XTBL_MODULE_INFO), (PVOID*)&ModuleInfo);
     if(Status != STATUS_EFI_SUCCESS)
     {
@@ -278,6 +283,7 @@ BlLoadModule(IN PWCHAR ModuleName)
         return Status;
     }
 
+    /* Zero module information block */
     RtlZeroMemory(ModuleInfo, sizeof(XTBL_MODULE_INFO));
 
     /* Setup PE/COFF EFI image headers */
@@ -286,7 +292,7 @@ BlLoadModule(IN PWCHAR ModuleName)
     SectionHeader = (PPECOFF_IMAGE_SECTION_HEADER)((PUCHAR)&PeHeader->OptionalHeader +
                                                    PeHeader->FileHeader.SizeOfOptionalHeader);
 
-    /* Look for .moddeps section */
+    /* Look for .moddeps and .modinfo sections */
     for(SectionIndex = 0; SectionIndex < PeHeader->FileHeader.NumberOfSections; SectionIndex++)
     {
         /* Check section name */
@@ -305,7 +311,7 @@ BlLoadModule(IN PWCHAR ModuleName)
                 if(Status != STATUS_EFI_SUCCESS)
                 {
                     /* Failed to load module, print error message and return status code */
-                    BlDebugPrint(L"Failed to load dependency module '%S', status = 0x%lx\n", DepsData, Status);
+                    BlDebugPrint(L"Failed to load dependency module '%S', (Status Code: 0x%lx)\n", DepsData, Status);
                     return STATUS_EFI_UNSUPPORTED;
                 }
 
@@ -324,7 +330,7 @@ BlLoadModule(IN PWCHAR ModuleName)
         }
     }
 
-    /* Store module name */
+    /* Finally, store module name */
     ModuleInfo->ModuleName = ModuleName;
 
     /* Setup module device path */
@@ -345,15 +351,51 @@ BlLoadModule(IN PWCHAR ModuleName)
                                                      ModuleData, ModuleSize, &ModuleHandle);
     if(Status != STATUS_EFI_SUCCESS)
     {
-        /* Failed to load module image */
+        /* Check if caused by secure boot */
+        if(Status == STATUS_EFI_ACCESS_DENIED && BlpStatus.SecureBoot >= 1)
+        {
+            /* SecureBoot signature validation failed */
+            BlDebugPrint(L"ERROR: SecureBoot signature validation failed, module '%S' will not be loaded\n", ModuleName);
+        }
+        else
+        {
+            /* Failed to load module */
+            BlDebugPrint(L"ERROR: Unable to load module '%S' (Status Code: 0x%lx)\n", ModuleName, Status);
+        }
+
+        /* Return error status code */
         return Status;
     }
 
+    /* Access module interface for further module type check */
+    Status = EfiSystemTable->BootServices->OpenProtocol(ModuleHandle, &LIPGuid, (PVOID *)&LoadedImage,
+                                                        EfiImageHandle, NULL, EFI_OPEN_PROTOCOL_GET_PROTOCOL);
+    if(Status != STATUS_EFI_SUCCESS)
+    {
+        /* Failed to open LoadedImage protocol */
+        BlDebugPrint(L"ERROR: Unable to access module interface (Status Code: 0x%lx)\n", Status);
+        return Status;
+    }
+
+    /* Some firmwares do not allow to start drivers which are not of 'boot system driver' type, so check it */
+    if(LoadedImage->ImageCodeType != EfiBootServicesCode)
+    {
+        /* Different type set, probably 'runtime driver', refuse to load it */
+        BlDebugPrint(L"ERROR: Loaded module is not a boot system driver\n");
+
+        /* Close protocol and skip module */
+        EfiSystemTable->BootServices->CloseProtocol(LoadedImage, &LIPGuid, LoadedImage, NULL);
+    }
+
+    /* Close loaded image protocol */
+    EfiSystemTable->BootServices->CloseProtocol(LoadedImage, &LIPGuid, LoadedImage, NULL);
+
     /* Start EFI image */
     Status = EfiSystemTable->BootServices->StartImage(ModuleHandle, NULL, NULL);
     if(Status != STATUS_EFI_SUCCESS)
     {
         /* Failed to start module image */
+        BlDebugPrint(L"ERROR: Failed to start module '%S' (Status Code: 0x%lx)\n", ModuleName, Status);
         return Status;
     }