potisanのプログラミングメモ

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

C++20&WinAPI PEファイルのインポート情報を列挙する

C++ Windows

PEファイル（Windowsの主なEXE、DLLファイル）のインポート情報を列挙するクラスのコードです。std::spanを多用していますが、所々に不正なメモリを読み込むバグの余地は残っていると思います。正常なPEファイルを確認する分には特に問題ありません。

修正履歴

2021/07/22

  1. 戻り値がbool型なのにTRUE、FALSE、{}等を返すものを修正しました。
  2. std::optional<T>の.value()を->演算子または*演算子に置き換えました。
  3. 条件式のミスを修正しました。
  4. ReadOnlyFileHandleAndView、PEFileCoreInfoにコピーコンストラクタとムーブコンストラクタを追加しました。

2021/7/23

  1. PEFileCoreInfoムーブコンストラクタのコンパイルエラーを修正しました。

2021/7/24

  1. PEFileCoreInfoの関数を追加・修正しました。

Source.cpp

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

#include "PEFileImportInfo.hpp"

#include <iostream>

int wmain()
{
    PEFileCoreInfo Info(LR"(C:\Windows\Explorer.exe)");

    auto ImgDosHdr = Info.GetImageDOSHeader();
    auto ImgFileHdr = Info.GetImageFileHeader();
    auto ImgDataDirs = Info.GetImageDataDirectories();
    auto ImgSecHdrs = Info.GetImageSectionHeaders();

    PEFileImportInfo ImportInfo(Info);

    for (auto& ImpDesc : ImportInfo.GetImageImportDescriptors())
    {
        std::cout << ImportInfo.GetName(ImpDesc) << std::endl;
        for (auto& FunctionInfo : ImportInfo.GetFunctionsAuto(ImpDesc))
        {
            auto name = std::get<std::string>(FunctionInfo);
            auto ordinal = std::get<WORD>(FunctionInfo);
            if (!name.empty())
                std::cout << name << ", ";
            else
                std::cout << " #" << ordinal << ", ";
        }
        std::cout << std::endl << std::endl;
    }

    return 0;
}

ReadOnlyFileHandleAndView.hpp

// ReadOnlyFileHandleAndView.hpp
#pragma once

#include <memory>
#include <string>
#include <span>

class ReadOnlyFileHandleAndView final
{
public:
    constexpr ReadOnlyFileHandleAndView() noexcept
    {
    }

    ReadOnlyFileHandleAndView(
        std::wstring_view FileName,
        DWORD ShareMode = FILE_SHARE_READ) noexcept
    {
        Reset(FileName, ShareMode);
    }

    ReadOnlyFileHandleAndView(const ReadOnlyFileHandleAndView&) = delete;

    ReadOnlyFileHandleAndView(ReadOnlyFileHandleAndView&& Source)
        : m_FileHandle(std::move(Source.m_FileHandle))
        , m_FileSize(std::exchange(Source.m_FileSize, 0))
        , m_FileMapping(std::move(Source.m_FileMapping))
        , m_MapViewOfFile(std::move(Source.m_MapViewOfFile))
    {
    }

    ~ReadOnlyFileHandleAndView()
    {
        Clear();
    }

    void Clear() noexcept
    {
        m_MapViewOfFile.reset();
        m_FileMapping.reset();
        m_FileHandle.reset();
    }

    bool IsEmpty() const noexcept
    {
        return !m_FileHandle && !m_FileMapping && !m_MapViewOfFile;
    }

    bool Reset(
        std::wstring_view FileName,
        DWORD ShareMode = FILE_SHARE_READ) noexcept
    {
        m_FileHandle.reset(CreateFileW(FileName.data(),
            GENERIC_READ, ShareMode, nullptr, OPEN_EXISTING, 0, nullptr));
        if (!m_FileHandle) {
            Clear();
            return false;
        }
        DWORD FileSizeHigh = 0;
        DWORD FileSizeLow = ::GetFileSize(m_FileHandle.get(), &FileSizeHigh);
        m_FileSize = FileSizeLow | ((LONGLONG)FileSizeHigh << 32); // LittleEndian
        m_FileMapping.reset(CreateFileMappingW(m_FileHandle.get(),
            nullptr, PAGE_READONLY, 0, 0, nullptr));
        if (!m_FileMapping) {
            Clear();
            return false;
        }
        m_MapViewOfFile.reset(MapViewOfFile(m_FileMapping.get(),
            FILE_MAP_READ, 0, 0, 0));
        if (!m_MapViewOfFile) {
            Clear();
            return false;
        }

        return true;
    }

