添加项目文件。

WinSysCore/ProcessModuleTracker.cpp

@@ -0,0 +1,213 @@
+#include "pch.h"
+#include "ProcessModuleTracker.h"
+#include <TlHelp32.h>
+#include <Psapi.h>
+#include <ImageHlp.h>
+#include "Helpers.h"
+
+#pragma comment(lib,"imagehlp")
+
+using namespace WinSys;
+
+struct ProcessModuleTracker::Impl {
+	std::vector<std::shared_ptr<ModuleInfo>> _modules, _newModules, _unloadedModules;
+	std::unordered_map<std::wstring, std::shared_ptr<ModuleInfo>> _moduleMap;
+	DWORD _pid;
+	wil::unique_handle _handle;
+	BOOL _isWow64;
+
+	explicit Impl(DWORD pid) :_pid(pid) {
+		_handle.reset(::OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ | SYNCHRONIZE, FALSE, pid));
+		if (_handle)
+			::IsWow64Process(_handle.get(), &_isWow64);
+	}
+
+	explicit Impl(HANDLE hProcess) :_handle(hProcess) {
+		::IsWow64Process(_handle.get(), &_isWow64);
+	}
+
+	const std::vector<std::shared_ptr<ModuleInfo>>& GetModules() const {
+		return _modules;
+	}
+
+	const std::vector<std::shared_ptr<ModuleInfo>>& GetNewModules() const {
+		return _newModules;
+	}
+
+	const std::vector<std::shared_ptr<ModuleInfo>>& GetUnloadedModules() const {
+		return _unloadedModules;
+	}
+
+	std::shared_ptr<ModuleInfo> FillModule(const MODULEENTRY32& me) {
+		auto mi = std::make_shared<ModuleInfo>();
+		mi->Name = me.szModule;
+		mi->Path = me.szExePath;
+		mi->Base = (void*)me.hModule;
+		mi->ModuleSize = me.modBaseSize;
+		return mi;
+	}
+
+	std::shared_ptr<ModuleInfo> FillModule(const MEMORY_BASIC_INFORMATION& mbi) {
+		auto mi = std::make_shared<ModuleInfo>();
+		mi->ModuleSize = 0;
+		WCHAR name[MAX_PATH];
+		mi->ImageBase = 0;
+		if (::GetMappedFileName(_handle.get(), mbi.AllocationBase, name, _countof(name))) {
+			mi->Path = Helpers::GetDosNameFromNtName(name);
+			mi->Name = ::wcsrchr(name, L'\\') + 1;
+			BYTE buffer[1 << 12];
+			if (::ReadProcessMemory(_handle.get(), mbi.BaseAddress, buffer, sizeof(buffer), nullptr)) {
+				auto nt = ::ImageNtHeader(buffer);
+				if (nt) {
+					auto machine = nt->FileHeader.Machine;
+					if (machine == IMAGE_FILE_MACHINE_ARM || machine == IMAGE_FILE_MACHINE_I386) {
+						auto oh = (IMAGE_OPTIONAL_HEADER32*)&nt->OptionalHeader;
+						mi->ImageBase = UlongToPtr(oh->ImageBase);
+						mi->Characteristics = (DllCharacteristics)oh->DllCharacteristics;
+					}
+					else {
+						mi->ImageBase = (PVOID)nt->OptionalHeader.ImageBase;
+						mi->Characteristics = (DllCharacteristics)nt->OptionalHeader.DllCharacteristics;
+					}
+				}
+			}
+		}
+		mi->Base = mbi.AllocationBase;
+		mi->Type = mbi.Type == MEM_MAPPED ? MapType::Data : MapType::Image;
+		return mi;
+	}
+
+	uint32_t EnumModules() {
+		return _handle ? EnumModulesWithVirtualQuery() : EnumModulesWithToolHelp();
+	}
+
+	uint32_t EnumModulesWithVirtualQuery() {
+		bool first = _modules.empty();
+		if (first) {
+			_modules.reserve(128);
+			_moduleMap.reserve(128);
+			_newModules.reserve(8);
+			_unloadedModules.reserve(8);
+		}
+		else {
+			_newModules.clear();
+			_unloadedModules.clear();
+		}
+
+		auto existing = _moduleMap;
+
+		MEMORY_BASIC_INFORMATION mbi;
+		const BYTE* address = nullptr;
+		ModuleInfo* pmi = nullptr;
+		while (::VirtualQueryEx(_handle.get(), address, &mbi, sizeof(mbi)) > 0) {
+			if (mbi.State == MEM_COMMIT && mbi.Type != MEM_PRIVATE) {
+				if (mbi.AllocationBase == mbi.BaseAddress) {
+					auto mi = FillModule(mbi);
+					pmi = mi.get();
+					auto key = mi->Name + L"#" + std::to_wstring((size_t)mi->Base);
+					if (first) {
+						_moduleMap.insert({ key,mi });
+						_modules.push_back(std::move(mi));
+					}
+					else {
+						auto it = _moduleMap.find(key);
+						if (it == _moduleMap.end()) {
+							// new module
+							_modules.push_back(mi);
+							_newModules.push_back(mi);
+							_moduleMap.insert({ key,mi });
+						}
+						else {
+							existing.erase(key);
+						}
+					}
+					pmi->ModuleSize += (ULONG)mbi.RegionSize;
+				}
+			}
+			address += mbi.RegionSize;
+		}
+
+		for (auto& [key, mi] : existing) {
+			_moduleMap.erase(key);
+			_unloadedModules.push_back(mi);
+		}
+
+		return static_cast<uint32_t>(_modules.size());
+	}
+
+	uint32_t EnumModulesWithToolHelp() {
+		wil::unique_handle hSnapshot(::CreateToolhelp32Snapshot(TH32CS_SNAPMODULE | TH32CS_SNAPMODULE32,_pid));
+		if (!hSnapshot)
+			return 0;
+
+		bool first = _modules.empty();
+		if (first) {
+			_modules.reserve(128);
+			_newModules.reserve(8);
+			_unloadedModules.reserve(8);
+		}
+		else {
+			_newModules.clear();
+			_unloadedModules.clear();
+		}
+
+		MODULEENTRY32 me;
+		me.dwSize = sizeof(me);
+		::Module32First(hSnapshot.get(), &me);
+
+		auto existing = _moduleMap;
+		do
+		{
+			if (first) {
+				auto mi = FillModule(me);
+				_moduleMap.insert({ mi->Path,mi });
+				_modules.push_back(std::move(mi));
+			}
+			else {
+				auto it = _moduleMap.find(me.szExePath);
+				if (it == _moduleMap.end()) {
+					// new module
+					auto mi = FillModule(me);
+					_modules.push_back(mi);
+					_newModules.push_back(mi);
+				}
+				else {
+					existing.erase(me.szExePath);
+				}
+			}
+		} while (::Module32Next(hSnapshot.get(),&me));
+
+		for (auto& [key, mi] : existing)
+			_unloadedModules.push_back(mi);
+
+		return static_cast<uint32_t>(_modules.size());
+	}
+};
+
+WinSys::ProcessModuleTracker::~ProcessModuleTracker() = default;
+
+WinSys::ProcessModuleTracker::ProcessModuleTracker(DWORD pid) :_impl(new Impl(pid)) {
+}
+
+WinSys::ProcessModuleTracker::ProcessModuleTracker(HANDLE hProcess) : _impl(new Impl(hProcess)) {
+}
+
+uint32_t WinSys::ProcessModuleTracker::EnumModules() {
+	return _impl->EnumModules();
+}
+
+const std::vector<std::shared_ptr<ModuleInfo>>& WinSys::ProcessModuleTracker::GetModules() const {
+	return _impl->GetModules();
+}
+
+const std::vector<std::shared_ptr<ModuleInfo>>& WinSys::ProcessModuleTracker::GetNewModules() const {
+	return _impl->GetNewModules();
+}
+
+const std::vector<std::shared_ptr<ModuleInfo>>& WinSys::ProcessModuleTracker::GetUnloadedModules() const {
+	return _impl->GetUnloadedModules();
+}
+
+bool ProcessModuleTracker::IsRunning() const{
+	return ::WaitForSingleObject(_impl->_handle.get(), 0) == WAIT_TIMEOUT;
+}