C++ pe_base::set_directory_size方法代码示例

//Resources rebuilder
//resource_directory - root resource directory
//resources_section - section where resource directory will be placed (must be attached to PE image)
//offset_from_section_start - offset from resources_section raw data start
//resource_directory is non-constant, because it will be sorted
//save_to_pe_headers - if true, new resource directory information will be saved to PE image headers
//auto_strip_last_section - if true and resources are placed in the last section, it will be automatically stripped
//number_of_id_entries and number_of_named_entries for resource directories are recalculated and not used
const image_directory rebuild_resources(pe_base& pe, resource_directory& info, section& resources_section, uint32_t offset_from_section_start, bool save_to_pe_header, bool auto_strip_last_section)
{
    //Check that resources_section is attached to this PE image
    if(!pe.section_attached(resources_section))
        throw pe_exception("Resource section must be attached to PE file", pe_exception::section_is_not_attached);
    
    //Check resource directory correctness
    if(info.get_entry_list().empty())
        throw pe_exception("Empty resource directory", pe_exception::incorrect_resource_directory);
    
    uint32_t aligned_offset_from_section_start = pe_utils::align_up(offset_from_section_start, sizeof(uint32_t));
    uint32_t needed_size_for_structures = aligned_offset_from_section_start - offset_from_section_start; //Calculate needed size for resource tables and data
    uint32_t needed_size_for_strings = 0;
    uint32_t needed_size_for_data = 0;

    calculate_resource_data_space(info, aligned_offset_from_section_start, needed_size_for_structures, needed_size_for_strings);

    {
        uint32_t current_data_pos = aligned_offset_from_section_start + needed_size_for_structures + needed_size_for_strings;
        calculate_resource_data_space(info, needed_size_for_structures, needed_size_for_strings, needed_size_for_data, current_data_pos);
    }

    uint32_t needed_size = needed_size_for_structures + needed_size_for_strings + needed_size_for_data;

    //Check if resources_section is last one. If it's not, check if there's enough place for resource data
    if(&resources_section != &*(pe.get_image_sections().end() - 1) && 
        (resources_section.empty() || pe_utils::align_up(resources_section.get_size_of_raw_data(), pe.get_file_alignment())
        < needed_size + aligned_offset_from_section_start))
        throw pe_exception("Insufficient space for resource directory", pe_exception::insufficient_space);

    std::string& raw_data = resources_section.get_raw_data();

    //This will be done only if resources_section is the last section of image or for section with unaligned raw length of data
    if(raw_data.length() < needed_size + aligned_offset_from_section_start)
        raw_data.resize(needed_size + aligned_offset_from_section_start); //Expand section raw data

    uint32_t current_structures_pos = aligned_offset_from_section_start;
    uint32_t current_strings_pos = current_structures_pos + needed_size_for_structures;
    uint32_t current_data_pos = current_strings_pos + needed_size_for_strings;
    rebuild_resource_directory(pe, resources_section, info, current_structures_pos, current_data_pos, current_strings_pos, aligned_offset_from_section_start);
    
    //Adjust section raw and virtual sizes
    pe.recalculate_section_sizes(resources_section, auto_strip_last_section);

    image_directory ret(pe.rva_from_section_offset(resources_section, aligned_offset_from_section_start), needed_size);

    //If auto-rewrite of PE headers is required
    if(save_to_pe_header)
    {
        pe.set_directory_rva(image_directory_entry_resource, ret.get_rva());
        pe.set_directory_size(image_directory_entry_resource, ret.get_size());
    }

    return ret;
}

