1 files changed, 299 insertions, 0 deletions
diff --git a/tools/pe_bliss/pe_relocations.cpp b/tools/pe_bliss/pe_relocations.cpp
new file mode 100644
index 0000000000..2d674e1346
--- /dev/null
+++ b/tools/pe_bliss/pe_relocations.cpp
@@ -0,0 +1,299 @@
+#include <string.h>
+#include "pe_relocations.h"
+#include "pe_properties_generic.h"
+
+namespace pe_bliss
+{
+using namespace pe_win;
+
+//RELOCATIONS
+//Default constructor
+relocation_entry::relocation_entry()
+	:rva_(0), type_(0)
+{}
+
+//Constructor from relocation item (WORD)
+relocation_entry::relocation_entry(uint16_t relocation_value)
+	:rva_(relocation_value & ((1 << 12) - 1)), type_(relocation_value >> 12)
+{}
+
+//Constructor from relative rva and relocation type
+relocation_entry::relocation_entry(uint16_t rrva, uint16_t type)
+	:rva_(rrva), type_(type)
+{}
+
+//Returns RVA of relocation
+uint16_t relocation_entry::get_rva() const
+{
+	return rva_;
+}
+
+//Returns type of relocation
+uint16_t relocation_entry::get_type() const
+{
+	return type_;
+}
+
+//Sets RVA of relocation
+void relocation_entry::set_rva(uint16_t rva)
+{
+	rva_ = rva;
+}
+
+//Sets type of relocation
+void relocation_entry::set_type(uint16_t type)
+{
+	type_ = type;
+}
+
+//Returns relocation item (rrva + type)
+uint16_t relocation_entry::get_item() const
+{
+	return rva_ | (type_ << 12);
+}
+
+//Sets relocation item (rrva + type)
+void relocation_entry::set_item(uint16_t item)
+{
+	rva_ = item & ((1 << 12) - 1);
+	type_ = item >> 12;
+}
+
+//Returns relocation list
+const relocation_table::relocation_list& relocation_table::get_relocations() const
+{
+	return relocations_;
+}
+
+//Adds relocation to table
+void relocation_table::add_relocation(const relocation_entry& entry)
+{
+	relocations_.push_back(entry);
+}
+
+//Default constructor
+relocation_table::relocation_table()
+	:rva_(0)
+{}
+
+//Constructor from RVA of relocation table
+relocation_table::relocation_table(uint32_t rva)
+	:rva_(rva)
+{}
+
+//Returns RVA of block
+uint32_t relocation_table::get_rva() const
+{
+	return rva_;
+}
+
+//Sets RVA of block
+void relocation_table::set_rva(uint32_t rva)
+{
+	rva_ = rva;
+}
+
+//Returns changeable relocation list
+relocation_table::relocation_list& relocation_table::get_relocations()
+{
+	return relocations_;
+}
+
+//Get relocation list of pe file, supports one-word sized relocations only
+//If list_absolute_entries = true, IMAGE_REL_BASED_ABSOLUTE will be listed
+const relocation_table_list get_relocations(const pe_base& pe, bool list_absolute_entries)
+{
+	relocation_table_list ret;
+
+	//If image does not have relocations
+	if(!pe.has_reloc())
+		return ret;
+
+	//Check the length in bytes of the section containing relocation directory
+	if(pe.section_data_length_from_rva(pe.get_directory_rva(image_directory_entry_basereloc),
+		pe.get_directory_rva(image_directory_entry_basereloc), section_data_virtual, true)
+		< sizeof(image_base_relocation))
+		throw pe_exception("Incorrect relocation directory", pe_exception::incorrect_relocation_directory);
+
+	unsigned long current_pos = pe.get_directory_rva(image_directory_entry_basereloc);
+	//First IMAGE_BASE_RELOCATION table
+	image_base_relocation reloc_table = pe.section_data_from_rva<image_base_relocation>(current_pos, section_data_virtual, true);
+
+	if(reloc_table.SizeOfBlock % 2)
+		throw pe_exception("Incorrect relocation directory", pe_exception::incorrect_relocation_directory);
+
+	unsigned long reloc_size = pe.get_directory_size(image_directory_entry_basereloc);
+	unsigned long read_size = 0;
+
+	//reloc_table.VirtualAddress is not checked (not so important)
+	while(reloc_table.SizeOfBlock && read_size < reloc_size)
+	{
+		//Create relocation table
+		relocation_table table;
+		//Save RVA
+		table.set_rva(reloc_table.VirtualAddress);
+
+		if(!pe_utils::is_sum_safe(current_pos, reloc_table.SizeOfBlock))
+			throw pe_exception("Incorrect relocation directory", pe_exception::incorrect_relocation_directory);
+
+		//List all relocations
+		for(unsigned long i = sizeof(image_base_relocation); i < reloc_table.SizeOfBlock; i += sizeof(uint16_t))
+		{
+			relocation_entry entry(pe.section_data_from_rva<uint16_t>(current_pos + i, section_data_virtual, true));
+			if(list_absolute_entries || entry.get_type() != image_rel_based_absolute)
+				table.add_relocation(entry);
+		}
+
+		//Save table
+		ret.push_back(table);
+		
+		//Go to next relocation block
+		if(!pe_utils::is_sum_safe(current_pos, reloc_table.SizeOfBlock))
+			throw pe_exception("Incorrect relocation directory", pe_exception::incorrect_relocation_directory);
+
+		current_pos += reloc_table.SizeOfBlock;
+		read_size += reloc_table.SizeOfBlock;
+		reloc_table = pe.section_data_from_rva<image_base_relocation>(current_pos, section_data_virtual, true);
+	}
+
+	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(needed_size + current_reloc_data_pos); //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;
+}
+
+//Recalculates image base with the help of relocation tables
+void rebase_image(pe_base& pe, const relocation_table_list& tables, uint64_t new_base)
+{
+	pe.get_pe_type() == pe_type_32
+		? rebase_image_base<pe_types_class_32>(pe, tables, new_base)
+		: rebase_image_base<pe_types_class_64>(pe, tables, new_base);
+}
+
+//RELOCATIONS
+//Recalculates image base with the help of relocation tables
+//Recalculates VAs of DWORDS/QWORDS in image according to relocations
+//Notice: if you move some critical structures like TLS, image relocations will not fix new
+//positions of TLS VAs. Instead, some bytes that now doesn't belong to TLS will be fixed.
+//It is recommended to rebase image in the very beginning and move all structures afterwards.
+template<typename PEClassType>
+void rebase_image_base(pe_base& pe, const relocation_table_list& tables, uint64_t new_base)
+{
+	//Get current image base value
+	typename PEClassType::BaseSize image_base;
+	pe.get_image_base(image_base);
+
+	//ImageBase difference
+	typename PEClassType::BaseSize base_rel = static_cast<typename PEClassType::BaseSize>(static_cast<int64_t>(new_base) - image_base);
+
+	//We need to fix addresses from relocation tables
+	//Enumerate relocation tables
+	for(relocation_table_list::const_iterator it = tables.begin(); it != tables.end(); ++it)
+	{
+		const relocation_table::relocation_list& relocs = (*it).get_relocations();
+
+		uint32_t base_rva = (*it).get_rva();
+
+		//Enumerate relocations
+		for(relocation_table::relocation_list::const_iterator rel = relocs.begin(); rel != relocs.end(); ++rel)
+		{
+			//Skip ABSOLUTE entries
+			if((*rel).get_type() == pe_win::image_rel_based_absolute)
+				continue;
+			
+			//Recalculate value by RVA and rewrite it
+			uint32_t current_rva = base_rva + (*rel).get_rva();
+			typename PEClassType::BaseSize value = pe.section_data_from_rva<typename PEClassType::BaseSize>(current_rva, section_data_raw, true);
+			value += base_rel;
+			memcpy(pe.section_data_from_rva(current_rva, true), &value, sizeof(value));
+		}
+	}
+
+	//Finally, save new image base
+	pe.set_image_base_64(new_base);
+}
+}