From: Tomasz Konojacki <me@xenu.pl>
Date: Mon, 18 May 2020 08:04:29 +0000 (+0200)
Subject: sdsasd
X-Git-Url: http://git.osdn.net/view?p=btop%2Fsystem-metrics.git;a=commitdiff_plain;h=c98da1c06e2bd4245ef232f9f897d6e59ed69a22

sdsasd
---

diff --git a/cxx/system_metrics.hxx b/cxx/system_metrics.hxx
index 20dc6c5..05ab542 100644
--- a/cxx/system_metrics.hxx
+++ b/cxx/system_metrics.hxx
@@ -22,6 +22,7 @@
 
 #pragma once
 
+#include <cstdef>
 #include <memory>
 #include <string>
 #include <vector>
diff --git a/cxx/win32/win32.cxx b/cxx/win32/win32.cxx
index 7d2e87b..9f9a76c 100644
--- a/cxx/win32/win32.cxx
+++ b/cxx/win32/win32.cxx
@@ -29,24 +29,26 @@
 
 #include <psapi.h>
 #include <windows.h>
+#include <ntstatus.h>
 #include "3rdparty/mingw-w64/winternl.h"
 
 struct sm_private {
     uint64_t page_size;
     std::vector<SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION> proc_perf_buf;
     std::vector<SYSTEM_PAGEFILE_INFORMATION> page_buf;
-    std::vector<WCHAR> path_buf;
+    std::vector<std::byte> proc_info_buf;
 };
 
 static inline uint64_t filetime_to_uint64t(FILETIME &ft);
-bool set_process_path(struct process &p, HANDLE proc_handle, std::vector<WCHAR> &buf);
+static inline void unicode_string_to_string(const UNICODE_STRING &src, std::string &dest);
+static void wstr_to_string(const PWSTR src, int src_len, std::string &dest);
 
 system_metrics::system_metrics()
     :_priv(new sm_private())
 {
     uint64_t n_cpus = this->number_of_cpus();
     this->_priv->proc_perf_buf.resize(n_cpus + 1);
-    this->_priv->path_buf.resize(33000);
+    this->_priv->proc_info_buf.resize(300000);
 }
 
 system_metrics::~system_metrics() = default;
@@ -155,41 +157,46 @@ system_metrics::mem_info()
 std::vector<struct process>
 system_metrics::process_list()
 {
-    std::vector<DWORD> pids(1000);
-    
-    for (;;) {
-        DWORD ret_size_bytes,
-              pids_size_bytes = static_cast<DWORD>(pids.size() * sizeof(pids[0]));
+    auto &buf = this->_priv->proc_info_buf;
 
-        if (!K32EnumProcesses(pids.data(), pids_size_bytes, &ret_size_bytes))
-            throw win32_exception(__func__, "K32EnumProcesses", GetLastError());
+    for (;;) {
+        ULONG ret_size_bytes;
+        NTSTATUS status = NtQuerySystemInformation(
+            SystemProcessInformation,
+            buf.data(),
+            static_cast<ULONG>(buf.size()),
+            &ret_size_bytes
+        );
 
-        if (ret_size_bytes < pids_size_bytes) {
-            pids.resize(ret_size_bytes / sizeof(pids[0]));
-            break;
+        if (status == STATUS_INFO_LENGTH_MISMATCH) {
+            buf.resize(static_cast<size_t>(
+                static_cast<double>(buf.size()) * 1.5
+            ));
+            continue;
         }
-            
-        pids.resize(static_cast<size_t>(
-            static_cast<double>(pids.size()) * 1.5
-        ));
+
+        if (!NT_SUCCESS(status) && status != STATUS_INFO_LENGTH_MISMATCH)
+            throw nt_exception(__func__, "NtQuerySystemInformation", status);
+
+        break;
     }
 
+    SYSTEM_PROCESS_INFORMATION *info = reinterpret_cast<SYSTEM_PROCESS_INFORMATION*>(&buf[0]);
     std::vector<struct process> ret;
-    for (auto const &pid: pids) {
-        HANDLE proc_handle = OpenProcess(
-            MAXIMUM_ALLOWED, FALSE, pid
-        );
-        if (!proc_handle)
-            continue;
-        
+    for (;;) {
         struct process p;
 
-        p.pid = pid;
-        
-        if (!set_process_path(p, proc_handle, this->_priv->path_buf))
-            continue;
+        p.pid = reinterpret_cast<uint64_t>(info->UniqueProcessId);
+        unicode_string_to_string(info->ImageName, p.path);
 
         ret.push_back(p);
+
+        if (info->NextEntryOffset)
+            info = reinterpret_cast<SYSTEM_PROCESS_INFORMATION*>(
+                reinterpret_cast<std::byte*>(info) + info->NextEntryOffset
+            );
+        else
+            break;
     }
     return ret;
 }
@@ -206,32 +213,36 @@ filetime_to_uint64t(FILETIME &ft)
     }.QuadPart;
 }
 
-bool
-set_process_path(struct process &p, HANDLE proc_handle, std::vector<WCHAR> &buf)
+static inline void
+unicode_string_to_string(const UNICODE_STRING &src, std::string &dest)
 {
-    DWORD ret_size = static_cast<DWORD>(buf.size());
-    BOOL succ = QueryFullProcessImageNameW(
-        proc_handle, 0, buf.data(), &ret_size
+    wstr_to_string(
+        src.Buffer, static_cast<int>(src.Length / sizeof(WCHAR)), dest
     );
-    if (!succ)
-        return false;
-    
+}
+
+static void
+wstr_to_string(const PWSTR src, int src_len, std::string &dest)
+{
+    if (!src || src_len == 0) {
+        dest = "";
+        return;
+    }
+
     int out_len = WideCharToMultiByte(
-        CP_UTF8, 0, buf.data(), ret_size, NULL, 0, NULL, NULL
+        CP_UTF8, 0, src, src_len, NULL, 0, NULL, NULL
     );
     if (!out_len)
         throw win32_exception(__func__, "WideCharToMultiByte", GetLastError());
-    
-    p.path.resize(out_len);
+
+    dest.resize(out_len);
 
     out_len = WideCharToMultiByte(
-        CP_UTF8, 0, buf.data(), ret_size, &p.path[0], out_len, NULL, NULL
+        CP_UTF8, 0, src, src_len, dest.data(), out_len, NULL, NULL
     );
     if (!out_len)
         throw win32_exception(__func__, "WideCharToMultiByte", GetLastError());
 
     // just in case
-    p.path.resize(out_len);
-
-    return true;
+    dest.resize(out_len);
 }