chore: Add MA_ALLOCATE_AUTO_OR_RETURN

This is used to replace the recently removed `PoolAllocator`. It handles casting, null checking, and declaring auto resource inline.
trungnt2910 · Nov 19, 2024 · 6c77717 · 6c77717
1 parent a1da7c2
commit 6c77717
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/lxmonika/include_private/monika_private.h b/lxmonika/include_private/monika_private.h
@@ -234,6 +234,15 @@ extern SIZE_T MapLxssProviderIndex;
     }                                                                           \
     while (FALSE)
 
+#define MA_ALLOCATE_AUTO_OR_RETURN(Type, Name, Flags, NumberOfBytes, Tag)       \
+    Type Name;                                                                  \
+    Name = (Type)ExAllocatePool2((Flags), (NumberOfBytes), (Tag));              \
+    if (Name == NULL)                                                           \
+    {                                                                           \
+        return STATUS_NO_MEMORY;                                                \
+    }                                                                           \
+    AUTO_RESOURCE(Name, [](auto p__){ ExFreePoolWithTag(p__, (Tag)); })
+
 // The built-in ASSERT causes compile errors in earlier WDK versions.
 #define MA_ASSERT(exp)                                                          \
     do                                                                          \

diff --git a/lxmonika/src/module.cpp b/lxmonika/src/module.cpp
@@ -36,14 +36,10 @@ MdlpFindModuleByName(
         // As we do not know the size yet, it will return STATUS_INFO_LENGTH_MISMATCH.
         ZwQuerySystemInformation(SystemModuleInformation, (PVOID)&uLen, 0, &uLen);
 
-        PRTL_PROCESS_MODULES pModules = (PRTL_PROCESS_MODULES)ExAllocatePool2(
+        MA_ALLOCATE_AUTO_OR_RETURN(
+            PRTL_PROCESS_MODULES, pModules,
             POOL_FLAG_PAGED, uLen, MDL_POOL_TAG
         );
-        if (pModules == NULL)
-        {
-            return STATUS_NO_MEMORY;
-        }
-        AUTO_RESOURCE(pModules, [](auto p) { ExFreePoolWithTag(p, MDL_POOL_TAG); });
 
         status = ZwQuerySystemInformation(SystemModuleInformation, (PVOID)pModules, uLen, &uLen);
         if (!NT_SUCCESS(status))