//imports - bound imported modules list
//imports_section - section where export directory will be placed (must be attached to PE image)
//offset_from_section_start - offset from imports_section raw data start
//save_to_pe_headers - if true, new bound import directory information will be saved to PE image headers
//auto_strip_last_section - if true and bound imports are placed in the last section, it will be automatically stripped
const image_directory rebuild_bound_imports(pe_base& pe, const bound_import_module_list& imports, section& imports_section, uint32_t offset_from_section_start, bool save_to_pe_header, bool auto_strip_last_section)
{
    //Check that exports_section is attached to this PE image
    if(!pe.section_attached(imports_section))
        throw pe_exception("Bound import section must be attached to PE file", pe_exception::section_is_not_attached);

    uint32_t directory_pos = pe_utils::align_up(offset_from_section_start, sizeof(uint32_t));
    uint32_t needed_size = sizeof(image_bound_import_descriptor) /* Ending null descriptor */;
    uint32_t needed_size_for_strings = 0;

    //Calculate needed size for bound import data
    for(bound_import_module_list::const_iterator it = imports.begin(); it != imports.end(); ++it)
    {
        const bound_import& import = *it;
        needed_size += sizeof(image_bound_import_descriptor);
        needed_size_for_strings += static_cast<uint32_t>((*it).get_module_name().length()) + 1 /* nullbyte */;

        const bound_import::ref_list& refs = import.get_module_ref_list();
        for(bound_import::ref_list::const_iterator ref_it = refs.begin(); ref_it != refs.end(); ++ref_it)
        {
            needed_size_for_strings += static_cast<uint32_t>((*ref_it).get_module_name().length()) + 1 /* nullbyte */;
            needed_size += sizeof(image_bound_forwarder_ref);
        }
    }

    needed_size += needed_size_for_strings;

    //Check if imports_section is last one. If it's not, check if there's enough place for bound import data
    if(&imports_section != &*(pe.get_image_sections().end() - 1) &&
            (imports_section.empty() || pe_utils::align_up(imports_section.get_size_of_raw_data(), pe.get_file_alignment()) < needed_size + directory_pos))
        throw pe_exception("Insufficient space for bound import directory", pe_exception::insufficient_space);

    std::string& raw_data = imports_section.get_raw_data();

    //This will be done only if imports_section is the last section of image or for section with unaligned raw length of data
    if(raw_data.length() < needed_size + directory_pos)
        raw_data.resize(needed_size + directory_pos); //Expand section raw data

    uint32_t current_pos_for_structures = directory_pos;
    uint32_t current_pos_for_strings = current_pos_for_structures + needed_size - needed_size_for_strings;

    for(bound_import_module_list::const_iterator it = imports.begin(); it != imports.end(); ++it)
    {
        const bound_import& import = *it;
        image_bound_import_descriptor descriptor;
        descriptor.NumberOfModuleForwarderRefs = static_cast<uint16_t>(import.get_module_ref_list().size());
        descriptor.OffsetModuleName = static_cast<uint16_t>(current_pos_for_strings - directory_pos);
        descriptor.TimeDateStamp = import.get_timestamp();

        memcpy(&raw_data[current_pos_for_structures], &descriptor, sizeof(descriptor));
        current_pos_for_structures += sizeof(descriptor);

        size_t length = import.get_module_name().length() + 1 /* nullbyte */;
        memcpy(&raw_data[current_pos_for_strings], import.get_module_name().c_str(), length);
        current_pos_for_strings += static_cast<uint32_t>(length);

        const bound_import::ref_list& refs = import.get_module_ref_list();
        for(bound_import::ref_list::const_iterator ref_it = refs.begin(); ref_it != refs.end(); ++ref_it)
        {
            const bound_import_ref& ref = *ref_it;
            image_bound_forwarder_ref ref_descriptor = {0};
            ref_descriptor.OffsetModuleName = static_cast<uint16_t>(current_pos_for_strings - directory_pos);
            ref_descriptor.TimeDateStamp = ref.get_timestamp();

            memcpy(&raw_data[current_pos_for_structures], &ref_descriptor, sizeof(ref_descriptor));
            current_pos_for_structures += sizeof(ref_descriptor);

            length = ref.get_module_name().length() + 1 /* nullbyte */;
            memcpy(&raw_data[current_pos_for_strings], ref.get_module_name().c_str(), length);
            current_pos_for_strings += static_cast<uint32_t>(length);
        }
    }

    //Adjust section raw and virtual sizes
    pe.recalculate_section_sizes(imports_section, auto_strip_last_section);

    image_directory ret(pe.rva_from_section_offset(imports_section, directory_pos), needed_size);

    //If auto-rewrite of PE headers is required
    if(save_to_pe_header)
    {
        pe.set_directory_rva(image_directory_entry_bound_import, ret.get_rva());
        pe.set_directory_size(image_directory_entry_bound_import, ret.get_size());
    }

    return ret;
}

