workingset_enum.cpp

Go to the documentation of this file.

#include "workingset_enum.h"
 
#include <iostream>
 
#include <psapi.h>
#pragma comment(lib,"psapi.lib")
 
#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)
        {
            const SIZE_T page_info_size = sizeof(MEMORY_BASIC_INFORMATION);
 
            while (start_va < mask) {
                //std::cout << "Checking: " << std::hex << start_va << std::endl;
                memset(&page_info, 0, page_info_size);
                const SIZE_T out = VirtualQueryEx(processHandle, (LPCVOID)start_va, &page_info, page_info_size);
                const bool is_read = (out == page_info_size) ? true : false;
                const DWORD error = is_read ? ERROR_SUCCESS : GetLastError();
                if (error == ERROR_INVALID_PARAMETER) {
                    //nothing more to read
#ifdef _DEBUG
                    std::cout << "Nothing more to read: " << std::hex << start_va << std::endl;
#endif
                    break;
                }
                if (error == ERROR_ACCESS_DENIED) {
                    std::cerr << "[ERROR] Cannot query the memory region. " << std::hex << start_va << " Error: " << std::dec << error << std::endl;
                    break;
                }
                if (!is_read) {
                    // on any other error:
                    std::cerr << "[WARNING] Cannot query the memory region. " << std::hex<< start_va << " Error: " << std::dec << error << ", skipping the page..." << std::endl;
                    start_va += PAGE_SIZE;
                    continue;
                }
                if ((page_info.State & MEM_FREE) || (page_info.State & MEM_COMMIT) == 0) {
                    if (page_info.RegionSize != 0) {
                        //std::cout << "Free:  " << std::hex << start_va << " RegionSize:" << page_info.RegionSize << std::endl;
                        start_va += page_info.RegionSize;
                        continue;
                    }
                }
                if (page_info.RegionSize == 0) {
                    start_va += PAGE_SIZE;
                    continue;
                }
                //std::cout << "Commited:  " << std::hex << start_va << " RegionSize: " << page_info.RegionSize << " Err: " << std::dec << error << std::endl;
                return true;
            }
            return false;
        }
 
    };
};
 
size_t pesieve::util::enum_workingset(HANDLE processHandle, std::set<mem_region_info> &region_bases)
{
    region_bases.clear();
 
    MEMORY_BASIC_INFORMATION page_info = { 0 };
    ULONGLONG next_va = 0;
    while (get_next_commited_region(processHandle, next_va, page_info))
    {
        ULONGLONG base = (ULONGLONG)page_info.BaseAddress;
        next_va = base + page_info.RegionSize; //end of the region
        mem_region_info curr_info((ULONGLONG)page_info.AllocationBase, base, page_info.RegionSize);
 
        if (region_bases.find(curr_info) != region_bases.end()) {
            // don't let it stuck on adding the same region over and over again
            break;
        }
        region_bases.insert(curr_info);
    }
    return region_bases.size();
}
 
DWORD pesieve::util::count_workingset_entries(HANDLE processHandle)
{
    DWORD number_of_entries = 1;
    DWORD buffer_size = sizeof(PSAPI_WORKING_SET_INFORMATION) + (number_of_entries * sizeof(PSAPI_WORKING_SET_BLOCK));
    PSAPI_WORKING_SET_INFORMATION* buffer = reinterpret_cast<PSAPI_WORKING_SET_INFORMATION*>(calloc(1, buffer_size));
    if (!buffer) {
        return 0; //this should not happen
    }
    DWORD res = QueryWorkingSet(processHandle, buffer, buffer_size);
    if (res == FALSE && GetLastError() == ERROR_BAD_LENGTH) {
        // ERROR_BAD_LENGTH is normal: we didn't provide the buffer that could fit all the entries
        res = TRUE;
    }
    number_of_entries = static_cast<DWORD>(buffer->NumberOfEntries);
    free(buffer); buffer = NULL;
 
    if (!res) {
        return 0;
    }
#ifdef _DEBUG
    std::cout << "Number of entries: " << std::dec << number_of_entries << std::endl;
#endif
    return number_of_entries;
}