Merge pull request #456 from processhacker/master

[pull] master from processhacker:master
CeruleanSky · Apr 17, 2022 · fcdca33 · fcdca33
2 parents 5d15951 + 457584b
commit fcdca33
Showing 1 changed file with 73 additions and 0 deletions.
diff --git a/phlib/include/phnativeinl.h b/phlib/include/phnativeinl.h
@@ -583,6 +583,22 @@ PhSetProcessQuotaLimits(
         );
 }
 
+FORCEINLINE
+NTSTATUS
+PhGetProcessAffinityMask(
+    _In_ HANDLE ProcessHandle,
+    _Out_ PKAFFINITY AffinityMask
+    )
+{
+    return NtQueryInformationProcess(
+        ProcessHandle,
+        ProcessAffinityMask,
+        AffinityMask,
+        sizeof(KAFFINITY),
+        NULL
+        );
+}
+
 /**
  * Sets a process' affinity mask.
  *
@@ -604,6 +620,35 @@ PhSetProcessAffinityMask(
         );
 }
 
+FORCEINLINE
+NTSTATUS
+PhGetProcessGroupAffinity(
+    _In_ HANDLE ProcessHandle,
+    _Inout_ PUSHORT GroupCount,
+    _Out_ PUSHORT GroupArray
+    )
+{
+    NTSTATUS status;
+    ULONG returnLength;
+
+    // rev from GetProcessGroupAffinity (dmex)
+    status = NtQueryInformationProcess(
+        ProcessHandle,
+        ProcessGroupInformation,
+        GroupArray,
+        sizeof(USHORT) * *GroupCount,
+        &returnLength
+        );
+
+    // (int)(status + 0x80000000) < 0 || status == STATUS_BUFFER_TOO_SMALL
+    if (NT_SUCCESS(status) || status == STATUS_BUFFER_TOO_SMALL)
+    {
+        *GroupCount = (USHORT)returnLength >> 1;
+    }
+
+    return status;
+}
+
 FORCEINLINE
 NTSTATUS
 PhGetProcessIsCFGuardEnabled(
@@ -1270,6 +1315,34 @@ PhGetThreadIsTerminated(
     return status;
 }
 
+FORCEINLINE
+NTSTATUS
+PhGetThreadAffinityMask(
+    _In_ HANDLE ThreadHandle,
+    _Out_ PKAFFINITY AffinityMask
+    )
+{
+    NTSTATUS status;
+    THREAD_BASIC_INFORMATION basicInfo;
+
+    status = PhGetThreadBasicInformation(ThreadHandle, &basicInfo);
+
+    if (NT_SUCCESS(status))
+    {
+        *AffinityMask = basicInfo.AffinityMask;
+    }
+
+    return status;
+
+    //return NtQueryInformationThread(
+    //    ThreadHandle,
+    //    ThreadAffinityMask,
+    //    &AffinityMask,
+    //    sizeof(KAFFINITY),
+    //    NULL
+    //    );
+}
+
 /**
  * Sets a thread's affinity mask.
  *