//.........这里部分代码省略.........
	image_config_section_struct.DeCommitFreeBlockThreshold = static_cast<typename PEClassType::BaseSize>(info.get_decommit_free_block_threshold());
	image_config_section_struct.DeCommitTotalFreeThreshold = static_cast<typename PEClassType::BaseSize>(info.get_decommit_total_free_threshold());
	image_config_section_struct.MaximumAllocationSize = static_cast<typename PEClassType::BaseSize>(info.get_max_allocation_size());
	image_config_section_struct.VirtualMemoryThreshold = static_cast<typename PEClassType::BaseSize>(info.get_virtual_memory_threshold());
	image_config_section_struct.ProcessHeapFlags = info.get_process_heap_flags();
	image_config_section_struct.ProcessAffinityMask = static_cast<typename PEClassType::BaseSize>(info.get_process_affinity_mask());
	image_config_section_struct.CSDVersion = info.get_service_pack_version();
	image_config_section_struct.EditList = static_cast<typename PEClassType::BaseSize>(info.get_edit_list_va());
	image_config_section_struct.SecurityCookie = static_cast<typename PEClassType::BaseSize>(info.get_security_cookie_va());
	image_config_section_struct.SEHandlerCount = static_cast<typename PEClassType::BaseSize>(info.get_se_handler_rvas().size());
	

	if(write_se_handlers)
	{
		if(info.get_se_handler_rvas().empty())
		{
			write_se_handlers = false;
			image_config_section_struct.SEHandlerTable = 0;
		}
		else
		{
			typename PEClassType::BaseSize va;
			pe.rva_to_va(pe.rva_from_section_offset(image_config_section, current_pos_of_se_handlers), va);
			image_config_section_struct.SEHandlerTable = va;
		}
	}
	else
	{
		image_config_section_struct.SEHandlerTable = static_cast<typename PEClassType::BaseSize>(info.get_se_handler_table_va());
	}

	if(write_lock_prefixes)
	{
		if(info.get_lock_prefix_rvas().empty())
		{
			write_lock_prefixes = false;
			image_config_section_struct.LockPrefixTable = 0;
		}
		else
		{
			typename PEClassType::BaseSize va;
			pe.rva_to_va(pe.rva_from_section_offset(image_config_section, current_pos_of_lock_prefixes), va);
			image_config_section_struct.LockPrefixTable = va;
		}
	}
	else
	{
		image_config_section_struct.LockPrefixTable = static_cast<typename PEClassType::BaseSize>(info.get_lock_prefix_table_va());
	}

	//Write image config section
	memcpy(&raw_data[image_config_data_pos], &image_config_section_struct, sizeof(image_config_section_struct));

	if(write_se_handlers)
	{
		//Sort SE Handlers list
		image_config_info::se_handler_list sorted_list = info.get_se_handler_rvas();
		std::sort(sorted_list.begin(), sorted_list.end());

		//Write SE Handlers table
		for(image_config_info::se_handler_list::const_iterator it = sorted_list.begin(); it != sorted_list.end(); ++it)
		{
			uint32_t se_handler_rva = *it;
			memcpy(&raw_data[current_pos_of_se_handlers], &se_handler_rva, sizeof(se_handler_rva));
			current_pos_of_se_handlers += sizeof(se_handler_rva);
		}
	}

	if(write_lock_prefixes)
	{
		//Write Lock Prefixes VA list
		for(image_config_info::lock_prefix_rva_list::const_iterator it = info.get_lock_prefix_rvas().begin(); it != info.get_lock_prefix_rvas().end(); ++it)
		{
			typename PEClassType::BaseSize lock_prefix_va;
			pe.rva_to_va(*it, lock_prefix_va);
			memcpy(&raw_data[current_pos_of_lock_prefixes], &lock_prefix_va, sizeof(lock_prefix_va));
			current_pos_of_lock_prefixes += sizeof(lock_prefix_va);
		}

		{
			//Ending null VA
			typename PEClassType::BaseSize lock_prefix_va = 0;
			memcpy(&raw_data[current_pos_of_lock_prefixes], &lock_prefix_va, sizeof(lock_prefix_va));
		}
	}

	//Adjust section raw and virtual sizes
	pe.recalculate_section_sizes(image_config_section, auto_strip_last_section);

	image_directory ret(pe.rva_from_section_offset(image_config_section, image_config_data_pos), sizeof(typename PEClassType::ConfigStruct));

	//If auto-rewrite of PE headers is required
	if(save_to_pe_header)
	{
		pe.set_directory_rva(image_directory_entry_load_config, ret.get_rva());
		pe.set_directory_size(image_directory_entry_load_config, ret.get_size());
	}

	return ret;
}

