From fc0ea264e6efdc640f3d4c6541f8eff1143ce2ff Mon Sep 17 00:00:00 2001
From: hasherezade <hasherezade@op.pl>
Date: Sun, 12 Jul 2020 23:10:55 +0200
Subject: [PATCH] [BUGFIX] Fixed scanning workingset by 32 bit scanner. Allow
 for ERROR_BAD_LENGHT

---
 scanners/scanner.cpp      |  7 ++++++-
 utils/workingset_enum.cpp | 31 ++++++++++++++-----------------
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/scanners/scanner.cpp b/scanners/scanner.cpp
index 47dd83cb9..6de4097d7 100644
--- a/scanners/scanner.cpp
+++ b/scanners/scanner.cpp
@@ -189,7 +189,12 @@ size_t pesieve::ProcessScanner::scanWorkingSet(ProcessScanReport &pReport) //thr
 {
 	PSAPI_WORKING_SET_INFORMATION wsi_1 = { 0 };
 	BOOL result = QueryWorkingSet(this->processHandle, (LPVOID)&wsi_1, sizeof(PSAPI_WORKING_SET_INFORMATION));
-	if (result == FALSE) {
+	if (result == FALSE && GetLastError() != ERROR_BAD_LENGTH) {
+		/**
+		Allow to proceed on ERROR_BAD_LENGTH.
+		ERROR_BAD_LENGTH may occur if the scanner is 32 bit and running on a 64 bit system.
+		In case of any different error, break.
+		*/
 		throw std::runtime_error("Could not query the working set. ");
 		return 0;
 	}
diff --git a/utils/workingset_enum.cpp b/utils/workingset_enum.cpp
index c3a1cb8eb..b48336e00 100644
--- a/utils/workingset_enum.cpp
+++ b/utils/workingset_enum.cpp
@@ -2,13 +2,19 @@
 
 #include <iostream>
 
+#ifdef _WIN64
+	const ULONGLONG mask = ULONGLONG(-1);
+#else
+	const ULONGLONG mask = DWORD(-1);
+#endif
+
 namespace pesieve {
 	namespace util {
 
 		bool get_next_commited_region(HANDLE processHandle, ULONGLONG start_va, MEMORY_BASIC_INFORMATION &page_info)
 		{
-			while (true) {
-				//std::cout << "Checking: " << std::hex << start_va << " vs " << std::hex << max_va << std::endl;
+			while (start_va < mask) {
+				//std::cout << "Checking: " << std::hex << start_va << std::endl;
 				memset(&page_info, 0, sizeof(MEMORY_BASIC_INFORMATION));
 				SIZE_T out = VirtualQueryEx(processHandle, (LPCVOID)start_va, &page_info, sizeof(page_info));
 				const DWORD error = GetLastError();
@@ -23,15 +29,12 @@ namespace pesieve {
 					std::cerr << "[WARNING] Cannot query the memory region. Error: " << std::dec << error << std::endl;
 					break;
 				}
-				if (error == ERROR_BAD_LENGTH) {
-#ifdef _DEBUG
-					if (sizeof(page_info) != sizeof(MEMORY_BASIC_INFORMATION64)){
-						std::cerr << "[WARNING] Use 64-bit scanner. Error:" << std::dec << error << std::endl;
-					}
-#endif
-					break;
-				}
-				if (out != sizeof(page_info) || error != ERROR_SUCCESS) {
+				/*
+				Allow to proceed on ERROR_BAD_LENGTH, if the filled MEMORY_BASIC_INFORMATION is as expected.
+				(ERROR_BAD_LENGTH may occur if the scanner is 32 bit and running on a 64 bit system.)
+				Otherwise - also on different error - skip.
+				*/
+				if (out != sizeof(page_info) || error != ERROR_BAD_LENGTH) {
 					std::cerr << "[WARNING] Cannot query the memory region. Error: " << std::dec << error << std::endl;
 					start_va += PAGE_SIZE;
 					continue;
@@ -58,12 +61,6 @@ namespace pesieve {
 
 size_t pesieve::util::enum_workingset(HANDLE processHandle, std::set<ULONGLONG> &region_bases)
 {
-#ifdef _WIN64
-	ULONGLONG mask = ULONGLONG(-1);
-#else
-	ULONGLONG mask = DWORD(-1);
-#endif
-
 	region_bases.clear();
 
 	MEMORY_BASIC_INFORMATION page_info = { 0 };