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*/