    HANDLE GetFileHandle() const noexcept
    {
        return m_FileHandle.get();
    }

    LONGLONG GetFileSize() const noexcept
    {
        return m_FileSize;
    }

    HANDLE GetFileMappingHandle() const noexcept
    {
        return m_FileMapping.get();
    }

    LPCVOID GetPointer() const noexcept
    {
        return m_MapViewOfFile.get();
    }

    std::span<const BYTE> GetSpan() const noexcept
    {
        return std::span<const BYTE>(
            reinterpret_cast<LPCBYTE>(m_MapViewOfFile.get()),
            (size_t)m_FileSize);
    }

private:
    struct HANDLEDeleter final
    {
        void operator()(HANDLE h) const noexcept
        {
            CloseHandle(h);
        }
    };
    struct MapViewOfFileHANDLEDeleter final
    {
        void operator()(LPVOID BaseAddress) const noexcept
        {
            UnmapViewOfFile(BaseAddress);
        }
    };
    using unique_handle = std::unique_ptr<std::remove_pointer_t<HANDLE>, HANDLEDeleter>;
    using unique_map_view_of_file = std::unique_ptr<std::remove_pointer_t<HANDLE>, MapViewOfFileHANDLEDeleter>;

    unique_handle m_FileHandle;
    LONGLONG m_FileSize;
    unique_handle m_FileMapping;
    unique_map_view_of_file m_MapViewOfFile;
};

PEFileCoreInfo.hpp

// PEFileCoreInfo
#pragma once

#include "ReadOnlyFileHandleAndView.hpp"

#include <optional>
#include <vector>

class PEFileCoreInfo final
{
public:
    constexpr PEFileCoreInfo() noexcept
    {
    }

    explicit PEFileCoreInfo(std::wstring_view FileName, DWORD ShareMode = FILE_SHARE_READ) noexcept
        : m_HandleAndView(FileName, ShareMode)
    {
    }

    PEFileCoreInfo(const PEFileCoreInfo&) = delete;

    PEFileCoreInfo(PEFileCoreInfo&& Source)
        : m_HandleAndView(std::move(Source.m_HandleAndView))
    {
    }

    ~PEFileCoreInfo()
    {
        Clear();
    }

    void Clear() noexcept
    {
        m_HandleAndView.Clear();
    }

    bool IsEmpty() const noexcept
    {
        return m_HandleAndView.IsEmpty();
    }

    std::span<const BYTE> GetData() const noexcept
    {
        if (IsEmpty()) return {};
        return m_HandleAndView.GetSpan();
    }

    std::span<const BYTE> GetData(size_t Offset, size_t Count = std::dynamic_extent) const noexcept
    {
        if (IsEmpty()) return {};
        return m_HandleAndView.GetSpan().subspan(Offset, Count);
    }

    template <typename T>
    std::optional<std::reference_wrapper<const T>> GetOneDataAs(size_t Offset) const noexcept
    {
        if (IsEmpty()) return {};
        if (GetData().size() - Offset < sizeof(T))
            return {};
        return *reinterpret_cast<const T*>(GetData().subspan(Offset).data());
    }

    template <typename T>
    std::span<const T> GetDataAs() const noexcept
    {
        auto Data = GetData();
        return std::span<const T>(
            reinterpret_cast<const T*>(Data.data()),
            reinterpret_cast<const T*>(Data.data() + Data.size()));
    }

