---
Some implementations of uefi don't seem to provide a memory key when
not called with a proper buffer, so I have reused the existing memory
map buffer since that does not seem to be in use at any point after
exiting boot services.

 efi/mmap.ha   |  6 +++---
 efi/systab.ha | 16 +++++++++++++---
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/efi/mmap.ha b/efi/mmap.ha
index f526dd6..e75fd21 100644
--- a/efi/mmap.ha
+++ b/efi/mmap.ha
@@ -33,6 +33,7 @@ export type MEMORY_DESCRIPTOR = struct {
 // in practice we may have less than 128 items here.
 def MMAP_STATIC: size = 128 * size(MEMORY_DESCRIPTOR);
 let mmap_buf: [MMAP_STATIC]u8 = [0...];
+let mmap_key: u64 = 0;
 
 export type mmap_iter = struct {
 	i: size,
@@ -43,11 +44,10 @@ export type mmap_iter = struct {
 // Iterates over the EFI memory map.
 export fn iter_mmap() (mmap_iter | error) = {
 	let msize = 0z;
-	let mkey = 0u64;
 	let dsize = 0z;
 	let dver = 0u32;
 	const err = systab.BootServices.GetMemoryMap(&msize,
-		null, &mkey, &dsize, &dver);
+		null, &mmap_key, &dsize, &dver);
 	assert(err == STATUS::BUFFER_TOO_SMALL);
 
 	if (msize > len(mmap_buf)) {
@@ -56,7 +56,7 @@ export fn iter_mmap() (mmap_iter | error) = {
 	assert(dver == 1);
 
 	efi_call(systab.BootServices.GetMemoryMap(&msize,
-			&mmap_buf[0], &mkey, &dsize, &dver))?;
+			&mmap_buf[0], &mmap_key, &dsize, &dver))?;
 	return mmap_iter {
 		i = 0,
 		n = msize / dsize,
diff --git a/efi/systab.ha b/efi/systab.ha
index 3bf6ff2..a036e37 100644
--- a/efi/systab.ha
+++ b/efi/systab.ha
@@ -110,12 +110,22 @@ export type CONFIGURATION_TABLE = struct {
 
 // Instructs the EFI boot services to exit.
 export fn exit_boot_services() void = {
-	let map_key = 0u64;
 	let msize = 0z, dsize = 0z, dver = 0u32;
+
+	const err = systab.BootServices.ExitBootServices(image, mmap_key);
+
+	if(err == STATUS::SUCCESS){
+		return;
+	};
+
 	const err = systab.BootServices.GetMemoryMap(
-		&msize, null, &map_key, &dsize, &dver);
+		&msize, null, &mmap_key, &dsize, &dver);
+
 	assert(err == STATUS::BUFFER_TOO_SMALL);
 
-	const err = systab.BootServices.ExitBootServices(image, map_key);
+	const err = systab.BootServices.GetMemoryMap(
+		&msize, &mmap_buf[0], &mmap_key, &dsize, &dver);
+
+	const err = systab.BootServices.ExitBootServices(image, mmap_key);
 	assert(err == STATUS::SUCCESS);
 };
-- 
2.40.0