From f9e15717650cbd27b0ed1d7b313f873f2d801fc2 Mon Sep 17 00:00:00 2001
From: hasherezade <hasherezade@gmail.com>
Date: Sat, 31 Aug 2024 09:51:34 -0700
Subject: [PATCH] [REFACT] Use custom buffer to store thread info

---
 CMakeLists.txt         |  1 +
 utils/custom_buffer.h  | 41 +++++++++++++++++++++++++++++++++++++++++
 utils/threads_util.cpp | 38 ++++++++++++++++++++++----------------
 3 files changed, 64 insertions(+), 16 deletions(-)
 create mode 100644 utils/custom_buffer.h

diff --git a/CMakeLists.txt b/CMakeLists.txt
index d9be2f22d..4722fd0ba 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -180,6 +180,7 @@ set (utils_hdrs
 	utils/byte_buffer.h
 	utils/code_patterns.h
 	utils/custom_mutex.h
+	utils/custom_buffer.h
 	utils/process_symbols.h
 )
 
diff --git a/utils/custom_buffer.h b/utils/custom_buffer.h
new file mode 100644
index 000000000..4305057c7
--- /dev/null
+++ b/utils/custom_buffer.h
@@ -0,0 +1,41 @@
+#pragma once
+#include <windows.h>
+
+namespace pesieve {
+	namespace util {
+
+		struct AutoBuffer
+		{
+			AutoBuffer() : buf(nullptr), max_size(0), buf_size(0) { }
+
+			~AutoBuffer() {
+				if (buf) {
+					::free(buf);
+					buf = nullptr;
+				}
+				max_size = 0;
+				buf_size = 0;
+			}
+
+			BYTE* alloc(size_t _buf_size)
+			{
+				if (_buf_size > max_size) {
+					BYTE* allocated = (BYTE*)::realloc((void*)buf, _buf_size);
+					if (!allocated) {
+						return nullptr;
+					}
+					buf = allocated;
+					max_size = _buf_size;
+				}
+				buf_size = _buf_size;
+				::memset(buf, 0, max_size);
+				return buf;
+			}
+
+			BYTE* buf;
+			size_t max_size;
+			size_t buf_size;
+		};
+
+	}; //namespace util
+}; //namespace pesieve
diff --git a/utils/threads_util.cpp b/utils/threads_util.cpp
index f8e6b4885..af64ec04f 100644
--- a/utils/threads_util.cpp
+++ b/utils/threads_util.cpp
@@ -3,41 +3,47 @@
 #include <peconv.h>
 #include <tlhelp32.h>
 #include "../utils/ntddk.h"
+#include "custom_buffer.h"
 
 #ifdef _DEBUG
 #include <iostream>
+
+void print_info(const pesieve::util::thread_info &threadi)
+{
+	std::cout << std::dec << "TID: " << threadi.tid;
+	if (threadi.is_extended) {
+		std::cout << std::hex << " Start: " << threadi.ext.start_addr << " State: " << threadi.ext.state;
+		if (threadi.ext.state == Waiting) {
+			std::cout << " Reason: " << threadi.ext.wait_reason << " Time: " << threadi.ext.wait_time;
+		}
+	}
+	std::cout << "\n";
+}
 #endif
 
 bool pesieve::util::fetch_threads_info(DWORD pid, std::vector<thread_info>& threads_info)
 {
-	BYTE* buffer = nullptr;
-	ULONG buffer_size = 0;
-	ULONG ret_len = 0;
+	AutoBuffer bBuf;
 
 	NTSTATUS status = STATUS_UNSUCCESSFUL;
 	while (status != STATUS_SUCCESS) {
-		status = NtQuerySystemInformation(SystemProcessInformation, buffer, buffer_size, &ret_len);
+		ULONG ret_len = 0;
+		status = NtQuerySystemInformation(SystemProcessInformation, bBuf.buf, bBuf.buf_size, &ret_len);
 		if (status == STATUS_INFO_LENGTH_MISMATCH) {
-			free(buffer);
-			buffer = nullptr;
-			buffer_size = 0;
-			buffer = (BYTE*)calloc(ret_len, 1);
-			if (!buffer) {
+			if (!bBuf.alloc(ret_len)) {
 				return false;
 			}
-			buffer_size = ret_len;
 			continue; // try again
 		}
 		break; //other error, or success
 	};
 
 	if (status != STATUS_SUCCESS) {
-		free(buffer);
 		return false;
 	}
 
 	bool found = false;
-	SYSTEM_PROCESS_INFORMATION* info = (SYSTEM_PROCESS_INFORMATION*)buffer;
+	SYSTEM_PROCESS_INFORMATION* info = (SYSTEM_PROCESS_INFORMATION*)bBuf.buf;
 	while (info) {
 		if (info->UniqueProcessId == pid) {
 			found = true;
@@ -55,13 +61,12 @@ bool pesieve::util::fetch_threads_info(DWORD pid, std::vector<thread_info>& thre
 			break;
 		}
 		info = (SYSTEM_PROCESS_INFORMATION*)((ULONG_PTR)info + info->NextEntryOffset);
-		if (!peconv::validate_ptr(buffer, buffer_size, info, sizeof(SYSTEM_PROCESS_INFORMATION))) {
+		if (!peconv::validate_ptr(bBuf.buf, bBuf.buf_size, info, sizeof(SYSTEM_PROCESS_INFORMATION))) {
 			break;
 		}
 	}
 
 	if (!found) {
-		free(buffer);
 		return false;
 	}
 
@@ -76,9 +81,10 @@ bool pesieve::util::fetch_threads_info(DWORD pid, std::vector<thread_info>& thre
 		threadi.ext.wait_reason = info->Threads[i].WaitReason;
 		threadi.ext.wait_time  = info->Threads[i].WaitTime;
 		threads_info.push_back(threadi);
+#ifdef _DEBUG
+		print_info(threadi);
+#endif
 	}
-
-	free(buffer);
 	return true;
 }