//.........这里部分代码省略.........
		dir.Characteristics = info.get_characteristics();
		dir.MajorVersion = info.get_major_version();
		dir.MinorVersion = info.get_minor_version();
		dir.TimeDateStamp = info.get_timestamp();
		dir.NumberOfFunctions = max_ordinal - ordinal_base + 1;
		dir.NumberOfNames = number_of_names;
		dir.Base = ordinal_base;
		dir.AddressOfFunctions = pe.rva_from_section_offset(exports_section, current_pos_of_function_addresses);
		dir.AddressOfNameOrdinals = pe.rva_from_section_offset(exports_section, current_pos_of_function_name_ordinals);
		dir.AddressOfNames = pe.rva_from_section_offset(exports_section, current_pos_of_function_names_rvas);
		dir.Name = pe.rva_from_section_offset(exports_section, directory_pos + sizeof(image_export_directory));

		//Save it
		memcpy(&raw_data[directory_pos], &dir, sizeof(dir));
	}

	//Sve library name
	memcpy(&raw_data[directory_pos + sizeof(image_export_directory)], info.get_name().c_str(), info.get_name().length() + 1);

	//A map to sort function names alphabetically
	typedef std::map<std::string, uint16_t> funclist; //function name; function name ordinal
	funclist funcs;

	uint32_t last_ordinal = ordinal_base;
	//Enumerate all exported functions
	for(exported_functions_list::const_iterator it = exports.begin(); it != exports.end(); ++it)
	{
		const exported_function& func = (*it);

		//If we're skipping some ordinals...
		if(func.get_ordinal() > last_ordinal)
		{
			//Fill this function RVAs data with zeros
			uint32_t len = sizeof(uint32_t) * (func.get_ordinal() - last_ordinal - 1);
			if(len)
			{
				memset(&raw_data[current_pos_of_function_addresses], 0, len);
				current_pos_of_function_addresses += len;
			}
			
			//Save last encountered ordinal
			last_ordinal = func.get_ordinal();
		}
		
		//If function is named, save its name ordinal and name in sorted alphabetically order
		if(func.has_name())
			funcs.insert(std::make_pair(func.get_name(), static_cast<uint16_t>(func.get_ordinal() - ordinal_base))); //Calculate name ordinal

		//If function is forwarded to another DLL
		if(func.is_forwarded())
		{
			//Write its forwarded name and its RVA
			uint32_t function_rva = pe.rva_from_section_offset(exports_section, current_pos_of_function_forwards);
			memcpy(&raw_data[current_pos_of_function_addresses], &function_rva, sizeof(function_rva));
			current_pos_of_function_addresses += sizeof(function_rva);

			memcpy(&raw_data[current_pos_of_function_forwards], func.get_forwarded_name().c_str(), func.get_forwarded_name().length() + 1);
			current_pos_of_function_forwards += static_cast<uint32_t>(func.get_forwarded_name().length() + 1);
		}
		else
		{
			//Write actual function RVA
			uint32_t function_rva = func.get_rva();
			memcpy(&raw_data[current_pos_of_function_addresses], &function_rva, sizeof(function_rva));
			current_pos_of_function_addresses += sizeof(function_rva);
		}
	}
	
	//Enumerate sorted function names
	for(funclist::const_iterator it = funcs.begin(); it != funcs.end(); ++it)
	{
		//Save function name RVA
		uint32_t function_name_rva = pe.rva_from_section_offset(exports_section, current_pos_of_function_names);
		memcpy(&raw_data[current_pos_of_function_names_rvas], &function_name_rva, sizeof(function_name_rva));
		current_pos_of_function_names_rvas += sizeof(function_name_rva);

		//Save function name
		memcpy(&raw_data[current_pos_of_function_names], (*it).first.c_str(), (*it).first.length() + 1);
		current_pos_of_function_names += static_cast<uint32_t>((*it).first.length() + 1);

		//Save function name ordinal
		uint16_t name_ordinal = (*it).second;
		memcpy(&raw_data[current_pos_of_function_name_ordinals], &name_ordinal, sizeof(name_ordinal));
		current_pos_of_function_name_ordinals += sizeof(name_ordinal);
	}
	
	//Adjust section raw and virtual sizes
	pe.recalculate_section_sizes(exports_section, auto_strip_last_section);
	
	image_directory ret(pe.rva_from_section_offset(exports_section, directory_pos), needed_size);

	//If auto-rewrite of PE headers is required
	if(save_to_pe_header)
	{
		pe.set_directory_rva(image_directory_entry_export, ret.get_rva());
		pe.set_directory_size(image_directory_entry_export, ret.get_size());
	}

	return ret;
}

