portable PE64 header utils

pe64.hpp

#pragma once

namespace os
{
   enum e_image_signature: const std::uint16_t
   {
      dos_magic = 0x5a4d,  // 'MZ'
      nts_magic = 0x4550,  // 'PE00'
      opt_magic = 0x020b
   };

   struct image_dos_header_t
   {
      std::uint16_t m_magic;     // Magic number
      std::uint16_t m_cblp;      // Bytes on last page of file
      std::uint16_t m_cp;        // Pages in file
      std::uint16_t m_crlc;      // Relocations
      std::uint16_t m_cparhdr;   // Size of header in paragraphs
      std::uint16_t m_minalloc;  // Minimum extra paragraphs needed
      std::uint16_t m_maxalloc;  // Maximum extra paragraphs needed
      std::uint16_t m_ss;        // Initial (relative) SS value
      std::uint16_t m_sp;        // Initial SP value
      std::uint16_t m_csum;      // Checksum
      std::uint16_t m_ip;        // Initial IP value
      std::uint16_t m_cs;        // Initial (relative) CS value
      std::uint16_t m_lfarlc;    // File address of relocation table
      std::uint16_t m_ovno;      // Overlay number
      std::uint16_t m_res0[4];   // Reserved words
      std::uint16_t m_oemid;     // OEM identifier (for m_oeminfo)
      std::uint16_t m_oeminfo;   // OEM information; m_oemid specific
      std::uint16_t m_res1[10];  // Reserved words
      std::int32_t m_lfanew;     // File address of new exe header

      [[ nodiscard ]]
      bool is_valid( )
      {
         return m_magic == e_image_signature::dos_magic;
      }
   };

   struct image_nt_headers_t
   {
      std::uint32_t m_signature;

      std::uint16_t m_machine;
      std::uint16_t m_number_of_sections;
      std::uint32_t m_time_date_stamp;
      std::uint32_t m_pointer_to_symbol_table;
      std::uint32_t m_number_of_symbols;
      std::int16_t m_size_of_optional_header;
      std::int16_t m_characteristics;

      std::uint16_t m_magic;
      std::uint8_t m_major_linker_version;
      std::uint8_t m_minor_linker_version;
      std::uint32_t m_size_of_code;
      std::uint32_t m_size_of_initialized_data;
      std::uint32_t m_size_of_uninitialized_data;
      std::uint32_t m_address_of_entry_point;
      std::uint32_t m_base_of_code;
      std::uint64_t m_image_base;
      std::uint32_t m_section_alignment;
      std::uint32_t m_file_alignment;
      std::uint16_t m_major_operating_system_version;
      std::uint16_t m_minor_operating_system_version;
      std::uint16_t m_major_image_version;
      std::uint16_t m_minor_image_version;
      std::uint16_t m_major_subsystem_version;
      std::uint16_t m_minor_subsystem_version;
      std::uint32_t m_win32_version_value;
      std::uint32_t m_size_of_image;
      std::uint32_t m_size_of_headers;
      std::uint32_t m_check_sum;
      std::uint16_t m_subsystem;
      std::uint16_t m_dll_characteristics;
      std::uint64_t m_size_of_stack_reserve;
      std::uint64_t m_size_of_stack_commit;
      std::uint64_t m_size_of_heap_reserve;
      std::uint64_t m_size_of_heap_commit;
      std::uint32_t m_loader_flags;
      std::uint32_t m_number_of_rva_and_sizes;

      struct image_data_directory_t
      {
         std::uint32_t m_virtual_address;
         std::uint32_t m_size;
      };

      image_data_directory_t m_export;          // Export Directory
      image_data_directory_t m_import;          // Import Directory
      image_data_directory_t m_resource;        // Resource Directory
      image_data_directory_t m_exception;       // Exception Directory
      image_data_directory_t m_security;        // Security Directory
      image_data_directory_t m_base_reloc;      // Base Relocation Table
      image_data_directory_t m_debug;           // Debug Directory
      image_data_directory_t m_architecture;    // Architecture Specific Data
      image_data_directory_t m_global_ptr;      // RVA of GP
      image_data_directory_t m_tls;             // TLS Directory
      image_data_directory_t m_load_config;     // Load Configuration Directory
      image_data_directory_t m_bound_import;    // Bound Import Directory in headers
      image_data_directory_t m_iat;             // Import Address Table
      image_data_directory_t m_delay_import;    // Delay Load Import Descriptors
      image_data_directory_t m_com_descriptor;  // COM Runtime descriptor
      image_data_directory_t m_unused;

