From 47e616d951b7a5aa379d7727d394d9ddc08ebd3b Mon Sep 17 00:00:00 2001
From: David Spruill <daspr@microsoft.com>
Date: Mon, 20 Jul 2026 17:46:19 -0400
Subject: Fix memory tagging part 2

---
 video/KMDOD/bdd.hxx     | 14 +++++++++++---
 video/KMDOD/bdd_ddi.cxx |  3 +--
 video/KMDOD/blthw.cxx   |  3 ++-
 video/KMDOD/memory.cxx  | 24 ++++++++++++++++++++++--
 4 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/video/KMDOD/bdd.hxx b/video/KMDOD/bdd.hxx
index 5bfcb6f4..8fd3cdbb 100644
--- a/video/KMDOD/bdd.hxx
+++ b/video/KMDOD/bdd.hxx
@@ -626,12 +626,20 @@ IsEdidChecksumValid(_In_reads_bytes_(EDID_V1_BLOCK_SIZE) const BYTE* pEdid);
 // Memory handling
 //
 
-// Defaulting the value of Flags means that any call to new Foo()
+enum class BDD_POOL_TYPE
+{
+    Paged,
+    NonPaged
+};
+
+// Defaulting the value of PoolType means that any call to new Foo()
 // will raise a compiler error for being ambiguous. This is to help keep
 // any calls to allocate memory from accidentally NOT going through
 // these functions.
-void* __cdecl operator new(size_t Size, POOL_FLAGS Flags = POOL_FLAG_PAGED);
-void* __cdecl operator new[](size_t Size, POOL_FLAGS Flags = POOL_FLAG_PAGED);
+void* __cdecl operator new(size_t Size, BDD_POOL_TYPE PoolType = BDD_POOL_TYPE::Paged);
+void* __cdecl operator new[](size_t Size, BDD_POOL_TYPE PoolType = BDD_POOL_TYPE::Paged);
+void  __cdecl operator delete(void* pObject, BDD_POOL_TYPE PoolType);
+void  __cdecl operator delete[](void* pObject, BDD_POOL_TYPE PoolType);
 void  __cdecl operator delete(void* pObject);
 void  __cdecl operator delete(void* pObject, size_t s);
 void  __cdecl operator delete[](void* pObject);
diff --git a/video/KMDOD/bdd_ddi.cxx b/video/KMDOD/bdd_ddi.cxx
index cdbe8594..f5f51cf8 100644
--- a/video/KMDOD/bdd_ddi.cxx
+++ b/video/KMDOD/bdd_ddi.cxx
@@ -103,8 +103,7 @@ BddDdiAddDevice(
     }
     *ppDeviceContext = NULL;
 
-    void* pBDDMemory = operator new(sizeof(BASIC_DISPLAY_DRIVER), POOL_FLAG_NON_PAGED);
-    BASIC_DISPLAY_DRIVER* pBDD = pBDDMemory ? new(pBDDMemory) BASIC_DISPLAY_DRIVER(pPhysicalDeviceObject) : NULL;
+    BASIC_DISPLAY_DRIVER* pBDD = new(BDD_POOL_TYPE::NonPaged) BASIC_DISPLAY_DRIVER(pPhysicalDeviceObject);
     if (pBDD == NULL)
     {
         BDD_LOG_LOW_RESOURCE0("pBDD failed to be allocated");
diff --git a/video/KMDOD/blthw.cxx b/video/KMDOD/blthw.cxx
index 0b1ff0b1..f9e8a51b 100644
--- a/video/KMDOD/blthw.cxx
+++ b/video/KMDOD/blthw.cxx
@@ -281,7 +281,8 @@ BDD_HWBLT::ExecutePresentDisplayOnly(
     SIZE_T sizeRects = NumDirtyRects*sizeof(RECT);
     SIZE_T size = sizeof(DoPresentMemory) + sizeMoves + sizeRects;
 
-    DoPresentMemory* ctx = reinterpret_cast<DoPresentMemory*>(operator new[](size, POOL_FLAG_PAGED));
+    DoPresentMemory* ctx = reinterpret_cast<DoPresentMemory*>
+                                (new (BDD_POOL_TYPE::Paged) BYTE[size]);
 
     if (!ctx)
     {
diff --git a/video/KMDOD/memory.cxx b/video/KMDOD/memory.cxx
index 49e56553..54f6a6dc 100644
--- a/video/KMDOD/memory.cxx
+++ b/video/KMDOD/memory.cxx
@@ -14,11 +14,12 @@
 //
 // New and delete operators
 //
-void* __cdecl operator new(size_t Size, POOL_FLAGS Flags)
+void* __cdecl operator new(size_t Size, BDD_POOL_TYPE PoolType)
 {
     PAGED_CODE();
 
     Size = (Size != 0) ? Size : 1;
+    POOL_FLAGS Flags = PoolType == BDD_POOL_TYPE::NonPaged ? POOL_FLAG_NON_PAGED : POOL_FLAG_PAGED;
     
     // Note that ExAllocatePool2 replaces ExAllocatePool* APIs in OS's starting
     // with Windows 10, version 2004. If your driver targets previous versions it
@@ -35,11 +36,12 @@ void* __cdecl operator new(size_t Size, POOL_FLAGS Flags)
     return pObject;
 }
 
-void* __cdecl operator new[](size_t Size, POOL_FLAGS Flags)
+void* __cdecl operator new[](size_t Size, BDD_POOL_TYPE PoolType)
 {
     PAGED_CODE();
 
     Size = (Size != 0) ? Size : 1;
+    POOL_FLAGS Flags = PoolType == BDD_POOL_TYPE::NonPaged ? POOL_FLAG_NON_PAGED : POOL_FLAG_PAGED;
     
     void* pObject = ExAllocatePool2(Flags, Size, BDDTAG);
 
@@ -53,6 +55,24 @@ void* __cdecl operator new[](size_t Size, POOL_FLAGS Flags)
     return pObject;
 }
 
+void __cdecl operator delete(void* pObject, BDD_POOL_TYPE PoolType)
+{
+    PAGED_CODE();
+
+    UNREFERENCED_PARAMETER(PoolType);
+
+    ::operator delete(pObject);
+}
+
+void __cdecl operator delete[](void* pObject, BDD_POOL_TYPE PoolType)
+{
+    PAGED_CODE();
+
+    UNREFERENCED_PARAMETER(PoolType);
+
+    ::operator delete[](pObject);
+}
+
 void __cdecl operator delete(void* pObject)
 {
     PAGED_CODE();
-- 
cgit v1.3.1