From 308408a96296fc0f1b06f1ab9b2eea4e6bc5aa33 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E9=80=9A=E6=B4=B2?= <carter.li@eoitek.com>
Date: Mon, 21 Nov 2022 13:05:53 +0800
Subject: [PATCH] Processes: improve performance (Windows)

---
 CMakeLists.txt                                |  6 +--
 src/detection/processes/processes_windows.cpp | 46 ++++++++++++++++++-
 .../terminalshell/terminalshell_windows.cpp   |  2 +-
 3 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 21e412f6e..283c47298 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -66,7 +66,7 @@ cmake_dependent_option(ENABLE_OPENCL "Enable opencl" ON "LINUX OR BSD OR WIN32"
 cmake_dependent_option(ENABLE_LIBCJSON "Enable libcjson" ON "LINUX OR WIN32" OFF)
 cmake_dependent_option(ENABLE_FREETYPE "Enable freetype" ON "ANDROID" OFF)
 cmake_dependent_option(ENABLE_THREADS "Enable multithreading" ON "Threads_FOUND AND NOT ANDROID" OFF)
-cmake_dependent_option(USE_WIN_FAST_PPID_DETECTION "Use internal NTAPI instead of querying WMI to get PPID" ON "WIN32" OFF)
+cmake_dependent_option(USE_WIN_NTAPI "Allow using internal NTAPI" ON "WIN32" OFF)
 
 option(BUILD_TESTS "Build tests" OFF) # Also create test executables
 option(SET_TWEAK "Add tweak to project version" ON) # This is set to off by github actions for release builds
@@ -546,8 +546,8 @@ elseif(WIN32)
         PRIVATE "ntdll"
         PRIVATE "version"
     )
-    if(USE_WIN_FAST_PPID_DETECTION)
-        target_compile_definitions(libfastfetch PRIVATE FF_USE_WIN_FAST_PPID_DETECTION)
+    if(USE_WIN_NTAPI)
+        target_compile_definitions(libfastfetch PRIVATE FF_USE_WIN_NTAPI)
     endif()
 endif()
 
diff --git a/src/detection/processes/processes_windows.cpp b/src/detection/processes/processes_windows.cpp
index 68de794a6..052162d2b 100644
--- a/src/detection/processes/processes_windows.cpp
+++ b/src/detection/processes/processes_windows.cpp
@@ -1,12 +1,54 @@
 extern "C" {
 #include "processes.h"
 }
-#include "util/windows/wmi.hpp"
+
+#ifdef FF_USE_WIN_NTAPI
+
+#include <winternl.h>
+
+static inline void wrapFree(SYSTEM_PROCESS_INFORMATION** ptr)
+{
+    free(*ptr);
+}
 
 uint32_t ffDetectProcesses(FFinstance* instance, FFstrbuf* error)
 {
     FF_UNUSED(instance);
 
+    ULONG size = 0;
+    if(NtQuerySystemInformation(SystemProcessInformation, nullptr, 0, &size) != (NTSTATUS)0xC0000004 /*STATUS_INFO_LENGTH_MISMATCH*/)
+    {
+        ffStrbufAppendS(error, "NtQuerySystemInformation(SystemProcessInformation, NULL) failed");
+        return 0;
+    }
+    size += sizeof(SystemProcessInformation) * 5; //What if new processes are created during two syscalls?
+
+    SYSTEM_PROCESS_INFORMATION* __attribute__((__cleanup__(wrapFree))) pstart = (SYSTEM_PROCESS_INFORMATION*)malloc(size);
+    if(!pstart)
+    {
+        ffStrbufAppendF(error, "malloc(%u) failed", (unsigned)size);
+        return 0;
+    }
+
+    if(!NT_SUCCESS(NtQuerySystemInformation(SystemProcessInformation, pstart, size, nullptr)))
+    {
+        ffStrbufAppendS(error, "NtQuerySystemInformation(SystemProcessInformation, pstart) failed");
+        return 0;
+    }
+
+    uint32_t result = 1; //Init with 1 because we test for ptr->NextEntryOffset
+    for (auto ptr = pstart; ptr->NextEntryOffset; ptr = (SYSTEM_PROCESS_INFORMATION*)((uint8_t*)ptr + ptr->NextEntryOffset))
+        ++result;
+
+    return result;
+}
+
+#else
+
+#include "util/windows/wmi.hpp"
+
+uint32_t ffDetectProcesses(FFinstance* instance, FFstrbuf* error)
+{
     FFWmiQuery query(L"SELECT NumberOfProcesses FROM Win32_OperatingSystem", error);
     if(!query)
         return 0;
@@ -23,3 +65,5 @@ uint32_t ffDetectProcesses(FFinstance* instance, FFstrbuf* error)
         return 0;
     }
 }
+
+#endif
diff --git a/src/detection/terminalshell/terminalshell_windows.cpp b/src/detection/terminalshell/terminalshell_windows.cpp
index 6f0139a2e..3aca7e7b4 100644
--- a/src/detection/terminalshell/terminalshell_windows.cpp
+++ b/src/detection/terminalshell/terminalshell_windows.cpp
@@ -10,7 +10,7 @@ extern "C" {
 
 #include <chrono>
 
-#ifdef FF_USE_WIN_FAST_PPID_DETECTION
+#ifdef FF_USE_WIN_NTAPI
 
 #include <winternl.h>