Untitled

#include <cstdint>
#include <memory>

#include <pe/pe_parser.h>

#define TRY(expr)                               \
  ({                                            \
    const auto TRY_expr = (expr);               \
    if (!TRY_expr.has_value()) {                \
      return std::unexpected(TRY_expr.error()); \
    }                                           \
    TRY_expr.value();                           \
  })

#define TRY_OPT(expr, error)          \
  ({                                  \
    const auto TRY_OPT_expr = (expr); \
    if (!TRY_OPT_expr.has_value()) {  \
      return std::unexpected(error);  \
    }                                 \
    TRY_OPT_expr.value();             \
  })

std::expected<void, ParserError> PEParser::parse_file_headers() {
  // Read the DOS header and verify that it's correct.
  m_fs.read(reinterpret_cast<char *>(&m_parsed->m_dos), sizeof(DosHeader));
  if (m_parsed->m_dos.m_magic != kDosSignature) {
    return std::unexpected(ParserError::kInvalidDosSignature);
  }

  // Read the NT header.
  m_fs.seekg(m_parsed->m_dos.m_lfanew);
  m_fs.read(reinterpret_cast<char *>(&m_parsed->m_nt), sizeof(NtHeader));
  if (m_parsed->m_nt.m_signature != kNtSignature) {
    return std::unexpected(ParserError::kInvalidNtSignature);
  }

  // Read the optional NT header (it's not optional at all, thanks Microsoft.)
  const auto old_pos{m_fs.tellg()};
  BaseOptional base_opt{};

  m_fs.read(reinterpret_cast<char *>(&base_opt), sizeof(BaseOptional));
  m_fs.seekg(old_pos);

  // Verify that it's either a 32- or 64-bit PE file.
  if (base_opt.m_magic == kOptionalHeaderMagic32) {
    auto &new_opt{m_parsed->m_optional.emplace<Optional32>()};
    m_fs.read(reinterpret_cast<char *>(&new_opt), sizeof(Optional32));
  } else if (base_opt.m_magic == kOptionalHeaderMagic64) {
    auto &new_opt{m_parsed->m_optional.emplace<Optional64>()};
    m_fs.read(reinterpret_cast<char *>(&new_opt), sizeof(Optional64));
  } else {
    return std::unexpected(ParserError::kInvalidOptionalMagic);
  }

  // Read section headers...
  m_fs.seekg(old_pos + static_cast<ptrdiff_t>(m_parsed->m_nt.m_size_of_optional_header));
  m_parsed->m_sections.resize(m_parsed->m_nt.m_number_of_sections);

  for (auto &sect : m_parsed->m_sections) {
    m_fs.read(reinterpret_cast<char *>(&sect.m_header), sizeof(SectionHeader));
    const auto data_size{sect.m_header.m_size_of_raw_data};
    if (data_size == 0) {
      continue;
    }

    // ... and their data too :)
    const auto old_pos{m_fs.tellg()};
    sect.m_data.resize(data_size);
    m_fs.seekg(sect.m_header.m_ptr_to_raw_data);
    m_fs.read(reinterpret_cast<char *>(sect.m_data.data()), data_size);
    m_fs.seekg(old_pos);
  }

  return {};
}

std::expected<void, ParserError> PEParser::parse_relocations() {
  // Get the base relocation directory, need to convert from RVA to file offset.
  const auto &reloc_dir{m_parsed->get_data_directory(kDirBaseReloc)};
  m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(reloc_dir.m_rva), ParserError::kInvalidRva));

  while (true) {
    // Keep reading relocation blocks until we hit one that's invalid. There is
    // always an invalid one at the end :^)
    BaseRelocation reloc{};
    m_fs.read(reinterpret_cast<char *>(&reloc), sizeof(BaseRelocation));
    if (reloc.m_virtual_address == 0 || reloc.m_size_of_block < sizeof(BaseRelocation)) {
      break;
    }

    // Read all base relocations in this relocation block.
    auto &reloc_block{m_parsed->m_relocs.emplace_back()};
    reloc_block.m_rva = reloc.m_virtual_address;
    reloc_block.m_relocs.resize((reloc.m_size_of_block - sizeof(BaseRelocation)) / sizeof(uint16_t));

    for (auto &base_reloc : reloc_block.m_relocs) {
      m_fs.read(reinterpret_cast<char *>(&base_reloc), sizeof(uint16_t));
    }
  }

  return {};
}

