| // Copyright 2012 Google Inc. All Rights Reserved. |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| |
| #include "syzygy/pe/image_layout.h" |
| |
| #include <limits> |
| |
| #include "base/logging.h" |
| #include "base/files/file_util.h" |
| #include "syzygy/block_graph/typed_block.h" |
| #include "syzygy/common/align.h" |
| #include "syzygy/pe/pe_file.h" |
| #include "syzygy/pe/pe_utils.h" |
| |
| namespace pe { |
| |
| using block_graph::BlockGraph; |
| using block_graph::BlockVector; |
| using block_graph::ConstTypedBlock; |
| using core::RelativeAddress; |
| |
| typedef std::vector<const BlockGraph::Section*> Sections; |
| typedef std::map<BlockGraph::SectionId, BlockVector> SectionBlocks; |
| |
| namespace { |
| |
| // Generates a sorted list of sections. Uses the default block graph ordering, |
| // but ensures that rsrc and reloc are second-last and last, respectively. |
| // Also ensures that those two sections are unique. |
| bool GetOrderedSections(BlockGraph* block_graph, |
| Sections* sections) { |
| DCHECK(block_graph != NULL); |
| DCHECK(sections != NULL); |
| |
| sections->clear(); |
| sections->reserve(block_graph->sections().size()); |
| |
| const BlockGraph::Section* rsrc = NULL; |
| const BlockGraph::Section* reloc = NULL; |
| BlockGraph::SectionMap::const_iterator section_it = |
| block_graph->sections().begin(); |
| for (; section_it != block_graph->sections().end(); ++section_it) { |
| const BlockGraph::Section* section = §ion_it->second; |
| if (section->name() == kResourceSectionName) { |
| if (rsrc != NULL) { |
| LOG(ERROR) << "Found more than one " << kResourceSectionName |
| << " section."; |
| return false; |
| } |
| rsrc = section; |
| continue; |
| } |
| if (section->name() == kRelocSectionName) { |
| if (reloc != NULL) { |
| LOG(ERROR) << "Found more than one reloc section."; |
| return false; |
| } |
| reloc = section; |
| continue; |
| } |
| sections->push_back(section); |
| } |
| |
| sections->push_back(rsrc); |
| sections->push_back(reloc); |
| |
| DCHECK_EQ(sections->size(), block_graph->sections().size()); |
| |
| return true; |
| } |
| |
| // Lays out a block in the given address space. Takes care of alignment |
| // and incrementing the insert_at pointer. |
| bool LayoutBlock(BlockGraph::Block* block, |
| BlockGraph::AddressSpace* address_space, |
| RelativeAddress* insert_at) { |
| DCHECK(block != NULL); |
| DCHECK(address_space != NULL); |
| DCHECK(insert_at != NULL); |
| |
| RelativeAddress aligned = insert_at->AlignUp(block->alignment()); |
| if (!address_space->InsertBlock(aligned, block)) { |
| LOG(ERROR) << "Failed to insert block \"" << block->name() << "\" at " |
| << *insert_at << "."; |
| return false; |
| } |
| |
| *insert_at = aligned + block->size(); |
| return true; |
| } |
| |
| // Returns true if the block contains only zeros, and may safely be left |
| // implicitly initialized. |
| bool BlockIsZeros(const BlockGraph::Block* block) { |
| if (block->references().size() != 0) |
| return false; |
| const uint8_t* data = block->data(); |
| if (data == NULL) |
| return true; |
| for (size_t i = 0; i < block->data_size(); ++i, ++data) { |
| if (*data != 0) |
| return false; |
| } |
| return true; |
| } |
| |
| // Compares two blocks. Uses the source address of the first byte of each block. |
| // If one or both of the blocks has no such address then we sort such that empty |
| // (blocks that are all zeros and may be safely initialized) are pushed to the |
| // end of the section. Finally, we break ties using the block id. |
| bool BlockCompare(const BlockGraph::Block* block1, |
| const BlockGraph::Block* block2) { |
| const BlockGraph::Block::SourceRanges::RangePair* pair1 = |
| block1->source_ranges().FindRangePair(0, 1); |
| const BlockGraph::Block::SourceRanges::RangePair* pair2 = |
| block2->source_ranges().FindRangePair(0, 1); |
| |
| // If we have addresses, sort using them first. |
| if (pair1 != NULL && pair2 != NULL) { |
| RelativeAddress addr1 = pair1->second.start(); |
| RelativeAddress addr2 = pair2->second.start(); |
| if (addr1 < addr2) |
| return true; |
| if (addr2 < addr1) |
| return false; |
| } |
| |
| // Next, sort based on the contents. Blocks containing all zeros get pushed |
| // to the end of the section. |
| bool is_zeros1 = BlockIsZeros(block1); |
| bool is_zeros2 = BlockIsZeros(block2); |
| if (is_zeros1 != is_zeros2) |
| return is_zeros2; |
| |
| // Finally we break ties using the block ID. |
| return block1->id() < block2->id(); |
| } |
| |
| // Returns the length of data in a block that must be explicitly specified. |
| // Any data after this length may be implicitly initialized as zeroes. |
| size_t GetExplicitLength(const BlockGraph::Block* block) { |
| size_t length = 0; |
| |
| // Get the offset of the last byte of the last reference, if there are any. |
| if (block->references().size() > 0) { |
| BlockGraph::Block::ReferenceMap::const_reverse_iterator last_ref = |
| block->references().rbegin(); |
| length = last_ref->first + last_ref->second.size(); |
| } |
| |
| // If there is any explicit data beyond the last reference we need to |
| // manually check it. |
| if (block->data_size() > length) { |
| const uint8_t* data = block->data(); |
| DCHECK(data != NULL); |
| |
| // Walk the data backwards from the end looking for the first non-zero byte. |
| size_t i = block->data_size(); |
| data += i - 1; |
| while (i > length && *data == 0) { |
| --i; |
| --data; |
| } |
| length = i; |
| } |
| |
| return length; |
| } |
| |
| } // namespace |
| |
| void CopySectionHeadersToImageLayout( |
| size_t num_sections, |
| const IMAGE_SECTION_HEADER* section_headers, |
| std::vector<ImageLayout::SectionInfo>* sections) { |
| DCHECK(num_sections > 0); |
| DCHECK(section_headers != NULL); |
| DCHECK(sections != NULL); |
| |
| sections->clear(); |
| sections->reserve(num_sections); |
| for (size_t i = 0; i < num_sections; ++i) { |
| sections->push_back(pe::ImageLayout::SectionInfo()); |
| pe::ImageLayout::SectionInfo& section = sections->back(); |
| |
| section.name = PEFile::GetSectionName(section_headers[i]); |
| section.addr.set_value(section_headers[i].VirtualAddress); |
| section.size = section_headers[i].Misc.VirtualSize; |
| section.data_size = section_headers[i].SizeOfRawData; |
| section.characteristics = section_headers[i].Characteristics; |
| } |
| } |
| |
| bool CopyHeaderToImageLayout(const BlockGraph::Block* nt_headers_block, |
| ImageLayout* layout) { |
| ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers; |
| if (!nt_headers.Init(0, nt_headers_block)) { |
| LOG(ERROR) << "NT Headers too short."; |
| return false; |
| } |
| |
| ConstTypedBlock<IMAGE_SECTION_HEADER> section_headers; |
| size_t size = sizeof(IMAGE_SECTION_HEADER) * |
| nt_headers->FileHeader.NumberOfSections; |
| if (!section_headers.InitWithSize(sizeof(IMAGE_NT_HEADERS), |
| size, |
| nt_headers_block)) { |
| LOG(ERROR) << "NT Headers too short to contain section headers."; |
| return false; |
| } |
| |
| CopySectionHeadersToImageLayout(nt_headers->FileHeader.NumberOfSections, |
| section_headers.Get(), |
| &layout->sections); |
| return true; |
| } |
| |
| ImageLayout::ImageLayout(BlockGraph* block_graph) |
| : blocks(block_graph) { |
| } |
| |
| bool BuildCanonicalImageLayout(ImageLayout* image_layout) { |
| DCHECK(image_layout != NULL); |
| |
| BlockGraph* block_graph = image_layout->blocks.graph(); |
| |
| // First, create an ordering for the sections. This will be the same as |
| // the ordering in the underlying block graph, but we enforce that |
| // .rsrc and .reloc come second last and last. |
| Sections sections; |
| if (!GetOrderedSections(block_graph, §ions)) |
| return false; |
| |
| // Get a list of all of the blocks for each section, and find the header |
| // blocks. |
| SectionBlocks section_blocks; |
| BlockGraph::Block* dos_header_block = NULL; |
| BlockGraph::Block* nt_headers_block = NULL; |
| BlockGraph::BlockMap::iterator block_it = |
| block_graph->blocks_mutable().begin(); |
| for (; block_it != block_graph->blocks_mutable().end(); ++block_it) { |
| BlockGraph::Block* block = &block_it->second; |
| |
| // Block has no section? Identify it as a header block. |
| if (block->section() == BlockGraph::kInvalidSectionId) { |
| if (dos_header_block == NULL && IsValidDosHeaderBlock(block)) { |
| dos_header_block = block; |
| } else if (nt_headers_block == NULL && IsValidNtHeadersBlock(block)) { |
| nt_headers_block = block; |
| } else { |
| LOG(ERROR) << "Found invalid header block."; |
| return false; |
| } |
| } else { |
| section_blocks[block->section()].push_back(block); |
| } |
| } |
| |
| // Ensure we found both header blocks. |
| if (dos_header_block == NULL || nt_headers_block == NULL) { |
| LOG(ERROR) << "Missing one or both header blocks."; |
| return false; |
| } |
| |
| // Output the header blocks. |
| RelativeAddress insert_at(0); |
| if (!LayoutBlock(dos_header_block, &image_layout->blocks, &insert_at) || |
| !LayoutBlock(nt_headers_block, &image_layout->blocks, &insert_at)) { |
| return false; |
| } |
| |
| // Get the section alignment from the headers. |
| ConstTypedBlock<IMAGE_DOS_HEADER> dos_header; |
| if (!dos_header.Init(0, dos_header_block)) { |
| LOG(ERROR) << "Failed to cast dos_header_block to IMAGE_DOS_HEADER."; |
| return false; |
| } |
| ConstTypedBlock<IMAGE_NT_HEADERS> nt_header; |
| if (!dos_header.Dereference(dos_header->e_lfanew, &nt_header)) { |
| LOG(ERROR) << "Failed to cast nt_headers_block to IMAGE_NT_HEADERS."; |
| return false; |
| } |
| size_t section_alignment = nt_header->OptionalHeader.SectionAlignment; |
| size_t file_alignment = nt_header->OptionalHeader.FileAlignment; |
| |
| image_layout->sections.clear(); |
| image_layout->sections.reserve(sections.size()); |
| |
| // Output the sections one at a time. |
| for (size_t i = 0; i < sections.size(); ++i) { |
| BlockVector& blocks = section_blocks[sections[i]->id()]; |
| |
| // NOTE: BlockCompare uses BlockIsZeros, which is a slightly simpler |
| // version of GetExplicitLength (checks for explicit length == 0). |
| // This called for both blocks in each block comparison while sorting, |
| // and then explicitly again during layout. It may be worthwhile to |
| // precalculate and cache the explicit length values. |
| |
| // Sort the blocks for the section. This sorts based on the source address |
| // if the block has a single source, then based on content (empty blocks |
| // pushed to the end of the section), and finally by block id. |
| std::sort(blocks.begin(), blocks.end(), &BlockCompare); |
| |
| // Align up for the section. |
| insert_at = insert_at.AlignUp(section_alignment); |
| RelativeAddress section_start = insert_at; |
| RelativeAddress section_data_end = insert_at; |
| |
| // Layout the blocks. Keep track of the end of any blocks that aren't |
| // strictly full of zeroes, in order to determine the data size. |
| for (size_t j = 0; j < blocks.size(); ++j) { |
| BlockGraph::Block* block = blocks[j]; |
| if (!LayoutBlock(block, &image_layout->blocks, &insert_at)) |
| return false; |
| |
| // Get the explicit length of this block. If it is non-zero update the |
| // end of the explicit data for this section. |
| size_t explicit_length = GetExplicitLength(block); |
| if (explicit_length > 0) |
| section_data_end = insert_at - block->size() + explicit_length; |
| } |
| |
| // Add this section to the image layout. |
| ImageLayout::SectionInfo section_info; |
| section_info.name = sections[i]->name(); |
| section_info.addr = section_start; |
| section_info.size = insert_at - section_start; |
| section_info.data_size = common::AlignUp(section_data_end - section_start, |
| file_alignment); |
| section_info.characteristics = sections[i]->characteristics(); |
| image_layout->sections.push_back(section_info); |
| } |
| |
| return true; |
| } |
| |
| bool CopyImageLayoutWithoutPadding(const ImageLayout& input_image_layout, |
| ImageLayout* output_image_layout) { |
| DCHECK(output_image_layout != NULL); |
| DCHECK_EQ(input_image_layout.blocks.graph(), |
| output_image_layout->blocks.graph()); |
| DCHECK_EQ(0u, output_image_layout->blocks.size()); |
| |
| output_image_layout->sections = input_image_layout.sections; |
| BlockGraph* block_graph = output_image_layout->blocks.graph(); |
| |
| // Remove the padding blocks from the decomposition. We also need to create |
| // a new version of the image layout not containing those blocks. |
| BlockGraph::AddressSpace::RangeMapConstIter block_it = |
| input_image_layout.blocks.begin(); |
| for (; block_it != input_image_layout.blocks.end(); ++block_it) { |
| BlockGraph::Block* block = block_it->second; |
| |
| // If it's a padding block, remove it from the block-graph and leave it |
| // out of the new image layout. |
| if ((block->attributes() & BlockGraph::PADDING_BLOCK)) { |
| if (!block_graph->RemoveBlock(block)) { |
| return false; |
| } |
| } else { |
| // If it's not a padding block, copy it over to the new image layout. |
| if (!output_image_layout->blocks.InsertBlock(block_it->first.start(), |
| block)) { |
| return false; |
| } |
| } |
| } |
| |
| return true; |
| } |
| |
| } // namespace pe |