source: dev/trunk/ab5.0/abdev/ab_common/src/ProcessAndModule.cpp@ 829

#include "stdafx.h"
#include "ProcessAndModule.h"
#include "OSVersion.h"

#include <memory>
#include <psapi.h>
#include <tlhelp32.h>

#pragma comment(lib, "psapi.lib")

namespace ActiveBasic { namespace Common {

struct HandleDeleter
{
    typedef HANDLE pointer;

    void operator ()(pointer h) const
    {
        ::CloseHandle(h);
    }
};

typedef std::unique_ptr<HANDLE, HandleDeleter> UniqueHandle;

boost::filesystem::path GetModuleFilePath(HANDLE hProcess, HMODULE hmod)
{   
    if (GetOSVersion().dwPlatformId == VER_PLATFORM_WIN32_NT)
    {
        WCHAR temp[MAX_PATH] = {};
        if (::GetModuleFileNameExW(hProcess, hmod, temp, MAX_PATH) != 0)
        {
            return temp;
        }
    }
    else
    {
        CHAR temp[MAX_PATH] = {};
        if (::GetModuleFileNameA(hmod, temp, MAX_PATH) != 0)
        {
            return boost::filesystem::path(Jenga::Common::ToWString(temp));
        }
    }
    throw std::runtime_error("GetModuleFileName or GetModuleFileNameEx failed");
}

namespace {

struct Toolhelp32Function
{
    decltype(::CreateToolhelp32Snapshot)* pCreateToolhelp32Snapshot;
    decltype(::Process32First)* pProcess32First;
    decltype(::Process32Next)* pProcess32Next;
    decltype(::Module32First)* pModule32First;
    decltype(::Module32Next)* pModule32Next;
};

Toolhelp32Function const& GetToolhelp32Function(HMODULE hmodKernel)
{
    static Toolhelp32Function const tf =
    {
        reinterpret_cast<decltype(::CreateToolhelp32Snapshot)*>(::GetProcAddress(hmodKernel, "CreateToolhelp32Snapshot")),
        reinterpret_cast<decltype(::Process32First)*>(::GetProcAddress(hmodKernel, "Process32First")),
        reinterpret_cast<decltype(::Process32Next)*>(::GetProcAddress(hmodKernel, "Process32Next")),
        reinterpret_cast<decltype(::Module32First)*>(::GetProcAddress(hmodKernel, "Module32First")),
        reinterpret_cast<decltype(::Module32Next)*>(::GetProcAddress(hmodKernel, "Module32Next")),
    };
    return tf;
}

UniqueHandle CreateToolhelp32SnapshotHelper(Toolhelp32Function const& tf, DWORD flags, DWORD processId = 0)
{
    UniqueHandle h(tf.pCreateToolhelp32Snapshot(flags, processId));
    if (h.get() == INVALID_HANDLE_VALUE)
    {
        h.release();
        return UniqueHandle();
    }
    return h;
}

} // unnamed namespace

std::vector<DWORD> GetProcesses()
{
    auto const& tf = GetToolhelp32Function(::GetModuleHandle(TEXT("kernel32")));
    if (tf.pCreateToolhelp32Snapshot)
    {
        UniqueHandle h = CreateToolhelp32SnapshotHelper(tf, TH32CS_SNAPPROCESS);
        if (h == nullptr)
        {
            return std::vector<DWORD>();
        }
        std::vector<DWORD> processes;
        PROCESSENTRY32 pe = {sizeof pe};
        if (tf.pProcess32First(h.get(), &pe))
        {
            do
            {
                processes.push_back(pe.th32ProcessID);
            } while (tf.pProcess32Next(h.get(), &pe));
        }
        return processes;
    }
    else
    {
        std::vector<DWORD> processes(8192);
        DWORD cbNeeded = 0;
        if (EnumProcesses(processes.data(), processes.size() * sizeof (DWORD), &cbNeeded))
        {
            processes.resize(cbNeeded / sizeof (DWORD));
            return processes;
        }
        return std::vector<DWORD>();
    }
}

boost::filesystem::path GetExecutableModulePath(DWORD processId)
{
    if (ActiveBasic::Common::Is9x())
    {
        auto const& tf = GetToolhelp32Function(::GetModuleHandle(TEXT("kernel32")));
        UniqueHandle h = CreateToolhelp32SnapshotHelper(tf, TH32CS_SNAPMODULE, processId);
        if (h == nullptr)
        {
            return boost::filesystem::path();
        }
        MODULEENTRY32 me = {sizeof me};
        if (tf.pModule32First(h.get(), &me))
        {
            return boost::filesystem::path(me.szExePath);
        }
        else
        {
            return boost::filesystem::path();
        }
    }
    else
    {
        UniqueHandle hProcess(::OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, processId));
        // PSAPIが使える環境はそちらを優先。
        WCHAR filename[MAX_PATH] = {};
        ::GetModuleFileNameExW(hProcess.get(), nullptr, filename, MAX_PATH);
        return boost::filesystem::path(filename);
    }
}

std::vector<HMODULE> GetModules(DWORD processId)
{
    auto const& tf = GetToolhelp32Function(::GetModuleHandle(TEXT("kernel32")));
    if (tf.pCreateToolhelp32Snapshot)
    {
        UniqueHandle h = CreateToolhelp32SnapshotHelper(tf, TH32CS_SNAPMODULE, processId);
        if (h == nullptr)
        {
            return std::vector<HMODULE>();
        }
        std::vector<HMODULE> modules;
        MODULEENTRY32 me = {sizeof me};
        if (tf.pModule32First(h.get(), &me))
        {
            do
            {
                modules.push_back(me.hModule);
            } while (tf.pModule32Next(h.get(), &me));
        }
        return modules;
    }
    else
    {
        UniqueHandle hProcess(::OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, processId));
        const DWORD MAX_MODULE = 1024;
        std::vector<HMODULE> module(MAX_MODULE);
        DWORD cbReturned = 0;
        EnumProcessModules(hProcess.get(), module.data(), sizeof (HMODULE) * MAX_MODULE, &cbReturned);
        module.resize(cbReturned / sizeof (HMODULE));
        return module;
    }
}

}}