[efi] Separate GetMemoryMap() wrapper from shim unlocker

Signed-off-by: Michael Brown <mcb30@ipxe.org>

diff --git a/src/interface/efi/efi_shim.c b/src/interface/efi/efi_shim.c
index 9b1b69e8071666927e6ebc28325b316f15b4fa40..14d4806f23cd24258f0f680a904d236268e44fc1 100644 (file)
--- a/src/interface/efi/efi_shim.c
+++ b/src/interface/efi/efi_shim.c
@@ -90,23 +90,13 @@ struct image_tag efi_shim __image_tag = {
 };
 
 /** Original GetMemoryMap() function */
-static EFI_GET_MEMORY_MAP efi_shim_orig_map;
+static EFI_GET_MEMORY_MAP efi_shim_orig_get_memory_map;
 
 /**
  * Unlock UEFI shim
  *
- * @v len              Memory map size
- * @v map              Memory map
- * @v key              Memory map key
- * @v desclen          Descriptor size
- * @v descver          Descriptor version
- * @ret efirc          EFI status code
- *
  */
-static EFIAPI EFI_STATUS efi_shim_unlock ( UINTN *len,
-                                          EFI_MEMORY_DESCRIPTOR *map,
-                                          UINTN *key, UINTN *desclen,
-                                          UINT32 *descver ) {
+static void efi_shim_unlock ( void ) {
        EFI_BOOT_SERVICES *bs = efi_systab->BootServices;
        uint8_t empty[0];
        union {
@@ -121,9 +111,30 @@ static EFIAPI EFI_STATUS efi_shim_unlock ( UINTN *len,
                u.lock->Verify ( empty, sizeof ( empty ) );
                DBGC ( &efi_shim, "SHIM unlocked via %p\n", u.lock );
        }
+}
+
+/**
+ * Wrap GetMemoryMap()
+ *
+ * @v len              Memory map size
+ * @v map              Memory map
+ * @v key              Memory map key
+ * @v desclen          Descriptor size
+ * @v descver          Descriptor version
+ * @ret efirc          EFI status code
+ */
+static EFIAPI EFI_STATUS efi_shim_get_memory_map ( UINTN *len,
+                                                  EFI_MEMORY_DESCRIPTOR *map,
+                                                  UINTN *key, UINTN *desclen,
+                                                  UINT32 *descver ) {
+
+       /* Unlock shim */
+       if ( ! efi_shim_require_loader )
+               efi_shim_unlock();
 
        /* Hand off to original GetMemoryMap() */
-       return efi_shim_orig_map ( len, map, key, desclen, descver );
+       return efi_shim_orig_get_memory_map ( len, map, key, desclen,
+                                             descver );
 }
 
 /**
@@ -216,27 +227,23 @@ int efi_shim_install ( struct image *shim, EFI_HANDLE handle,
        EFI_BOOT_SERVICES *bs = efi_systab->BootServices;
        int rc;
 
-       /* Intercept GetMemoryMap() via boot services table */
-       efi_shim_orig_map = bs->GetMemoryMap;
-       if ( ! efi_shim_require_loader )
-               bs->GetMemoryMap = efi_shim_unlock;
-
        /* Stop PXE base code */
        if ( ( ! efi_shim_allow_pxe ) &&
             ( ( rc = efi_shim_inhibit_pxe ( handle ) ) != 0 ) ) {
-               goto err_inhibit_pxe;
+               return rc;
        }
 
        /* Update command line */
        if ( ( rc = efi_shim_cmdline ( shim, cmdline ) ) != 0 )
-               goto err_cmdline;
+               return rc;
 
-       return 0;
+       /* Record original boot services functions */
+       efi_shim_orig_get_memory_map = bs->GetMemoryMap;
 
- err_cmdline:
- err_inhibit_pxe:
-       bs->GetMemoryMap = efi_shim_orig_map;
-       return rc;
+       /* Wrap relevant boot services functions */
+       bs->GetMemoryMap = efi_shim_get_memory_map;
+
+       return 0;
 }
 
 /**
@@ -246,6 +253,6 @@ int efi_shim_install ( struct image *shim, EFI_HANDLE handle,
 void efi_shim_uninstall ( void ) {
        EFI_BOOT_SERVICES *bs = efi_systab->BootServices;
 
-       /* Restore original GetMemoryMap() */
-       bs->GetMemoryMap = efi_shim_orig_map;
+       /* Restore original boot services functions */
+       bs->GetMemoryMap = efi_shim_orig_get_memory_map;
 }