    template <typename T>
    std::span<const T> GetDataByteSizeAs(size_t Offset, size_t ByteSize) const noexcept
    {
        auto Data = GetData(Offset);
        return std::span<const T>(
            reinterpret_cast<const T*>(Data.data()),
            reinterpret_cast<const T*>(Data.data() + ByteSize));
    }

    template <typename T>
    std::span<const T> GetDataTypedCountAs(size_t Offset, size_t TypedCount) const noexcept
    {
        return std::span<const T>(reinterpret_cast<const T*>(GetData(Offset).data()), TypedCount);
    }

    std::optional<std::reference_wrapper<const IMAGE_DOS_HEADER>> GetImageDOSHeader() const noexcept
    {
        auto Hdr = GetOneDataAs<IMAGE_DOS_HEADER>(0);
        if (!Hdr) return {};
        if (Hdr->get().e_magic != IMAGE_DOS_SIGNATURE) return {};
        return Hdr;
    }

    std::optional<DWORD> GetImageNTHeadersSignature() const noexcept
    {
        auto ImgDosHdr = GetImageDOSHeader();
        if (!ImgDosHdr) return {};
        return GetOneDataAs<DWORD>(ImgDosHdr->get().e_lfanew);
    }

    bool IsValidPEFile() const noexcept
    {
        return !IsEmpty() && HasValidImageNTHeadersSignature();
    }

    bool HasValidImageNTHeadersSignature() const noexcept
    {
        auto Sign = GetImageNTHeadersSignature();
        if (!Sign) return false;
        return *GetImageNTHeadersSignature() == IMAGE_NT_SIGNATURE;
    }

    std::optional<std::reference_wrapper<const IMAGE_FILE_HEADER>> GetImageFileHeader() const noexcept
    {
        if (!HasValidImageNTHeadersSignature())
            return {};
        return GetOneDataAs<IMAGE_FILE_HEADER>(
            GetImageDOSHeader()->get().e_lfanew + sizeof(DWORD));
    }

    std::optional<WORD> GetImageOptionalHeaderMagic() const noexcept
    {
        auto FileHdr = GetImageFileHeader();
        if (!FileHdr) return {};
        if (FileHdr->get().SizeOfOptionalHeader < sizeof(WORD)) return {};

        return GetOneDataAs<WORD>(
            GetImageDOSHeader()->get().e_lfanew + sizeof(DWORD) + sizeof(IMAGE_FILE_HEADER));
    }

    bool HasImageOptionalHeader32() const noexcept
    {
        auto FileHdr = GetImageFileHeader();
        if (!FileHdr) return false;
        auto NTHdrMagic = GetImageOptionalHeaderMagic();
        if (!NTHdrMagic) return false;

        return FileHdr->get().SizeOfOptionalHeader == sizeof(IMAGE_OPTIONAL_HEADER32)
            && *NTHdrMagic == IMAGE_NT_OPTIONAL_HDR32_MAGIC;
    }

    bool HasImageOptionalHeader64() const noexcept
    {
        auto FileHdr = GetImageFileHeader();
        if (!FileHdr) return false;
        auto NTHdrMagic = GetImageOptionalHeaderMagic();
        if (!NTHdrMagic) return false;

        return FileHdr->get().SizeOfOptionalHeader == sizeof(IMAGE_OPTIONAL_HEADER64)
            && *NTHdrMagic == IMAGE_NT_OPTIONAL_HDR64_MAGIC;
    }

    std::optional<std::reference_wrapper<const IMAGE_OPTIONAL_HEADER32>> GetImageOptionalHeader32() const noexcept
    {
        if (!HasImageOptionalHeader32()) return {};
        return GetOneDataAs<IMAGE_OPTIONAL_HEADER32>(
            GetImageDOSHeader()->get().e_lfanew + sizeof(DWORD) + sizeof(IMAGE_FILE_HEADER));
    }

    std::optional<std::reference_wrapper<const IMAGE_OPTIONAL_HEADER64>> GetImageOptionalHeader64() const noexcept
    {
        if (!HasImageOptionalHeader64()) return {};
        return GetOneDataAs<IMAGE_OPTIONAL_HEADER64>(
            GetImageDOSHeader()->get().e_lfanew + sizeof(DWORD) + sizeof(IMAGE_FILE_HEADER));
    }

