From 0f08233581e9688e91b89e8a171fb6e8ac95a2fa 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: Wed, 23 Nov 2022 19:11:10 +0800
Subject: [PATCH] Users: improve performance (Windows)

---
 CMakeLists.txt                        |  3 +-
 src/detection/users/users_windows.c   | 55 +++++++++++++++++++++++++++
 src/detection/users/users_windows.cpp | 35 -----------------
 3 files changed, 57 insertions(+), 36 deletions(-)
 create mode 100644 src/detection/users/users_windows.c
 delete mode 100644 src/detection/users/users_windows.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6675f802b..313d46e25 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -450,7 +450,7 @@ elseif(WIN32)
         src/detection/terminalfont/terminalfont_windows.c
         src/detection/terminalshell/terminalshell_windows.cpp
         src/detection/uptime/uptime_windows.c
-        src/detection/users/users_windows.cpp
+        src/detection/users/users_windows.c
         src/detection/wmtheme/wmtheme_windows.c
         src/util/windows/getline.c
         src/util/windows/pwd.c
@@ -555,6 +555,7 @@ elseif(WIN32)
         PRIVATE "version"
         PRIVATE "setupapi"
         PRIVATE "dxgi"
+        PRIVATE "wtsapi32"
     )
     if(USE_WIN_NTAPI)
         target_compile_definitions(libfastfetch PRIVATE FF_USE_WIN_NTAPI)
diff --git a/src/detection/users/users_windows.c b/src/detection/users/users_windows.c
new file mode 100644
index 000000000..46e5b4add
--- /dev/null
+++ b/src/detection/users/users_windows.c
@@ -0,0 +1,55 @@
+#include "users.h"
+#include "util/windows/unicode.h"
+
+#include <wtsapi32.h>
+
+//at the time of writing, <wtsapi32.h> of MinGW doesn't have the definition of WTSEnumerateSessionsExW
+typedef struct _WTS_SESSION_INFO_1W {
+    DWORD ExecEnvId;
+    WTS_CONNECTSTATE_CLASS State;
+    DWORD SessionId;
+    LPWSTR pSessionName;
+    LPWSTR pHostName;
+    LPWSTR pUserName;
+    LPWSTR pDomainName;
+    LPWSTR pFarmName;
+} WTS_SESSION_INFO_1W, * PWTS_SESSION_INFO_1W;
+
+BOOL
+WINAPI
+WTSEnumerateSessionsExW(
+    HANDLE hServer,
+    DWORD* pLevel,
+    DWORD Filter,
+    PWTS_SESSION_INFO_1W* ppSessionInfo,
+    DWORD* pCount);
+
+void ffDetectUsers(FFlist* users, FFstrbuf* error)
+{
+    WTS_SESSION_INFO_1W* sessionInfo;
+    DWORD sessionCount;
+    DWORD level = 1;
+
+    if(!WTSEnumerateSessionsExW(WTS_CURRENT_SERVER_HANDLE, &level, 0, &sessionInfo, &sessionCount))
+    {
+        ffStrbufAppendS(error, "WTSEnumerateSessionsW(WTS_CURRENT_SERVER_HANDLE) failed");
+        return;
+    }
+
+    for (DWORD i = 0; i < sessionCount; i++)
+    {
+        WTS_SESSION_INFO_1W* session = &sessionInfo[i];
+        if(session->State != WTSActive)
+            continue;
+
+        FF_STRBUF_AUTO_DESTROY domainName = ffStrbufFromWchar(session->pDomainName);
+        FF_STRBUF_AUTO_DESTROY userName = ffStrbufFromWchar(session->pUserName);
+
+        ffStrbufInitF((FFstrbuf*)ffListAdd(users), "%s\\%s", domainName.chars, userName.chars);
+    }
+
+    WTSFreeMemory(sessionInfo);
+
+    if(users->length == 0)
+        ffStrbufAppendS(error, "Unable to detect users");
+}
diff --git a/src/detection/users/users_windows.cpp b/src/detection/users/users_windows.cpp
deleted file mode 100644
index 9e50cd365..000000000
--- a/src/detection/users/users_windows.cpp
+++ /dev/null
@@ -1,35 +0,0 @@
-extern "C" {
-#include "users.h"
-}
-#include "util/windows/wmi.hpp"
-
-void ffDetectUsers(FFlist* users, FFstrbuf* error)
-{
-    FFWmiQuery query(L"SELECT Antecedent FROM Win32_LoggedOnUser", error);
-    if(!query)
-        return;
-
-next:
-    while(FFWmiRecord record = query.next())
-    {
-        FFstrbuf antecedent;
-        ffStrbufInit(&antecedent);
-        record.getString(L"Antecedent", &antecedent); // \\.\root\cimv2:Win32_Account.Domain="DOMAIN",Name="NAME"
-        ffStrbufTrimRight(&antecedent, '"'); // \\.\root\cimv2:Win32_Account.Domain="DOMAIN",Name="NAME
-        ffStrbufSubstrAfterFirstC(&antecedent, '"'); // DOMAIN",Name="NAME
-        uint32_t index = ffStrbufFirstIndexC(&antecedent, '"');
-        ffStrbufRemoveSubstr(&antecedent, index, ffStrbufLastIndexC(&antecedent, '"')); // DOMAIN"NAME
-        antecedent.chars[index] = '\\';
-
-        for(uint32_t i = 0; i < users->length; ++i)
-        {
-            if(ffStrbufComp((FFstrbuf*)ffListGet(users, i), &antecedent) == 0)
-                goto next;
-        }
-
-        *(FFstrbuf*)ffListAdd(users) = antecedent;
-    }
-
-    if(users->length == 0)
-        ffStrbufAppendS(error, "Unable to detect users");
-}