Skip to content

Instantly share code, notes, and snippets.

@yutopp yutopp/

Created Dec 16, 2012
Embed
What would you like to do?
pe_gen
#ifndef LINKER_PE_SECTION_GENERATOR_BASE_HPP
#define LINKER_PE_SECTION_GENERATOR_BASE_HPP
#pragma once
#include <iostream>
#include <string>
#include <unordered_map>
#include <boost/range/adaptors.hpp>
#include "../resource_format/coff.hpp"
#include "../resource_format/ms_import_lib.hpp"
#include "../prototype_sections.hpp"
namespace linker
{
namespace yff = ytl::fileformat;
//
template<typename Derived, typename AddressType, std::size_t Alignment>
class pe_section_genarator_base
{
public:
static std::size_t const alingnment = Alignment;
public:
typedef pe_section_genarator_base base_type;
typedef Derived derived_type;
typedef AddressType address_type;
typedef coff_accessor_type::symbol_type coff_symbol_type;
typedef coff_accessor_type::section_type coff_section_type;
private:
struct relocation_info
{
located_section_info location;
ytl::word_t relocation_type;
};
private:
struct holder_table
{
std::unordered_map<std::string, located_section_info> located_symbol_positions;
std::unordered_multimap<std::string, relocation_info> locate_waiting_tasks;
};
typedef std::unordered_map<std::shared_ptr<void>, holder_table> holder_info;
private:
typedef std::tuple<
std::size_t, // offset of host_section
ytl::word_t, // relocation type
section_id, // target_section id
std::size_t, // offset of target_section
bool // true: add base address, be VA. false: be RVA
> complete_info;
public:
// mapping COFF symbol
void mapping(
coff_holder_shared_pointer const holder,
std::string const& symbol_name,
coff_section_type const& section
)
{
located_section_info const lsi = write_to( holder, symbol_name, section );
solve_relocation_tasks( holder, symbol_name, lsi );
}
// mapping MS Import Library symbol
void mapping(
ms_import_lib_holder_shared_pointer const holder,
std::string const& symbol_name
)
{
write_to( holder, symbol_name );
solve_relocation_tasks( holder, symbol_name );
}
public:
template<typename TargetHolderPtr>
bool register_relocation(
TargetHolderPtr const target_holder,
std::string const& target_symbol_name,
coff_holder_shared_pointer const host_holder,
std::string const& host_symbol_name,
coff_accessor_type::section_type::relocation_type const& reloc
)
{
//<!-- debug begin
std::cout
<< "--- relocation information" << std::endl
<< "VirtualAddress : 0x" << std::hex << reloc.virtual_address << std::endl
<< "Type : 0x" << std::hex << reloc.type << std::endl
<< "SymbolTableIndex : 0x" << std::hex << reloc.symbol_table_index << std::endl
<< std::dec << std::endl;
// debug end-->
auto const host_position = holder_info_[host_holder].located_symbol_positions.find( host_symbol_name );
if ( host_position == holder_info_[host_holder].located_symbol_positions.cend() ) {
throw std::domain_error( "-4 critical error" );; // critical error
}
std::cout
<< "host: " << host_symbol_name << " -> " << target_symbol_name << std::endl
<< "head: " << host_position->second.offset_from_head << std::endl
<< "diff: " << reloc.virtual_address << std::endl;
located_section_info host_location = {
host_position->second.id,
host_position->second.offset_from_head + reloc.virtual_address
};
relocation_info host_relocation_info = { host_location, reloc.type };
auto const target_position = holder_info_[target_holder].located_symbol_positions.find( target_symbol_name );
if ( target_position != holder_info_[target_holder].located_symbol_positions.cend() ) {
// target_symbol has already been located!
located_section_info target_location = {
target_position->second.id,
target_position->second.offset_from_head
};
register_relocation_data( host_relocation_info, target_location );
return true;
} else {
// target_symbol has not been located yet...
// register relocation information
holder_info_[target_holder].locate_waiting_tasks.emplace( target_symbol_name, host_relocation_info );
return false;
}
}
//
void register_section_rva(
std::string const section_name,
address_type const rva
)
{
section_id const id = sections_.generate_section_id( section_name );
register_section_rva_by_id( id, rva );
}
//
void make_section_complete(
std::string const section_name,
address_type const image_base_address
)
{
section_id const id = sections_.generate_section_id( section_name );
make_section_complete_by_id( id, image_base_address );
}
//
void make_all_section_complete(
address_type const image_base_address
)
{
for( auto const& id : get_sections().section_ids_range() ) {
make_section_complete_by_id( id, image_base_address );
}
}
public:
prototype_sections const& get_sections() const
{
return sections_;
}
public:
void create_import_directory_section(
std::string const& section_name
)
{
std::size_t const dll_num = boost::distance( dll_and_name_table_ | boost::adaptors::map_keys );
// Calc Import Library Table Size
std::size_t const import_library_table_num = dll_num + 1/*null terminate*/;
std::size_t const import_library_table_size = import_library_table_num * sizeof( yff::pe::image::import_descriptor );
// Calc ILT/IAT Size and Offset
std::vector<std::size_t> ilt_iat_offsets;
std::size_t ilt_iat_offset = 0;
for( auto const& dll_name : dll_and_name_table_ | boost::adaptors::map_keys ) {
ilt_iat_offsets.emplace_back( ilt_iat_offset );
auto const symbols_range_pair = dll_and_name_table_.equal_range( dll_name );
auto const symbol_num
= boost::distance( boost::make_iterator_range( symbols_range_pair ) | boost::adaptors::map_values )
+ 1/*null terminate*/;
ilt_iat_offset += symbol_num * sizeof( address_type );
}
std::size_t const ilt_iat_size = ilt_iat_offset;
// H/N Table
std::vector<std::size_t> h_n_table_offsets;
prototype_sections::buffer_type h_n_table_buffer;
for( auto const& dll_name : dll_and_name_table_ | boost::adaptors::map_keys ) {
h_n_table_offsets.emplace_back( h_n_table_buffer.size() );
auto const symbols_range_pair = dll_and_name_table_.equal_range( dll_name );
for( auto const& symbol_name : boost::make_iterator_range( symbols_range_pair ) | boost::adaptors::map_values ) {
h_n_table_buffer.binary_expand( ytl::uint16_t( 0 ) ); // Hint
h_n_table_buffer << symbol_name; // Symbol name
h_n_table_buffer.align( 2, 0 ); //
}
}
std::size_t const h_n_table_size = h_n_table_buffer.size();
ytl::dumpbin32( h_n_table_buffer );
// Name
std::vector<std::size_t> name_table_offsets;
prototype_sections::buffer_type name_table_buffer;
for( auto const& dll_name : dll_and_name_table_ | boost::adaptors::map_keys ) {
name_table_offsets.emplace_back( name_table_buffer.size() );
name_table_buffer << dll_name;
}
// Offsets
std::size_t const ilt_offset = import_library_table_size;
std::size_t const iat_offset = ilt_offset + ilt_iat_size;
std::size_t const h_n_table_offset = iat_offset + ilt_iat_size;
std::size_t const name_table_offset = h_n_table_offset + h_n_table_size;
// create the [.idata] section
auto const id = sections_.generate_section_id( section_name );
auto& section = sections_.reference_section( section_name );
auto& buffer = section.buffer;
// Import Directory Table
yff::pe::image::import_descriptor const zero_import_desc = {};
for( std::size_t i=0; i<dll_num; ++i ) {
// yff::pe::image::import_descriptor :: original_first_thunk( dword_t )
add_lazy_relocate_rva(
id,
buffer.size() + sizeof( ytl::uint32_t ) * 0,
yff::coff::image::rel_i386_dir32,//tmp
id,
ilt_offset + ilt_iat_offsets.at( i ) // ILT
);
// yff::pe::image::import_descriptor :: name( dword_t )
add_lazy_relocate_rva(
id,
buffer.size() + sizeof( ytl::uint32_t ) * 3,
yff::coff::image::rel_i386_dir32,//tmp
id,
name_table_offset + name_table_offsets.at( i ) // Name
);
// yff::pe::image::import_descriptor :: name( first_thunk )
add_lazy_relocate_rva(
id,
buffer.size() + sizeof( ytl::uint32_t ) * 4,
yff::coff::image::rel_i386_dir32,//tmp
id,
iat_offset + ilt_iat_offsets.at( i ) // IAT
);
// make space
buffer.binary_expand( zero_import_desc );
}
buffer.binary_expand( zero_import_desc ); // NULL Terminate
// ILT and IAT
for( int i=0; i<2; ++i ) {
std::size_t index = 0;
for( auto const& dll_name : dll_and_name_table_ | boost::adaptors::map_keys ) {
for( auto const& v : boost::make_iterator_range( dll_and_name_table_.equal_range( dll_name ) ) ) {
// pointer to H/N table
add_lazy_relocate_rva(
id,
buffer.size(),
yff::coff::image::rel_i386_dir32,//tmp
id,
h_n_table_offset + h_n_table_offsets.at( index ) // H/N table
);
++index;
// make space
buffer.binary_expand( address_type( 0 ) );
}
// NULL Terminate
buffer.binary_expand( address_type( 0 ) );
}
}
// H/N Table
buffer << h_n_table_buffer;
// Name
buffer << name_table_buffer;
ytl::dumpbin32( buffer );
}
// ---- COFF
private:
located_section_info write_to(
coff_holder_shared_pointer const holder,
std::string const& symbol_name,
coff_section_type const& section
)
{
std::cout
<< "write section !" << std::endl
<< "size: " << section.get_binary().size() << std::endl;
section_id const id = sections_.generate_section_id( section.name() );
auto& buffer = sections_.get_buffer_at( id );
auto const section_offset = buffer.size();
// copy to buffer from symbol's section binary
buffer << section.get_binary();
buffer.align( alingnment );
//
located_section_info li = { id, section_offset };
holder_info_[holder].located_symbol_positions.emplace( symbol_name, li );
std::cout << "~~~?? wtite_to : " << symbol_name << " in " << holder << std::endl;
return li;
}
void solve_relocation_tasks(
coff_holder_shared_pointer const holder,
std::string const& symbol_name,
located_section_info const& lsi
)
{
auto const range_pair = holder_info_[holder].locate_waiting_tasks.equal_range( symbol_name );
for( auto const& v : boost::make_iterator_range( range_pair.first, range_pair.second ) ) {
register_relocation_data( v.second, lsi );
}
holder_info_[holder].locate_waiting_tasks.erase( range_pair.first, range_pair.second );
}
// ---- MS Import Library
private:
void write_to(
ms_import_lib_holder_shared_pointer const holder,
std::string const& symbol_name
)
{
std::cout
<< "write section !" << std::endl
<< "symbol_name : " << symbol_name << std::endl
<< "dll name : " << holder->accessor.get_dll_name() << std::endl
<< "import name : " << holder->accessor.get_import_name() << std::endl;
//
dll_and_name_table_.emplace( holder->accessor.get_dll_name(), holder->accessor.get_import_name() );
}
void solve_relocation_tasks(
ms_import_lib_holder_shared_pointer const holder,
std::string const& symbol_name
)
{
}
private:
void register_relocation_data(
relocation_info const& host,
located_section_info const& target
)
{
add_lazy_relocate_va(
host.location.id,
host.location.offset_from_head,
host.relocation_type,
target.id,
target.offset_from_head
);
}
//
void register_section_rva_by_id(
section_id const host_section_id,
address_type const rva
)
{
if ( rva < 0 ) {
throw std::runtime_error( "rva must be positive integer" );
}
auto& host_section = sections_.get_section_at( host_section_id );
host_section.rva = rva;
}
//
void make_section_complete_by_id(
section_id const host_section_id,
address_type const image_base_address
)
{
auto& host_section = sections_.get_section_at( host_section_id );
auto& host_buffer = host_section.buffer;
auto const& host_rva = host_section.rva; // optional
if ( !host_rva ) {
throw std::runtime_error( "-5 no rva in host" );
}
auto const range_pair = completer_.equal_range( host_section_id );
for( auto const& v : boost::make_iterator_range( range_pair ) | boost::adaptors::map_values ) {
auto const& host_offset = std::get<0>( v );
auto const& relocation_type = std::get<1>( v );
auto const& target_section_id = std::get<2>( v );
auto& target_section = sections_.get_section_at( target_section_id );
auto const& target_rva = target_section.rva; // optional
if ( !target_rva ) {
throw std::runtime_error( "-6 no rva in target" );
}
auto const& target_offset = std::get<3>( v );
bool const is_add_base_address = std::get<4>( v );
//
address_type* const rewite_address
= reinterpret_cast<address_type *>( &host_buffer[host_offset] );
// set address
*rewite_address
= static_cast<derived_type&>( *this ).resolve_address(
is_add_base_address ? image_base_address : 0, // branch RVA and VA
*rewite_address,
*host_rva,
host_offset,
relocation_type,
*target_rva,
target_offset
);
std::cout
<< std::hex
<< "address : 0x" << rewite_address << std::endl
<< "value : " << *rewite_address << std::endl
<< std::dec << std::endl;
}
completer_.erase( range_pair.first, range_pair.second );
}
private:
void add_lazy_relocate_va(
section_id host_section_id, // address of this section
std::size_t host_section_offset, // will be rewrote.
ytl::word_t type_of_relocation,
section_id target_section_id,
std::size_t target_section_offset
)
{
completer_.emplace(
host_section_id,
std::make_tuple(
host_section_offset,
type_of_relocation,
target_section_id,
target_section_offset,
true
)
);
}
void add_lazy_relocate_rva(
section_id host_section_id, // address of this section
std::size_t host_section_offset, // will be rewrote.
ytl::word_t type_of_relocation,
section_id target_section_id,
std::size_t target_section_offset
)
{
completer_.emplace(
host_section_id,
std::make_tuple(
host_section_offset,
type_of_relocation,
target_section_id,
target_section_offset,
false
)
);
}
private:
prototype_sections sections_;
holder_info holder_info_;
std::unordered_multimap<std::string, std::string> dll_and_name_table_;
std::unordered_multimap<section_id, complete_info> completer_;
};
} // namespace linker
#endif /*LINKER_PE_SECTION_GENERATOR_BASE_HPP*/
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.