    std::span<const IMAGE_DATA_DIRECTORY> GetImageDataDirectories() const noexcept
    {
        if (HasImageOptionalHeader32())
        {
            auto OptHdr32 = GetImageOptionalHeader32();
            return std::span(OptHdr32->get().DataDirectory, IMAGE_NUMBEROF_DIRECTORY_ENTRIES);
        }
        else if (HasImageOptionalHeader64())
        {
            auto OptHdr64 = GetImageOptionalHeader64();
            return std::span(OptHdr64->get().DataDirectory, IMAGE_NUMBEROF_DIRECTORY_ENTRIES);
        }
        else
        {
            return {};
        }
    }

    std::span<const IMAGE_SECTION_HEADER> GetImageSectionHeaders() const noexcept
    {
        auto DosHdr = GetImageDOSHeader();
        if (!DosHdr) return {};
        auto FileHdr = GetImageFileHeader();
        if (!FileHdr) return {};

        if (HasImageOptionalHeader32())
        {
            return GetDataTypedCountAs<IMAGE_SECTION_HEADER>(
                DosHdr->get().e_lfanew + sizeof(IMAGE_NT_HEADERS32),
                FileHdr->get().NumberOfSections);
        }
        else if (HasImageOptionalHeader64())
        {
            return GetDataTypedCountAs<IMAGE_SECTION_HEADER>(
                DosHdr->get().e_lfanew + sizeof(IMAGE_NT_HEADERS64),
                FileHdr->get().NumberOfSections);
        }
        else
        {
            return {};
        }
    }

    bool Reset(std::wstring_view FileName, DWORD ShareMode = FILE_SHARE_READ)
    {
        m_HandleAndView.Reset(FileName, ShareMode);

        auto Data = GetData();

        auto ImgDosHdr = GetImageDOSHeader();
        if (!ImgDosHdr || ImgDosHdr->get().e_magic != IMAGE_DOS_SIGNATURE) {
            Clear();
            return false;
        }
        auto ImgNTHdrSign = GetImageNTHeadersSignature();
        if (!ImgNTHdrSign || ImgNTHdrSign.value() != IMAGE_NT_SIGNATURE) {
            Clear();
            return false;
        }

        return true;
    }

    std::vector<std::string_view> GetSectionNames() const noexcept
    {
        const auto SecHdrs = GetImageSectionHeaders();
        std::vector<std::string_view> Names;
        Names.reserve(SecHdrs.size());
        for (DWORD i = 0; i < SecHdrs.size(); i++)
        {
            Names.emplace_back(reinterpret_cast<const char*>(SecHdrs[i].Name));
        }
        return std::move(Names);
    }

    std::span<const BYTE> GetSectionData(DWORD Index) const noexcept
    {
        auto SecHdrs = GetImageSectionHeaders();
        if (!(0 <= Index && Index <= SecHdrs.size())) return {};

        return GetSectionData(SecHdrs[Index]);
    }

    std::span<const BYTE> GetSectionData(const IMAGE_SECTION_HEADER& SectionHeader) const noexcept
    {
        return GetData().subspan(SectionHeader.PointerToRawData, SectionHeader.SizeOfRawData);
    }

    std::span<const BYTE> GetSectionDataForImageDataDirectory(DWORD Index) const noexcept
    {
        auto DataDirs = GetImageDataDirectories();
        if (!(0 <= Index && Index <= DataDirs.size()))
            return {};

        auto& DataDir = DataDirs[Index];
        auto SecHdr = FindSectionForVA(DataDir.VirtualAddress, DataDir.Size);
        if (!SecHdr) return {};
        auto& SecHdrRaw = SecHdr->get();
        return GetSectionData(SecHdrRaw).subspan(DataDir.VirtualAddress - SecHdrRaw.VirtualAddress);
    }

