From 43c6c75710ea98c7d9cf12a117a956851b7945e0 Mon Sep 17 00:00:00 2001
From: Kaimakan71 <undefined.foss@gmail.com>
Date: Sun, 1 Sep 2024 16:56:18 -0400
Subject: [PATCH] [BOOT:MM] Fix descriptor memory leaks

Sometimes when handling MmMdAddDescriptorToList() errors, descriptors
were not freed with MmMdFreeDescriptor().
---
 BOOT/ENVIRON/LIB/EFI/efimm.c |  5 +++++
 BOOT/ENVIRON/LIB/MM/mmmd.c   | 11 +++++++++--
 BOOT/ENVIRON/LIB/MM/mmpa.c   |  1 +
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/BOOT/ENVIRON/LIB/EFI/efimm.c b/BOOT/ENVIRON/LIB/EFI/efimm.c
index f0f3647..898b4a3 100644
--- a/BOOT/ENVIRON/LIB/EFI/efimm.c
+++ b/BOOT/ENVIRON/LIB/EFI/efimm.c
@@ -421,6 +421,7 @@ Return Value:
 
                 Status = MmMdAddDescriptorToList(Mdl, NtDescriptor, MDL_OPERATION_FLAGS_TRUNCATE);
                 if (!NT_SUCCESS(Status)) {
+                    MmMdFreeDescriptor(NtDescriptor);
                     goto exit;
                 }
 
@@ -456,6 +457,7 @@ Return Value:
 
                 Status = MmMdAddDescriptorToList(Mdl, NtDescriptor, MDL_OPERATION_FLAGS_TRUNCATE);
                 if (!NT_SUCCESS(Status)) {
+                    MmMdFreeDescriptor(NtDescriptor);
                     goto exit;
                 }
 
@@ -486,6 +488,7 @@ Return Value:
 
                     Status = MmMdAddDescriptorToList(Mdl, NtDescriptor, MDL_OPERATION_FLAGS_TRUNCATE);
                     if (!NT_SUCCESS(Status)) {
+                        MmMdFreeDescriptor(NtDescriptor);
                         goto exit;
                     }
                 }
@@ -507,6 +510,7 @@ Return Value:
 
                 Status = MmMdAddDescriptorToList(Mdl, NtDescriptor, MDL_OPERATION_FLAGS_TRUNCATE);
                 if (!NT_SUCCESS(Status)) {
+                    MmMdFreeDescriptor(NtDescriptor);
                     goto exit;
                 }
 
@@ -568,6 +572,7 @@ Return Value:
 
         Status = MmMdAddDescriptorToList(Mdl, NtDescriptor, MDL_OPERATION_FLAGS_TRUNCATE);
         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 70c1b85..a2fd346 100644
--- a/BOOT/ENVIRON/LIB/MM/mmmd.c
+++ b/BOOT/ENVIRON/LIB/MM/mmmd.c
@@ -152,6 +152,7 @@ Return Value:
 --*/
 
 {
+    NTSTATUS Status;
     PMEMORY_DESCRIPTOR PrevDescriptor, NextDescriptor, NewDescriptor;
     ULONGLONG DescriptorEnd, PrevDescriptorEnd, NextDescriptorEnd;
     ULONGLONG MappedFirstPage;
@@ -186,7 +187,10 @@ Return Value:
                     PrevDescriptor->Type
                 );
                 if (NewDescriptor != NULL) {
-                    MmMdAddDescriptorToList(Mdl, NewDescriptor, Flags);
+                    Status = MmMdAddDescriptorToList(Mdl, NewDescriptor, Flags);
+                    if (!NT_SUCCESS(Status)) {
+                        MmMdFreeDescriptor(NewDescriptor);
+                    }
                 }
             }
 
@@ -242,7 +246,10 @@ Return Value:
                     Descriptor->Type
                 );
                 if (NewDescriptor != NULL) {
-                    MmMdAddDescriptorToList(Mdl, NewDescriptor, Flags);
+                    Status = MmMdAddDescriptorToList(Mdl, NewDescriptor, Flags);
+                    if (!NT_SUCCESS(Status)) {
+                        MmMdFreeDescriptor(NewDescriptor);
+                    }
                 }
             }
 
diff --git a/BOOT/ENVIRON/LIB/MM/mmpa.c b/BOOT/ENVIRON/LIB/MM/mmpa.c
index 38a8dcc..8845966 100644
--- a/BOOT/ENVIRON/LIB/MM/mmpa.c
+++ b/BOOT/ENVIRON/LIB/MM/mmpa.c
@@ -141,6 +141,7 @@ Return Value:
 
         Status = MmMdAddDescriptorToList(&MmMdlReservedAllocated, NewDescriptor, 0x00);
         if (!NT_SUCCESS(Status)) {
+            MmMdFreeDescriptor(NewDescriptor);
             return Status;
         }