From 670a812649cca7528ab71bc990b6eb335963339e Mon Sep 17 00:00:00 2001
From: Rafal Kupiec <belliash@codingworkshop.eu.org>
Date: Sun, 28 Jan 2024 17:01:44 +0100
Subject: [PATCH] Store number of mappings when mapping virtual memory

---
 sdk/xtdk/bltypes.h | 1 +
 xtldr/memory.c     | 4 +++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/sdk/xtdk/bltypes.h b/sdk/xtdk/bltypes.h
index 0b14b24..e667d28 100644
--- a/sdk/xtdk/bltypes.h
+++ b/sdk/xtdk/bltypes.h
@@ -189,6 +189,7 @@ typedef struct _XTBL_MODULE_INFO
 typedef struct _XTBL_PAGE_MAPPING
 {
     LIST_ENTRY MemoryMap;
+    ULONG MapSize;
     PVOID PtePointer;
     SHORT PageMapLevel;
     PAGE_SIZE PageSize;
diff --git a/xtldr/memory.c b/xtldr/memory.c
index 8ae89d1..dc90a9c 100644
--- a/xtldr/memory.c
+++ b/xtldr/memory.c
@@ -178,6 +178,7 @@ BlInitializePageMap(OUT PXTBL_PAGE_MAPPING PageMap,
 {
     /* Initialize memory mappings */
     RtlInitializeListHead(&PageMap->MemoryMap);
+    PageMap->MapSize = 0;
 
     /* Set page map size/level and memory map address */
     PageMap->PageMapLevel = PageMapLevel;
@@ -457,8 +458,9 @@ BlMapVirtualMemory(IN OUT PXTBL_PAGE_MAPPING PageMap,
         ListEntry = ListEntry->Flink;
     }
 
-    /* Insert new mapping to the list */
+    /* Insert new mapping to the list and increase page map size */
     RtlInsertTailList(&PageMap->MemoryMap, &Mapping1->ListEntry);
+    PageMap->MapSize++;
 
     /* Return success */
     return STATUS_EFI_SUCCESS;