    std::optional<std::reference_wrapper<const IMAGE_SECTION_HEADER>> FindSectionForVA(
        DWORD VirtualAddress, DWORD Size = 0) const noexcept
    {
        for (auto& SecHdr : GetImageSectionHeaders())
        {
            if (SecHdr.VirtualAddress <= VirtualAddress
                && VirtualAddress <= SecHdr.VirtualAddress + SecHdr.SizeOfRawData)
            {
                if (VirtualAddress + Size <= SecHdr.VirtualAddress + SecHdr.SizeOfRawData)
                    return SecHdr;
                break;
            }
        }
        return {};
    }

    std::span<const BYTE> GetDataForRVA(size_t RVA) const noexcept
    {
        auto SecHdr = FindSectionForVA(RVA);
        if (!SecHdr) return {};
        auto& SecHdrRaw = SecHdr->get();
        return GetSectionData(SecHdrRaw).subspan(RVA - SecHdrRaw.VirtualAddress);
    }

    template <typename T>
    std::span<const T> GetDataForRVAAs(size_t RVA) const noexcept
    {
        auto Data = GetDataForRVA(RVA);
        return std::span(reinterpret_cast<const T*>(Data.data()), Data.size() / sizeof(T));
    }

    template <typename T>
    std::span<const T> GetDataForRVAAs(ULONGLONG RVA, size_t Count) const noexcept
    {
        auto Data = GetDataForRVA(RVA);
        return std::span(reinterpret_cast<const T*>(Data.data()), Count);
    }

    std::string_view GetNullTerminatedStringFromRVA(ULONGLONG RVA) const noexcept
    {
        auto Data = GetDataForRVA(RVA);
        auto found = std::find(Data.begin(), Data.end(), '\0');
        if (found == Data.end()) return {};
        auto len = std::distance(Data.begin(), found);
        return std::string_view(reinterpret_cast<const char*>(Data.data()), len);
    }

private:
    ReadOnlyFileHandleAndView m_HandleAndView;

};

PEFileImportInfo.hpp

// PEFileImportInfo.hpp
#pragma once

#include "PEFileCoreInfo.hpp"

class PEFileImportInfo final
{
public:
    explicit PEFileImportInfo(const PEFileCoreInfo& Info) noexcept
        : m_Info(Info)
    {
    }

    std::span<const IMAGE_IMPORT_DESCRIPTOR> GetImageImportDescriptors() const noexcept
    {
        auto ImpDirData = m_Info.GetSectionDataForImageDataDirectory(IMAGE_DIRECTORY_ENTRY_IMPORT);
        auto ImpDescs = std::span<const IMAGE_IMPORT_DESCRIPTOR>(
            reinterpret_cast<const IMAGE_IMPORT_DESCRIPTOR*>(ImpDirData.data()),
            reinterpret_cast<const IMAGE_IMPORT_DESCRIPTOR*>(ImpDirData.data() + ImpDirData.size()));
        for (DWORD i = 0; i < std::numeric_limits<DWORD>::max(); i++)
        {
            if (ImpDescs[i].Characteristics == 0)
                return ImpDescs.subspan(0, i - 1);
        }
        return {};
    }

    std::string_view GetName(const IMAGE_IMPORT_DESCRIPTOR& Descriptor) const noexcept
    {
        return m_Info.GetNullTerminatedStringFromRVA(Descriptor.Name);
    }

    std::vector<std::string_view> GetNames() const noexcept
    {
        auto impDescs = GetImageImportDescriptors();
        std::vector<std::string_view> names;
        names.reserve(impDescs.size());
        for (auto& ImpDesc : GetImageImportDescriptors())
        {
            names.emplace_back(GetName(ImpDesc));
        }
        return std::move(names);
    }

    std::vector<std::reference_wrapper<const IMAGE_THUNK_DATA32>> GetThunkData32(
        const IMAGE_IMPORT_DESCRIPTOR& Descriptor) const noexcept
    {
        if (!m_Info.HasImageOptionalHeader32()) return {};

        DWORD ThunkRVA = Descriptor.FirstThunk;
        std::vector<std::reference_wrapper<const IMAGE_THUNK_DATA32>> Result;
        for (DWORD i = 0; i < std::numeric_limits<DWORD>::max(); i++)
        {
            auto pThunk = m_Info.GetDataForRVAAs<IMAGE_THUNK_DATA32>(ThunkRVA, 1).data();
            if (pThunk->u1.Function == 0)
                break;
            Result.emplace_back(*pThunk);

            ThunkRVA += sizeof(DWORD);
        }
        return std::move(Result);
    }

