From 7fc654769ebc205929f981f6199876e948a51538 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Benjamin=20H=C3=B6glinger-Stelzer?= <nefarius@dhmx.at>
Date: Mon, 3 Jul 2023 01:54:12 +0200
Subject: [PATCH] Warnings and memory allocation fixes

---
 src/Domito.CodeIntegrity.cpp | 24 +++++++++++++++++++-----
 src/Domito.Memory.cpp        |  4 ++--
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/src/Domito.CodeIntegrity.cpp b/src/Domito.CodeIntegrity.cpp
index 49df98d..b566c75 100644
--- a/src/Domito.CodeIntegrity.cpp
+++ b/src/Domito.CodeIntegrity.cpp
@@ -288,7 +288,13 @@ DomitoCalculatePortableExecutableDigest(
     // TODO: Not sure if  16 * 512 * 512 is right. Do something better!
     //
     ULONG copySize = phDos->e_lfanew + sizeof(IMAGE_FILE_HEADER) + 4 + 0x40;
-    const PUCHAR pBuf = (PUCHAR)G_Memory.Allocate(16 * 512 * 512);
+#pragma warning(disable:4996)
+    const PUCHAR pBuf = (PUCHAR)ExAllocatePoolWithTag(
+        NonPagedPool,
+        16 * 512 * 512,
+        DOMITO_POOL_TAG
+    );
+#pragma warning(default:4996)
     if (!pBuf)
     {
         return STATUS_INSUFFICIENT_RESOURCES;
@@ -344,7 +350,9 @@ DomitoCalculatePortableExecutableDigest(
     //
     // Allocate a buffer to store the resulting hash
     //
-    pHash = (PUCHAR)G_Memory.Allocate(hashLength);
+#pragma warning(disable:4996)
+    pHash = (PUCHAR)ExAllocatePoolWithTag(NonPagedPool, hashLength, DOMITO_POOL_TAG);
+#pragma warning(default:4996)
     if (!pHash)
     {
         status = STATUS_INSUFFICIENT_RESOURCES;
@@ -495,7 +503,7 @@ DomitoCalculatePortableExecutableDigest(
 cleanup:
     if (pBuf)
     {
-        G_Memory.Free(pBuf);
+        ExFreePoolWithTag(pBuf, DOMITO_POOL_TAG);
     }
 
     if (hbHash)
@@ -555,7 +563,13 @@ DomitoValidateFileLegacyMode(
         if (SecurityDirectory->Size != 0u &&
             SecurityDirectory->VirtualAddress != 0u)
         {
-            certDirectory = G_Memory.Allocate(SecurityDirectory->Size);
+#pragma warning(disable:4996)
+            certDirectory = ExAllocatePoolWithTag(
+                PagedPool,
+                SecurityDirectory->Size,
+                DOMITO_POOL_TAG
+            );
+#pragma warning(default:4996)
             if (certDirectory == NULL)
             {
                 status = STATUS_INSUFFICIENT_RESOURCES;
@@ -659,7 +673,7 @@ DomitoValidateFileLegacyMode(
 
     if (certDirectory)
     {
-        G_Memory.Free(certDirectory);
+        ExFreePoolWithTag(certDirectory, DOMITO_POOL_TAG);
     }
 
     return status;
diff --git a/src/Domito.Memory.cpp b/src/Domito.Memory.cpp
index cd3300d..51bbefd 100644
--- a/src/Domito.Memory.cpp
+++ b/src/Domito.Memory.cpp
@@ -11,13 +11,13 @@
 static PVOID NTAPI DomitoDefaultMalloc(size_t s)
 {
 #pragma warning(disable:4996)
-    PVOID mem = ExAllocatePoolWithTag(NonPagedPool, s, DOMITO_POOL_TAG);
+    const PVOID mem = ExAllocatePoolWithTag(PagedPool, s, DOMITO_POOL_TAG);
     if (mem)
     {
         RtlZeroMemory(mem, s);
     }
     return mem;
-#pragma warninf(default:4996)
+#pragma warning(default:4996)
 }
 
 static void NTAPI DomitoDefaultFree(PVOID p)