From e68aadd2b9d42b7b8d4ffab58b8580e4266ed973 Mon Sep 17 00:00:00 2001
From: ami-GS <1991.daiki@gmail.com>
Date: Wed, 4 Sep 2024 23:03:51 -0700
Subject: [PATCH] VirtualAlloc and affinity settings

---
 inc/cxplat_winkernel.h | 35 ++++++++++++++++++++++++++++++++++
 inc/cxplat_winuser.h   | 43 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/inc/cxplat_winkernel.h b/inc/cxplat_winkernel.h
index 54dd6ad..8cbd65e 100644
--- a/inc/cxplat_winkernel.h
+++ b/inc/cxplat_winkernel.h
@@ -119,6 +119,9 @@ CxPlatLogAssert(
 #define CXPLAT_ALLOC_PAGED(Size, Tag) ExAllocatePool2(POOL_FLAG_PAGED | POOL_FLAG_UNINITIALIZED, Size, Tag)
 #define CXPLAT_ALLOC_NONPAGED(Size, Tag) ExAllocatePool2(POOL_FLAG_NON_PAGED | POOL_FLAG_UNINITIALIZED, Size, Tag)
 #define CXPLAT_FREE(Mem, Tag) ExFreePoolWithTag((void*)Mem, Tag)
+#define CXPLAT_VIRTUAL_ALLOC(Size, TypeFlags, ProtectionFlags, Tag) \
+    CXPLAT_ALLOC_NONPAGED(Size, Tag)
+#define CXPLAT_VIRTUAL_FREE(Mem, Size, Type, Tag) CXPLAT_FREE(Mem, Tag);
 
 #define CxPlatZeroMemory RtlZeroMemory
 #define CxPlatCopyMemory RtlCopyMemory
@@ -603,6 +606,38 @@ CxPlatThreadCreate(
 typedef ULONG_PTR CXPLAT_THREAD_ID;
 #define CxPlatCurThreadID() ((CXPLAT_THREAD_ID)PsGetCurrentThreadId())
 
+inline
+BOOLEAN
+CxPlatSetThreadNodeAffinity(
+    LONG NodeAffinity
+    )
+{
+    GROUP_AFFINITY group;
+
+    KeQueryNodeActiveAffinity((USHORT)NodeAffinity, &group, NULL);
+    KeSetSystemGroupAffinityThread(&group, NULL);
+
+    return TRUE;
+}
+
+inline
+BOOLEAN
+CxPlatSetThreadGroupAffinity(
+    LONG GroupNumber,
+    DWORD_PTR CpuAffinity
+    )
+{
+    GROUP_AFFINITY group = {0};
+
+    group.Mask = CpuAffinity;
+    group.Group = (USHORT)GroupNumber;
+
+    KeSetSystemGroupAffinityThread(&group, NULL);
+
+    return TRUE;
+}
+
+
 //
 // Crypto Interfaces
 //
diff --git a/inc/cxplat_winuser.h b/inc/cxplat_winuser.h
index 6a59c20..88c95a7 100644
--- a/inc/cxplat_winuser.h
+++ b/inc/cxplat_winuser.h
@@ -146,6 +146,9 @@ CxPlatFree(
 #define CXPLAT_ALLOC_PAGED(Size, Tag) CxPlatAlloc(Size, Tag)
 #define CXPLAT_ALLOC_NONPAGED(Size, Tag) CxPlatAlloc(Size, Tag)
 #define CXPLAT_FREE(Mem, Tag) CxPlatFree((void*)Mem, Tag)
+#define CXPLAT_VIRTUAL_ALLOC(Size, TypeFlags, ProtectionFlags, Tag) \
+    VirtualAlloc(NULL, Size, TypeFlags, ProtectionFlags)
+#define CXPLAT_VIRTUAL_FREE(Mem, Size, Type, Tag) VirtualFree(Mem, Size, Type);
 
 #define CxPlatZeroMemory RtlZeroMemory
 #define CxPlatCopyMemory RtlCopyMemory
@@ -597,6 +600,46 @@ CxPlatThreadWaitWithTimeout(
 typedef uint32_t CXPLAT_THREAD_ID;
 #define CxPlatCurThreadID() GetCurrentThreadId()
 
+inline
+BOOLEAN
+CxPlatSetThreadNodeAffinity(
+    LONG NodeAffinity
+    )
+{
+    GROUP_AFFINITY group;
+
+    if (!GetNumaNodeProcessorMaskEx((USHORT)NodeAffinity, &group)) {
+        CXPLAT_DBG_ASSERT(FALSE);
+        return FALSE;
+    }
+    if (!SetThreadGroupAffinity(GetCurrentThread(), &group, NULL)) {
+        CXPLAT_DBG_ASSERT(FALSE);
+        return FALSE;
+    }
+
+    return TRUE;
+}
+
+inline
+BOOLEAN
+CxPlatSetThreadGroupAffinity(
+    LONG GroupNumber,
+    DWORD_PTR CpuAffinity
+    )
+{
+    GROUP_AFFINITY group = {0};
+
+    group.Mask = CpuAffinity;
+    group.Group = (WORD)GroupNumber;
+    if (!SetThreadGroupAffinity(GetCurrentThread(), &group, NULL)) {
+        CXPLAT_DBG_ASSERT(FALSE);
+        return FALSE;
+    }
+
+    return TRUE;
+}
+
+
 //
 // Crypto Interfaces
 //