      [[ nodiscard ]]
      bool is_valid( )
      {
         return m_signature == e_image_signature::nts_magic && m_magic == e_image_signature::opt_magic;
      }
   };

   struct image_import_descriptor_t
   {
      union
      {
         std::uint32_t m_characteristics;       // 0 for terminating null import descriptor
         std::uint32_t m_original_first_thunk;  // RVA to original unbound IAT (PIMAGE_THUNK_DATA)
      };

      std::uint32_t m_time_date_stamp;          // 0 if not bound
      std::uint32_t m_forwarder_chain;          // -1 if no forwarders
      std::uint32_t m_name;
      std::uint32_t m_first_thunk;              // RVA to IAT (if bound this IAT has actual addresses)
   };

   struct image_base_relocation_t
   {
      std::uint32_t m_virtual_address;
      std::uint32_t m_size_of_block;
   };

   struct image_import_by_name_t
   {
      std::uint16_t m_hint;
      char m_name[1];
   };

   struct image_export_directory_t
   {
      std::uint32_t m_characteristics;
      std::uint32_t m_time_date_stamp;
      std::int16_t m_major_version;
      std::int16_t m_minor_version;
      std::uint32_t m_name;
      std::uint32_t m_base;
      std::uint32_t m_number_of_functions;
      std::uint32_t m_number_of_names;
      std::uint32_t m_address_of_functions;        // RVA from base of image
      std::uint32_t m_address_of_names;            // RVA from base of image
      std::uint32_t m_address_of_names_ordinals;   // RVA from base of image
   };

   struct image_section_header_t
   {
      char m_name[8];
      union
      {
         std::uint32_t m_physical_address;
         std::uint32_t m_virtual_size;
      };
      std::uint32_t m_virtual_address;
      std::uint32_t m_size_of_raw_data;
      std::uint32_t m_pointer_to_raw_data;
      std::uint32_t m_pointer_to_relocations;
      std::uint32_t m_pointer_to_line_numbers;
      std::int16_t m_number_of_relocations;
      std::int16_t m_number_of_line_numbers;
      std::uint32_t m_characteristics;
   };

   void for_each_section( auto callback, std::uint8_t* image_ptr )
   {
      auto dos_header{ ( os::image_dos_header_t* )( image_ptr ) };
      auto nt_headers{ ( os::image_nt_headers_t* )( image_ptr + dos_header->m_lfanew ) };
      if ( !dos_header->is_valid( )
        || !nt_headers->is_valid( ) )
         return;

      auto section_ctx{ ( os::image_section_header_t* )( ( std::uint8_t* )( &nt_headers->m_magic ) + nt_headers->m_size_of_optional_header ) };
      if ( !section_ctx->m_name
        || !section_ctx->m_pointer_to_raw_data )
         return;

      for ( std::int16_t idx{}; idx < nt_headers->m_number_of_sections; idx++, section_ctx++ ) 
      {
         if ( !callback( idx, section_ctx ) ) break;
      }
   }

