From 65e33fdad5bf9abf7ed0b6cfa4f0fb7a584e0640 Mon Sep 17 00:00:00 2001
From: Kaimakan71 <undefined.foss@gmail.com>
Date: Sun, 1 Sep 2024 16:49:19 -0400
Subject: [PATCH] [BOOT:MM] Finish MmMdRemoveRegionFromMdlEx()

---
 BOOT/ENVIRON/LIB/EFI/efimm.c |   2 +-
 BOOT/ENVIRON/LIB/MM/mmmd.c   | 151 ++++++++++++++++++++++++++++++-----
 2 files changed, 130 insertions(+), 23 deletions(-)

diff --git a/BOOT/ENVIRON/LIB/EFI/efimm.c b/BOOT/ENVIRON/LIB/EFI/efimm.c
index 3bec1aa..f0f3647 100644
--- a/BOOT/ENVIRON/LIB/EFI/efimm.c
+++ b/BOOT/ENVIRON/LIB/EFI/efimm.c
@@ -611,7 +611,7 @@ Return Value:
     //
     // Remove the current descriptor.
     //
-    Status = MmMdRemoveRegionFromMdlEx(Mdl, NtStartPage, NtPageCount, MDL_OPERATION_FLAGS_PHYSICAL, 0);
+    Status = MmMdRemoveRegionFromMdlEx(Mdl, NtStartPage, NtPageCount, MDL_OPERATION_FLAGS_PHYSICAL, NULL);
     if (!NT_SUCCESS(Status)) {
         MmMdFreeDescriptor(NtDescriptor);
         goto exit;
diff --git a/BOOT/ENVIRON/LIB/MM/mmmd.c b/BOOT/ENVIRON/LIB/MM/mmmd.c
index 492eeec..70c1b85 100644
--- a/BOOT/ENVIRON/LIB/MM/mmmd.c
+++ b/BOOT/ENVIRON/LIB/MM/mmmd.c
@@ -515,7 +515,7 @@ Return Value:
 NTSTATUS
 MmMdRemoveRegionFromMdlEx (
     IN PMEMORY_DESCRIPTOR_LIST Mdl,
-    IN ULONGLONG FirstPage,
+    IN ULONGLONG RemoveStart,
     IN ULONGLONG PageCount,
     IN ULONG Flags,
     OUT PMEMORY_DESCRIPTOR_LIST Unused
@@ -531,46 +531,153 @@ Arguments:
 
     Mdl - MDL to remove the region from.
 
-    FirstPage - The first page in the region.
+    RemoveStart - The first page in the region.
 
     PageCount - The number of pages in the region.
 
     Flags - MDL_OPERATION_FLAGS_*.
+        MDL_OPERATION_FLAGS_PHYSICAL if the region is physical.
+        MDL_OPERATION_FLAGS_VIRTUAL if the region is virtual.
 
     Unused - Unused.
 
 Return Value:
 
-    None.
+    STATUS_SUCCESS if successful,
+    STATUS_INVALID_PARAMETER if Flags value is invalid.
 
 --*/
 
 {
-    ULONGLONG RemoveEnd, DescriptorEnd;
-    PLIST_ENTRY Entry;
-    PMEMORY_DESCRIPTOR Descriptor;
-    MEMORY_DESCRIPTOR RemovedDescriptor;
+    NTSTATUS Status;
+    ULONG Offset;
+    BOOLEAN Mapped;
+    ULONGLONG RemoveEnd;
+    PLIST_ENTRY ListEntry;
+    ULONGLONG DescriptorStart, DescriptorEnd;
+    PMEMORY_DESCRIPTOR Descriptor, NewDescriptor;
 
-    (VOID)Flags;
     (VOID)Unused;
 
-    RemoveEnd = FirstPage + PageCount;
-    Entry = Mdl->Head->Flink;
-    while (Entry != Mdl->Head) {
-        Descriptor = (PMEMORY_DESCRIPTOR)Entry;
-        DescriptorEnd = Descriptor->FirstPage + Descriptor->PageCount;
-
-        RtlCopyMemory(&RemovedDescriptor, Descriptor, sizeof(MEMORY_DESCRIPTOR));
-
-        // if (FirstPage <= Descriptor->FirstPage && Descriptor->FirstPage < RemoveEnd) {
-        // }
+    Mapped = FALSE;
+    if (Flags & MDL_OPERATION_FLAGS_VIRTUAL) {
+        if (Mdl->Type == MDL_TYPE_PHYSICAL) {
+            Mapped = TRUE;
+        }
+    } else {
+        //
+        // If the MDL is virtual, the
+        // virtual flag must be set.
+        //
+        if (Mdl->Type == MDL_TYPE_VIRTUAL) {
+            return STATUS_INVALID_PARAMETER;
+        }
     }
 
-    //
-    // TODO: Implement the rest of this routine.
-    //
+    RemoveEnd = RemoveStart + PageCount;
+    ListEntry = Mdl->Head->Flink;
+    while (ListEntry != Mdl->Head) {
+        Descriptor = (PMEMORY_DESCRIPTOR)ListEntry;
 
-    return STATUS_SUCCESS;
+        if (Mapped) {
+            DescriptorStart = Descriptor->MappedFirstPage;
+        } else {
+            DescriptorStart = Descriptor->FirstPage;
+        }
+        DescriptorEnd = DescriptorStart + Descriptor->PageCount;
+
+        //
+        // Check if the region to be removed
+        // is inside the current descriptor.
+        //
+        if (RemoveStart <= DescriptorStart && RemoveEnd > DescriptorStart) {
+            //
+            // The region is around the start of the descriptor, or
+            // they have identical locations and sizes.
+            // | RemoveStart | DescriptorStart | RemoveEnd | DescriptorEnd |
+            // | Lower address ............................ Higher address |
+            //
+
+            if (RemoveEnd < DescriptorEnd) {
+                Offset = RemoveEnd - DescriptorStart;
+            } else {
+                Offset = DescriptorEnd - DescriptorStart;
+            }
+
+            //
+            // Shrink the descriptor.
+            //
+            Descriptor->FirstPage += Offset;
+            Descriptor->PageCount -= Offset;
+            if (Descriptor->MappedFirstPage) {
+                Descriptor->MappedFirstPage += Offset;
+            }
+
+            //
+            // Remove descriptor if now empty.
+            //
+            if (Descriptor->PageCount == 0) {
+                MmMdRemoveDescriptorFromList(Mdl, Descriptor);
+                MmMdFreeDescriptor(Descriptor);
+            }
+        } else if (RemoveStart < DescriptorEnd && RemoveEnd >= DescriptorEnd) {
+            //
+            // The region is around the end of the descriptor.
+            // | DescriptorStart | RemoveStart | DescriptorEnd | RemoveEnd |
+            // | Lower address ............................ Higher address |
+            //
+
+            //
+            // Simply shrink the descriptor.
+            //
+            Descriptor->PageCount -= DescriptorEnd - RemoveStart;
+        } else if (RemoveStart > DescriptorStart && RemoveEnd < DescriptorEnd) {
+            //
+            // The region is completely inside the descriptor.
+            // In this case, the descriptor must be split in two.
+            // | DescriptorStart | RemoveStart | RemoveEnd | DescriptorEnd |
+            // | Lower address ............................ Higher address |
+            //
+
+            //
+            // Create a new descriptor before the removed region.
+            //
+            NewDescriptor = MmMdInitDescriptor(
+                Descriptor->FirstPage,
+                Descriptor->MappedFirstPage,
+                RemoveStart - DescriptorStart,
+                Descriptor->Attributes,
+                Descriptor->Type
+            );
+
+            //
+            // Shrink and move the current descriptor.
+            //
+            Offset = NewDescriptor->PageCount + PageCount;
+            Descriptor->FirstPage += Offset;
+            Descriptor->PageCount -= Offset;
+            if (Descriptor->MappedFirstPage) {
+                Descriptor->MappedFirstPage += Offset;
+            }
+
+            //
+            // Now check if MmMdInitDescriptor() actually worked.
+            //
+            if (NewDescriptor == NULL) {
+                return STATUS_NO_MEMORY;
+            }
+
+            Status = MmMdAddDescriptorToList(Mdl, NewDescriptor, Flags);
+            if (!NT_SUCCESS(Status)) {
+                MmMdFreeDescriptor(NewDescriptor);
+                return Status;
+            }
+        }
+
+        ListEntry = ListEntry->Flink;
+    }
+
+    return Status;
 }
 
 NTSTATUS