potisanのプログラミングメモ

趣味のプログラマーがプログラミング関係で気になったことや調べたことをいつでも忘れられるようにメモするブログです。はてなブログ無料版なので記事の上の方はたぶん広告です。記事中にも広告挿入されるみたいです。

C++20&Win API&WIL Process Snapshot APIでプロセスのハンドル型・オブジェクト名を列挙する

Process Snapshot APIを使ってプロセスに所属するハンドルの型やオブジェクト名を列挙するサンプルコードです。

#include <algorithm>
#include <filesystem>
#include <string>
#include <memory>
#include <vector>

#define STRICT
#define NOMINMAX
#include <Windows.h>
#include <Psapi.h>

#include "wil/resource.h"

template <DWORD InitSize = 100, DWORD ExpandingSize = 100>
std::vector<DWORD> GetAllProcessIDs()
{
    std::vector<DWORD> processIds;

    for (DWORD size = InitSize;; size += ExpandingSize)
    {
        processIds.resize(size);
        DWORD cbNeeded;
        DWORD cbSize = size * sizeof(DWORD);
        if (!EnumProcesses(processIds.data(), cbSize, &cbNeeded))
        {
            return {};
        }
        if (cbNeeded < cbSize)
        {
            processIds.resize(cbNeeded / sizeof(DWORD));
            processIds.shrink_to_fit();
            return processIds;
        }
    }
}

template <DWORD InitSize = 100, DWORD ExpandingSize = 100>
std::wstring QueryFullProcessImageNameW(HANDLE processHandle, DWORD flags)
{
    std::wstring s;
    for (DWORD size = InitSize;; size += ExpandingSize)
    {
        s.resize(size);
        DWORD cch = size;
        if (QueryFullProcessImageNameW(processHandle, flags, s.data(), &cch))
        {
            s.resize(cch);
            s.shrink_to_fit();
            return s;
        }
        if (GetLastError() != ERROR_INSUFFICIENT_BUFFER)
        {
            return {};
        }
    }
}

std::vector<DWORD> FindProcessesByPathEquivalent(std::filesystem::path path)
{
    std::vector<DWORD> ids;
    std::ranges::copy_if(GetAllProcessIDs(), std::back_inserter(ids),
        [&path](DWORD processId)
        {
            wil::unique_process_handle processHandle(
                OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, processId));
            std::filesystem::path processPath(QueryFullProcessImageNameW(processHandle.get(), 0));
            if (processPath.empty()) return path.empty();
            if (path.empty()) return processPath.empty();
            return std::filesystem::equivalent(processPath, path);
        });
    return ids;
}

std::optional<DWORD> FindOneProcessByPathEquivalent(std::filesystem::path path)
{
    auto ids = FindProcessesByPathEquivalent(path);
    if (ids.size() != 1) return {};
    return ids[0];
}

#include <ProcessSnapshot.h>

// wil::unique_anyを参考にしたHPSS管理クラス
class unique_hpss {
private:
    HANDLE process_handle;
    HPSS hpss;
public:
    using pointer = HPSS;

    constexpr unique_hpss() noexcept
        : process_handle(nullptr), hpss(nullptr)
    {
    }

    constexpr unique_hpss(HANDLE process_handle, HPSS hpss) noexcept
        : process_handle(process_handle), hpss(hpss)
    {
    }

    constexpr ~unique_hpss() noexcept
    {
        reset();
    }

    unique_hpss(unique_hpss const&) = delete;
    unique_hpss& operator=(unique_hpss const&) = delete;

    constexpr unique_hpss(unique_hpss&& other) noexcept
    {
        process_handle = std::exchange(other.process_handle, nullptr);
        hpss = std::exchange(other.hpss, nullptr);
    }
    constexpr unique_hpss& operator=(unique_hpss&& other) noexcept
    {
        if (this != std::addressof(other))
        {
            process_handle = std::exchange(other.process_handle, nullptr);
            hpss = std::exchange(other.hpss, nullptr);
        }
        return *this;
    }
    constexpr unique_hpss& operator=(std::nullptr_t) noexcept
    {
        reset();
        return *this;
    }

    constexpr HANDLE get_process_handle() const noexcept { return process_handle; }
    constexpr HPSS get() const noexcept { return hpss; }
    constexpr HPSS get_both(HANDLE& process_handle) const noexcept
    {
        process_handle = this->process_handle;
        return hpss;
    }
    constexpr HPSS* put(HANDLE process_handle) noexcept
    {
        reset();
        this->process_handle = process_handle;
        return &hpss;
    }
    constexpr void reset() noexcept
    {
        if (hpss == nullptr)
            return;
        if (PssFreeSnapshot(process_handle, hpss) == ERROR_SUCCESS)
        {
            process_handle = nullptr;
            hpss = nullptr;
        }
    }
    constexpr void swap(unique_hpss& other) noexcept
    {
        auto self(std::move(*this));
        operator=(std::move(other));
        other = std::move(self);
    }
    constexpr explicit operator bool() const noexcept
    {
        return process_handle != nullptr && hpss != nullptr;
    }
};

using unique_pss_walk_marker = wil::unique_any<HPSSWALK, decltype(PssWalkMarkerFree), &PssWalkMarkerFree>;

constexpr std::wstring wstring_not_null(const wchar_t* s) noexcept
{
    return s != nullptr ? s : L"";
}

#include <iostream>
#include <format>

namespace wilutil
{
    inline auto OpenProcess(
        DWORD dwDesiredAccess,
        BOOL bInheritHandle,
        DWORD dwProcessId) noexcept
    {
        return wil::unique_process_handle(::OpenProcess(dwDesiredAccess, bInheritHandle, dwProcessId));
    }
}

int main()
{
    // Explorer.exeのプロセスIDを取得する。
    auto processId = FindOneProcessByPathEquivalent(L"C:\\Windows\\Explorer.exe");
    if (!processId) return -1;

    auto processHandle = wilutil::OpenProcess(PROCESS_ALL_ACCESS, FALSE, *processId);

    unique_hpss hpss;
    auto err = PssCaptureSnapshot(
        processHandle.get(),
        PSS_CAPTURE_HANDLES | PSS_CAPTURE_HANDLE_NAME_INFORMATION | PSS_CAPTURE_HANDLE_TYPE_SPECIFIC_INFORMATION,
        0,
        hpss.put(processHandle.get()));

    unique_pss_walk_marker hwm;
    err = PssWalkMarkerCreate(nullptr, hwm.put());

    for (PSS_HANDLE_ENTRY entry; ;)
    {
        if (PssWalkSnapshot(hpss.get(), PSS_WALK_HANDLES, hwm.get(), &entry, sizeof(entry)) != ERROR_SUCCESS)
        {
            break;
        }

        auto typeName = wstring_not_null(entry.TypeName);
        auto objectName = wstring_not_null(entry.ObjectName);
        std::wcout
            << std::format(L"{}, {}", typeName, objectName)
            << std::endl;
    }

    return 0;
}