    std::vector<std::reference_wrapper<const IMAGE_THUNK_DATA64>> GetThunkData64(
        const IMAGE_IMPORT_DESCRIPTOR& Descriptor) const noexcept
    {
        if (!m_Info.HasImageOptionalHeader64()) return {};

        DWORD ThunkRVA = Descriptor.FirstThunk;
        std::vector<std::reference_wrapper<const IMAGE_THUNK_DATA64>> Result;
        for (DWORD i = 0; i < std::numeric_limits<DWORD>::max(); i++)
        {
            auto pThunk = m_Info.GetDataForRVAAs<IMAGE_THUNK_DATA64>(ThunkRVA, 1).data();
            if (pThunk->u1.Function == 0)
                break;
            Result.emplace_back(*pThunk);

            ThunkRVA += sizeof(DWORDLONG);
        }
        return std::move(Result);
    }

    std::vector<std::tuple<WORD, std::string>> GetFunctions32(
        const IMAGE_IMPORT_DESCRIPTOR& Descriptor) const noexcept
    {
        auto ThunkData = GetThunkData32(Descriptor);
        std::vector<std::tuple<WORD, std::string>> Names;
        Names.reserve(ThunkData.size());
        for (auto& ThunkData1 : ThunkData)
        {
            if (IMAGE_SNAP_BY_ORDINAL32(ThunkData1.get().u1.Ordinal))
            {
                Names.emplace_back(std::make_tuple<WORD, std::string>(
                    IMAGE_ORDINAL32(ThunkData1.get().u1.Ordinal), {}));
            }
            else
            {
                auto pImpByName = m_Info.GetDataForRVAAs<IMAGE_IMPORT_BY_NAME>(
                    ThunkData1.get().u1.AddressOfData, 1).data();
                Names.emplace_back(std::make_tuple<WORD, std::string>(
                    (int)pImpByName->Hint, pImpByName->Name));
            }
        }
        return std::move(Names);
    }

    std::vector<std::tuple<WORD, std::string>> GetFunctions64(
        const IMAGE_IMPORT_DESCRIPTOR& Descriptor) const noexcept
    {
        auto ThunkData = GetThunkData64(Descriptor);
        std::vector<std::tuple<WORD, std::string>> Names;
        Names.reserve(ThunkData.size());
        for (auto& ThunkData1 : ThunkData)
        {
            if (IMAGE_SNAP_BY_ORDINAL64(ThunkData1.get().u1.Ordinal))
            {
                Names.emplace_back(std::make_tuple<WORD, std::string>(
                    IMAGE_ORDINAL64(ThunkData1.get().u1.Ordinal), {}));
            }
            else
            {
                auto pImpByName = m_Info.GetDataForRVAAs<IMAGE_IMPORT_BY_NAME>(
                    ThunkData1.get().u1.AddressOfData, 1).data();
                Names.emplace_back(std::make_tuple<WORD, std::string>(
                    (WORD)pImpByName->Hint, pImpByName->Name));
            }
        }
        return std::move(Names);
    }

    std::vector<std::tuple<WORD, std::string>> GetFunctionsAuto(
        const IMAGE_IMPORT_DESCRIPTOR& Descriptor) const noexcept
    {
        if (m_Info.HasImageOptionalHeader32())
            return GetFunctions32(Descriptor);
        else if (m_Info.HasImageOptionalHeader64())
            return GetFunctions64(Descriptor);
        else
            return {};
    }

private:
    const PEFileCoreInfo& m_Info;

private:
    PEFileImportInfo() = delete;
    PEFileImportInfo(const PEFileImportInfo&) = delete;
    PEFileImportInfo(PEFileImportInfo&&) = delete;

};