From 0cad657ce22de7d192e3cdbde4d81b41fd0fb954 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: Fri, 2 Dec 2022 01:01:17 +0800
Subject: [PATCH] Swap: correct implementation (Windows)

---
 CMakeLists.txt                      |  2 +-
 src/detection/swap/swap_windows.c   | 11 -----
 src/detection/swap/swap_windows.cpp | 64 +++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 12 deletions(-)
 delete mode 100644 src/detection/swap/swap_windows.c
 create mode 100644 src/detection/swap/swap_windows.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index a2049b950..bb119b853 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -465,7 +465,7 @@ elseif(WIN32)
         src/detection/packages/packages_windows.c
         src/detection/poweradapter/poweradapter_nosupport.c
         src/detection/processes/processes_windows.cpp
-        src/detection/swap/swap_windows.c
+        src/detection/swap/swap_windows.cpp
         src/detection/terminalfont/terminalfont_windows.c
         src/detection/terminalshell/terminalshell_windows.cpp
         src/detection/uptime/uptime_windows.c
diff --git a/src/detection/swap/swap_windows.c b/src/detection/swap/swap_windows.c
deleted file mode 100644
index 9fc1d86a4..000000000
--- a/src/detection/swap/swap_windows.c
+++ /dev/null
@@ -1,11 +0,0 @@
-#include "swap.h"
-
-void ffDetectSwapImpl(FFMemoryStorage* swap)
-{
-    MEMORYSTATUSEX statex = {
-        .dwLength = sizeof(statex),
-    };
-    GlobalMemoryStatusEx(&statex);
-    swap->bytesTotal = statex.ullTotalPageFile;
-    swap->bytesUsed = statex.ullTotalPageFile - statex.ullAvailPageFile;
-}
diff --git a/src/detection/swap/swap_windows.cpp b/src/detection/swap/swap_windows.cpp
new file mode 100644
index 000000000..5e7597602
--- /dev/null
+++ b/src/detection/swap/swap_windows.cpp
@@ -0,0 +1,64 @@
+extern "C" {
+#include "swap.h"
+#include "util/mallocHelper.h"
+}
+
+#ifdef FF_USE_WIN_NTAPI
+
+#include <winternl.h>
+#include <ntstatus.h>
+
+extern "C"
+void ffDetectSwapImpl(FFMemoryStorage* swap)
+{
+    SYSTEM_INFO sysInfo;
+    GetNativeSystemInfo(&sysInfo);
+
+    ULONG size = sizeof(SYSTEM_PAGEFILE_INFORMATION);
+    SYSTEM_PAGEFILE_INFORMATION* FF_AUTO_FREE pstart = (SYSTEM_PAGEFILE_INFORMATION*)malloc(size);
+    while(true)
+    {
+        NTSTATUS status = NtQuerySystemInformation(SystemPagefileInformation, pstart, size, &size);
+        if(status == STATUS_INFO_LENGTH_MISMATCH)
+        {
+            if(!(pstart = (SYSTEM_PAGEFILE_INFORMATION*)realloc(pstart, size)))
+            {
+                ffStrbufAppendF(&swap->error, "relloc(pstart, %lu) failed", size);
+                return;
+            }
+        }
+        else if(!NT_SUCCESS(status))
+        {
+            ffStrbufAppendF(&swap->error, "NtQuerySystemInformation(SystemPagefileInformation, %lu) failed", size);
+            return;
+        }
+        break;
+    }
+    swap->bytesUsed = (uint64_t)pstart->TotalUsed * sysInfo.dwPageSize;
+    swap->bytesTotal = (uint64_t)pstart->CurrentSize * sysInfo.dwPageSize;
+}
+
+#else
+
+#include "util/windows/wmi.hpp"
+
+extern "C"
+void ffDetectSwapImpl(FFMemoryStorage* swap)
+{
+    FFWmiQuery query(L"SELECT AllocatedBaseSize, CurrentUsage FROM Win32_PageFileUsage", &swap->error);
+    if(!query)
+        return;
+
+    if(FFWmiRecord record = query.next())
+    {
+        //MB
+        record.getUnsigned(L"AllocatedBaseSize", &swap->bytesTotal);
+        record.getUnsigned(L"CurrentUsage", &swap->bytesUsed);
+        swap->bytesTotal *= 1024 * 1024;
+        swap->bytesUsed *= 1024 * 1024;
+    }
+    else
+        ffStrbufInitS(&swap->error, "No Wmi result returned");
+}
+
+#endif