From 53265f6b7f043946d2511c6508528d6224289002 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: Sat, 31 Dec 2022 17:35:25 +0800
Subject: [PATCH] Users: fix memleaks (Windows)

---
 CMakeLists.txt                        |  1 +
 src/detection/users/users_windows.c   | 24 ++------------
 src/detection/users/wtsapi32_extend.h | 45 +++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 22 deletions(-)
 create mode 100644 src/detection/users/wtsapi32_extend.h

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 8a0b6966e..715211871 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -654,6 +654,7 @@ if(APPLE)
         PRIVATE "-weak_framework MediaRemote -F /System/Library/PrivateFrameworks"
     )
 elseif(WIN32)
+    target_compile_definitions(libfastfetch PRIVATE -D_WIN32_WINNT=0x0601)
     target_link_libraries(libfastfetch
         PRIVATE "dwmapi"
         PRIVATE "gdi32"
diff --git a/src/detection/users/users_windows.c b/src/detection/users/users_windows.c
index 0b4316c10..7ed1c70fa 100644
--- a/src/detection/users/users_windows.c
+++ b/src/detection/users/users_windows.c
@@ -2,27 +2,7 @@
 #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);
+#include "wtsapi32_extend.h"
 
 void ffDetectUsers(FFlist* users, FFstrbuf* error)
 {
@@ -48,7 +28,7 @@ void ffDetectUsers(FFlist* users, FFstrbuf* error)
         ffStrbufInitF((FFstrbuf*)ffListAdd(users), "%s\\%s", domainName.chars, userName.chars);
     }
 
-    WTSFreeMemory(sessionInfo);
+    WTSFreeMemoryExW(WTSTypeSessionInfoLevel1, sessionInfo, 1);
 
     if(users->length == 0)
         ffStrbufAppendS(error, "Unable to detect users");
diff --git a/src/detection/users/wtsapi32_extend.h b/src/detection/users/wtsapi32_extend.h
new file mode 100644
index 000000000..5c89beb2f
--- /dev/null
+++ b/src/detection/users/wtsapi32_extend.h
@@ -0,0 +1,45 @@
+#pragma once
+
+//at the time of writing, <wtsapi32.h> of MinGW doesn't have the definition of WTSEnumerateSessionsExW and friends
+
+#if (_WIN32_WINNT >= 0x0601)
+
+typedef struct _WTS_SESSION_INFO_1A {
+  DWORD ExecEnvId;
+  WTS_CONNECTSTATE_CLASS State;
+  DWORD SessionId;
+  LPSTR pSessionName;
+  LPSTR pHostName;
+  LPSTR pUserName;
+  LPSTR pDomainName;
+  LPSTR pFarmName;
+} WTS_SESSION_INFO_1A, *PWTS_SESSION_INFO_1A;
+
+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;
+
+#define WTS_SESSION_INFO_1 __MINGW_NAME_AW(WTS_SESSION_INFO_1)
+#define PWTS_SESSION_INFO_1 __MINGW_NAME_AW(PWTS_SESSION_INFO_1)
+
+WINBOOL WINAPI WTSEnumerateSessionsExA(HANDLE hServer,DWORD* pLevel,DWORD Filter,PWTS_SESSION_INFO_1A* ppSessionInfo,DWORD* pCount);
+WINBOOL WINAPI WTSEnumerateSessionsExW(HANDLE hServer,DWORD* pLevel,DWORD Filter,PWTS_SESSION_INFO_1W* ppSessionInfo,DWORD* pCount);
+#define WTSEnumerateSessionsEx __MINGW_NAME_AW(WTSEnumerateSessionsEx)
+
+typedef enum _WTS_TYPE_CLASS {
+  WTSTypeProcessInfoLevel0,
+  WTSTypeProcessInfoLevel1,
+  WTSTypeSessionInfoLevel1
+} WTS_TYPE_CLASS;
+BOOL WTSFreeMemoryExA(WTS_TYPE_CLASS WTSTypeClass,PVOID pMemory,ULONG NumberOfEntries);
+BOOL WTSFreeMemoryExW(WTS_TYPE_CLASS WTSTypeClass,PVOID pMemory,ULONG NumberOfEntries);
+#define WTSFreeMemoryEx __MINGW_NAME_AW(WTSFreeMemoryEx)
+
+#endif /*(_WIN32_WINNT >= 0x0601)*/