//Simple relocations rebuilder
//To keep PE file working, don't remove any of existing relocations in
//relocation_table_list returned by a call to get_relocations() function
//auto_strip_last_section - if true and relocations are placed in the last section, it will be automatically stripped
//offset_from_section_start - offset from the beginning of reloc_section, where relocations data will be situated
//If save_to_pe_header is true, PE header will be modified automatically
const image_directory rebuild_relocations(pe_base& pe, const relocation_table_list& relocs, section& reloc_section, uint32_t offset_from_section_start, bool save_to_pe_header, bool auto_strip_last_section)
{
    //Check that reloc_section is attached to this PE image
    if(!pe.section_attached(reloc_section))
        throw pe_exception("Relocations section must be attached to PE file", pe_exception::section_is_not_attached);

    uint32_t current_reloc_data_pos = pe_utils::align_up(offset_from_section_start, sizeof(uint32_t));

    uint32_t needed_size = current_reloc_data_pos - offset_from_section_start; //Calculate needed size for relocation tables
    uint32_t size_delta = needed_size;

    uint32_t start_reloc_pos = current_reloc_data_pos;

    //Enumerate relocation tables
    for(relocation_table_list::const_iterator it = relocs.begin(); it != relocs.end(); ++it)
    {
        needed_size += static_cast<uint32_t>((*it).get_relocations().size() * sizeof(uint16_t) /* relocations */ + sizeof(image_base_relocation) /* table header */);
        //End of each table will be DWORD-aligned
        if((start_reloc_pos + needed_size - size_delta) % sizeof(uint32_t))
            needed_size += sizeof(uint16_t); //Align it with IMAGE_REL_BASED_ABSOLUTE relocation
    }

    //Check if reloc_section is last one. If it's not, check if there's enough place for relocations data
    if(&reloc_section != &*(pe.get_image_sections().end() - 1) &&
            (reloc_section.empty() || pe_utils::align_up(reloc_section.get_size_of_raw_data(), pe.get_file_alignment()) < needed_size + current_reloc_data_pos))
        throw pe_exception("Insufficient space for relocations directory", pe_exception::insufficient_space);

    std::string& raw_data = reloc_section.get_raw_data();

    //This will be done only if reloc_section is the last section of image or for section with unaligned raw length of data
    if(raw_data.length() < needed_size + current_reloc_data_pos)
        raw_data.resize(pe_utils::align_up(needed_size + current_reloc_data_pos, pe.get_file_alignment())); //Expand section raw data

    //Enumerate relocation tables
    for(relocation_table_list::const_iterator it = relocs.begin(); it != relocs.end(); ++it)
    {
        //Create relocation table header
        image_base_relocation reloc;
        reloc.VirtualAddress = (*it).get_rva();
        const relocation_table::relocation_list& reloc_list = (*it).get_relocations();
        reloc.SizeOfBlock = static_cast<uint32_t>(sizeof(image_base_relocation) + sizeof(uint16_t) * reloc_list.size());
        if((reloc_list.size() * sizeof(uint16_t)) % sizeof(uint32_t)) //If we must align end of relocation table
            reloc.SizeOfBlock += sizeof(uint16_t);

        memcpy(&raw_data[current_reloc_data_pos], &reloc, sizeof(reloc));
        current_reloc_data_pos += sizeof(reloc);

        //Enumerate relocations in table
        for(relocation_table::relocation_list::const_iterator r = reloc_list.begin(); r != reloc_list.end(); ++r)
        {
            //Save relocations
            uint16_t reloc_value = (*r).get_item();
            memcpy(&raw_data[current_reloc_data_pos], &reloc_value, sizeof(reloc_value));
            current_reloc_data_pos += sizeof(reloc_value);
        }

        if(current_reloc_data_pos % sizeof(uint32_t)) //If end of table is not DWORD-aligned
        {
            memset(&raw_data[current_reloc_data_pos], 0, sizeof(uint16_t)); //Align it with IMAGE_REL_BASED_ABSOLUTE relocation
            current_reloc_data_pos += sizeof(uint16_t);
        }
    }

    image_directory ret(pe.rva_from_section_offset(reloc_section, start_reloc_pos), needed_size - size_delta);

    //Adjust section raw and virtual sizes
    pe.recalculate_section_sizes(reloc_section, auto_strip_last_section);

    //If auto-rewrite of PE headers is required
    if(save_to_pe_header)
    {
        pe.set_directory_rva(image_directory_entry_basereloc, ret.get_rva());
        pe.set_directory_size(image_directory_entry_basereloc, ret.get_size());

        pe.clear_characteristics_flags(image_file_relocs_stripped);
        pe.set_dll_characteristics(pe.get_dll_characteristics() | image_dllcharacteristics_dynamic_base);
    }

    return ret;
}