   void for_each_export( auto callback, std::uint8_t* image_ptr )
   {
      auto dos_header{ ( os::image_dos_header_t* )( image_ptr ) };
      auto nt_headers{ ( os::image_nt_headers_t* )( image_ptr + dos_header->m_lfanew ) };
      if ( !dos_header->is_valid( )
        || !nt_headers->is_valid( ) )
         return;

      if ( nt_headers->m_export.m_size )
      {
         auto exp_dir{ ( os::image_export_directory_t* )( image_ptr + nt_headers->m_export.m_virtual_address ) };
         if ( !exp_dir->m_address_of_functions 
           || !exp_dir->m_address_of_names 
           || !exp_dir->m_address_of_names_ordinals )
            return;
      
         auto name_ptr{ ( std::uint32_t* )( image_ptr + exp_dir->m_address_of_names ) };
         auto func_ptr{ ( std::uint32_t* )( image_ptr + exp_dir->m_address_of_functions ) };
         auto ords_ptr{ ( std::uint16_t* )( image_ptr + exp_dir->m_address_of_names_ordinals ) };

         for ( std::uint32_t idx{}; idx < exp_dir->m_number_of_names; idx++ )
         {
            if ( !callback( idx, name_ptr, func_ptr, ords_ptr ) ) break;
         }
      }
   }

   void for_each_import( auto callback, std::uint8_t* image_ptr )
   {
      auto dos_header{ ( os::image_dos_header_t* )( image_ptr ) };
      auto nt_headers{ ( os::image_nt_headers_t* )( image_ptr + dos_header->m_lfanew ) };
      if ( !dos_header->is_valid( )
        || !nt_headers->is_valid( ) )
         return;

      if ( nt_headers->m_import.m_size )
      {
         auto imp_dir{ ( os::image_import_descriptor_t* )( image_ptr + nt_headers->m_import.m_virtual_address ) };
         if ( !imp_dir->m_first_thunk 
           || !imp_dir->m_original_first_thunk )
            return;

         while ( imp_dir->m_name )
         {
            auto name_ptr{ ( std::uint64_t* )( image_ptr + imp_dir->m_original_first_thunk ) };
            auto func_ptr{ ( std::uint64_t* )( image_ptr + imp_dir->m_first_thunk ) };

            while ( *name_ptr )
            {
               if ( !callback( name_ptr, func_ptr ) ) return;

               name_ptr++;
               func_ptr++;
            }
            imp_dir++;
         }
      }
   }

   template< class type_t >
   [[ nodiscard ]]
   type_t get_image_section( std::uint8_t* image_ptr, std::uint32_t section_name_hash )
   {
      type_t section_ptr{};

      for_each_section( [&]( std::uint32_t idx, os::image_section_header_t* section_ctx )
      {
         if ( _fnv1( section_ctx->m_name ) == section_name_hash )
         {
            section_ptr = ( type_t )( image_ptr + section_ctx->m_pointer_to_raw_data );
            return false;
         }

         return true;
      }, image_ptr );

      return section_ptr;
   }

   template< class type_t >
   [[ nodiscard ]]
   type_t get_image_export( std::uint8_t* image_ptr, std::uint32_t export_name_hash )
   {
      type_t export_ptr{};

      for_each_export( [&]( std::uint32_t idx, std::uint32_t* name_ptr, std::uint32_t* func_ptr, std::uint16_t* ords_ptr )
      {
         auto cur_name{ image_ptr + name_ptr[idx] };
         auto cur_func{ image_ptr + func_ptr[ords_ptr[idx]] };
         if ( !cur_name || !cur_func )
            return true;

         if ( _fnv1( ( const char* )( cur_name ) ) == export_name_hash )
         {
            export_ptr = ( type_t )( cur_func );
            return false;
         }
         
         return true;
      }, image_ptr );
      
      return export_ptr;
   }

   template< class type_t >
   [[ nodiscard ]]
   type_t get_image_import( std::uint8_t* image_ptr, std::uint32_t import_name_hash )
   {
      type_t import_ptr{};

      for_each_import( [&]( std::uint64_t* name_ptr, std::uint64_t* func_ptr )
      {
         auto cur_imp{ ( os::image_import_by_name_t* )( image_ptr + *name_ptr ) };
         if ( !cur_imp->m_name )
            return true;
         
         if ( _fnv1( cur_imp->m_name ) == import_name_hash )
         {
            import_ptr = ( type_t )( *func_ptr );
            return false;
         }

         return true;
      }, image_ptr );

      return import_ptr;
   }
}

#define _align( addr, align ) ( ( ( addr ) + ( align ) - 1 ) & ~( ( align ) - 1 ) )