std::expected<void, ParserError> PEParser::parse_import_descriptors() {
  // Get the import directory, need to convert from RVA to file offset.
  const auto &import_dir{m_parsed->get_data_directory(kDirImport)};
  m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(import_dir.m_rva), ParserError::kInvalidRva));

  while (true) {
    ImportDescriptor import_desc{};
    m_fs.read(reinterpret_cast<char *>(&import_desc), sizeof(ImportDescriptor));
    const auto next_pos{m_fs.tellg()};
    if (import_desc.m_name == 0) {
      break;
    }

    ParsedImportThunk thunk{};

    // Read the import name, need to convert from RVA to file offset. This is null terminated,
    // so we need to use a stack buffer first.
    std::array<char, kMaxPathSize> module_name{};
    m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(import_desc.m_name), ParserError::kInvalidRva));
    m_fs.get(module_name.data(), module_name.size() - 1, '\0');

    thunk.m_module.assign_range(module_name);

    // Read all import thunks and the imports descriptors they point to.
    m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(import_desc.m_first_thunk), ParserError::kInvalidRva));
    if (m_parsed->get_optional_magic() == kOptionalHeaderMagic32) {
      TRY(fill_import_thunk32(thunk));
    } else if (m_parsed->get_optional_magic() == kOptionalHeaderMagic64) {
      TRY(fill_import_thunk64(thunk));
    }

    m_parsed->m_imports.emplace_back(thunk);
    m_fs.seekg(next_pos);
  }

  return {};
}

std::expected<void, ParserError> PEParser::fill_import_thunk32(ParsedImportThunk &thunk) {
  while (true) {
    ThunkData32 thunk_data{};
    m_fs.read(reinterpret_cast<char *>(&thunk_data), sizeof(ThunkData32));
    const auto next_pos{m_fs.tellg()};
    if (thunk_data.m_address_of_data == 0) {
      break;
    }

    ImportByName import_desc{};
    m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(thunk_data.m_address_of_data), ParserError::kInvalidRva));
    m_fs.read(reinterpret_cast<char *>(&import_desc), sizeof(uint16_t));
    m_fs.get(import_desc.m_name.data(), import_desc.m_name.size() - 1, '\0');

    ParsedImportByNameDescriptor parsed_desc{};
    parsed_desc.m_name.assign_range(import_desc.m_name);
    parsed_desc.m_offset = next_pos - static_cast<std::streampos>(sizeof(ThunkData32));

    thunk.m_imports.emplace_back(parsed_desc);
    m_fs.seekg(next_pos);
  }

  return {};
}

std::expected<void, ParserError> PEParser::fill_import_thunk64(ParsedImportThunk &thunk) {
  while (true) {
    ThunkData64 thunk_data{};
    m_fs.read(reinterpret_cast<char *>(&thunk_data), sizeof(ThunkData64));
    const auto next_pos{m_fs.tellg()};
    if (thunk_data.m_address_of_data == 0) {
      break;
    }

    ImportByName import_desc{};
    m_fs.seekg(TRY_OPT(m_parsed->rva_to_offset(thunk_data.m_address_of_data), ParserError::kInvalidRva));
    m_fs.read(reinterpret_cast<char *>(&import_desc), sizeof(uint16_t));
    m_fs.get(import_desc.m_name.data(), import_desc.m_name.size() - 1, '\0');

    ParsedImportByNameDescriptor parsed_desc{};
    parsed_desc.m_name.assign_range(import_desc.m_name);
    parsed_desc.m_offset = next_pos - static_cast<std::streampos>(sizeof(ThunkData64));

    thunk.m_imports.emplace_back(parsed_desc);
    m_fs.seekg(next_pos);
  }

  return {};
}

std::expected<SharedPEImage, ParserError> PEParser::parse(const std::filesystem::path &file_path) {
  m_fs.open(file_path, std::ios::binary);

  if (!m_fs.is_open()) {
    return std::unexpected(ParserError::kInvalidFile);
  }

  m_parsed = std::make_shared<PEImage>();

  TRY(parse_file_headers());
  TRY(parse_import_descriptors());
  TRY(parse_relocations());

  return m_parsed;
}