1 : // Copyright 2012 Google Inc. All Rights Reserved.
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // http://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 :
15 : #include "syzygy/pe/image_layout.h"
16 :
17 : #include <limits>
18 :
19 : #include "base/logging.h"
20 : #include "base/files/file_util.h"
21 : #include "syzygy/block_graph/typed_block.h"
22 : #include "syzygy/common/align.h"
23 : #include "syzygy/pe/pe_file.h"
24 : #include "syzygy/pe/pe_utils.h"
25 :
26 : namespace pe {
27 :
28 : using block_graph::BlockGraph;
29 : using block_graph::BlockVector;
30 : using block_graph::ConstTypedBlock;
31 : using core::RelativeAddress;
32 :
33 : typedef std::vector<const BlockGraph::Section*> Sections;
34 : typedef std::map<BlockGraph::SectionId, BlockVector> SectionBlocks;
35 :
36 : namespace {
37 :
38 : // Generates a sorted list of sections. Uses the default block graph ordering,
39 : // but ensures that rsrc and reloc are second-last and last, respectively.
40 : // Also ensures that those two sections are unique.
41 : bool GetOrderedSections(BlockGraph* block_graph,
42 E : Sections* sections) {
43 E : DCHECK(block_graph != NULL);
44 E : DCHECK(sections != NULL);
45 :
46 E : sections->clear();
47 E : sections->reserve(block_graph->sections().size());
48 :
49 E : const BlockGraph::Section* rsrc = NULL;
50 E : const BlockGraph::Section* reloc = NULL;
51 : BlockGraph::SectionMap::const_iterator section_it =
52 E : block_graph->sections().begin();
53 E : for (; section_it != block_graph->sections().end(); ++section_it) {
54 E : const BlockGraph::Section* section = §ion_it->second;
55 E : if (section->name() == kResourceSectionName) {
56 E : if (rsrc != NULL) {
57 i : LOG(ERROR) << "Found more than one " << kResourceSectionName
58 : << " section.";
59 i : return false;
60 : }
61 E : rsrc = section;
62 E : continue;
63 : }
64 E : if (section->name() == kRelocSectionName) {
65 E : if (reloc != NULL) {
66 i : LOG(ERROR) << "Found more than one reloc section.";
67 i : return false;
68 : }
69 E : reloc = section;
70 E : continue;
71 : }
72 E : sections->push_back(section);
73 E : }
74 :
75 E : sections->push_back(rsrc);
76 E : sections->push_back(reloc);
77 :
78 E : DCHECK_EQ(sections->size(), block_graph->sections().size());
79 :
80 E : return true;
81 E : }
82 :
83 : // Lays out a block in the given address space. Takes care of alignment
84 : // and incrementing the insert_at pointer.
85 : bool LayoutBlock(BlockGraph::Block* block,
86 : BlockGraph::AddressSpace* address_space,
87 E : RelativeAddress* insert_at) {
88 E : DCHECK(block != NULL);
89 E : DCHECK(address_space != NULL);
90 E : DCHECK(insert_at != NULL);
91 :
92 E : RelativeAddress aligned = insert_at->AlignUp(block->alignment());
93 E : if (!address_space->InsertBlock(aligned, block)) {
94 i : LOG(ERROR) << "Failed to insert block \"" << block->name() << "\" at "
95 : << *insert_at << ".";
96 i : return false;
97 : }
98 :
99 E : *insert_at = aligned + block->size();
100 E : return true;
101 E : }
102 :
103 : // Returns true if the block contains only zeros, and may safely be left
104 : // implicitly initialized.
105 i : bool BlockIsZeros(const BlockGraph::Block* block) {
106 i : if (block->references().size() != 0)
107 i : return false;
108 i : const uint8* data = block->data();
109 i : if (data == NULL)
110 i : return true;
111 i : for (size_t i = 0; i < block->data_size(); ++i, ++data) {
112 i : if (*data != 0)
113 i : return false;
114 i : }
115 i : return true;
116 i : }
117 :
118 : // Compares two blocks. Uses the source address of the first byte of each block.
119 : // If one or both of the blocks has no such address then we sort such that empty
120 : // (blocks that are all zeros and may be safely initialized) are pushed to the
121 : // end of the section. Finally, we break ties using the block id.
122 : bool BlockCompare(const BlockGraph::Block* block1,
123 E : const BlockGraph::Block* block2) {
124 : const BlockGraph::Block::SourceRanges::RangePair* pair1 =
125 E : block1->source_ranges().FindRangePair(0, 1);
126 : const BlockGraph::Block::SourceRanges::RangePair* pair2 =
127 E : block2->source_ranges().FindRangePair(0, 1);
128 :
129 : // If we have addresses, sort using them first.
130 E : if (pair1 != NULL && pair2 != NULL) {
131 E : RelativeAddress addr1 = pair1->second.start();
132 E : RelativeAddress addr2 = pair2->second.start();
133 E : if (addr1 < addr2)
134 E : return true;
135 E : if (addr2 < addr1)
136 E : return false;
137 : }
138 :
139 : // Next, sort based on the contents. Blocks containing all zeros get pushed
140 : // to the end of the section.
141 i : bool is_zeros1 = BlockIsZeros(block1);
142 i : bool is_zeros2 = BlockIsZeros(block2);
143 i : if (is_zeros1 != is_zeros2)
144 i : return is_zeros2;
145 :
146 : // Finally we break ties using the block ID.
147 i : return block1->id() < block2->id();
148 E : }
149 :
150 : // Returns the length of data in a block that must be explicitly specified.
151 : // Any data after this length may be implicitly initialized as zeroes.
152 E : size_t GetExplicitLength(const BlockGraph::Block* block) {
153 E : size_t length = 0;
154 :
155 : // Get the offset of the last byte of the last reference, if there are any.
156 E : if (block->references().size() > 0) {
157 : BlockGraph::Block::ReferenceMap::const_reverse_iterator last_ref =
158 E : block->references().rbegin();
159 E : length = last_ref->first + last_ref->second.size();
160 : }
161 :
162 : // If there is any explicit data beyond the last reference we need to
163 : // manually check it.
164 E : if (block->data_size() > length) {
165 E : const uint8* data = block->data();
166 E : DCHECK(data != NULL);
167 :
168 : // Walk the data backwards from the end looking for the first non-zero byte.
169 E : size_t i = block->data_size();
170 E : data += i - 1;
171 E : while (i > length && *data == 0) {
172 E : --i;
173 E : --data;
174 E : }
175 E : length = i;
176 : }
177 :
178 E : return length;
179 E : }
180 :
181 : } // namespace
182 :
183 : void CopySectionHeadersToImageLayout(
184 : size_t num_sections,
185 : const IMAGE_SECTION_HEADER* section_headers,
186 E : std::vector<ImageLayout::SectionInfo>* sections) {
187 E : DCHECK(num_sections > 0);
188 E : DCHECK(section_headers != NULL);
189 E : DCHECK(sections != NULL);
190 :
191 E : sections->clear();
192 E : sections->reserve(num_sections);
193 E : for (size_t i = 0; i < num_sections; ++i) {
194 E : sections->push_back(pe::ImageLayout::SectionInfo());
195 E : pe::ImageLayout::SectionInfo& section = sections->back();
196 :
197 E : section.name = PEFile::GetSectionName(section_headers[i]);
198 E : section.addr.set_value(section_headers[i].VirtualAddress);
199 E : section.size = section_headers[i].Misc.VirtualSize;
200 E : section.data_size = section_headers[i].SizeOfRawData;
201 E : section.characteristics = section_headers[i].Characteristics;
202 E : }
203 E : }
204 :
205 : bool CopyHeaderToImageLayout(const BlockGraph::Block* nt_headers_block,
206 E : ImageLayout* layout) {
207 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
208 E : if (!nt_headers.Init(0, nt_headers_block)) {
209 i : LOG(ERROR) << "NT Headers too short.";
210 i : return false;
211 : }
212 :
213 E : ConstTypedBlock<IMAGE_SECTION_HEADER> section_headers;
214 : size_t size = sizeof(IMAGE_SECTION_HEADER) *
215 E : nt_headers->FileHeader.NumberOfSections;
216 : if (!section_headers.InitWithSize(sizeof(IMAGE_NT_HEADERS),
217 : size,
218 E : nt_headers_block)) {
219 i : LOG(ERROR) << "NT Headers too short to contain section headers.";
220 i : return false;
221 : }
222 :
223 : CopySectionHeadersToImageLayout(nt_headers->FileHeader.NumberOfSections,
224 : section_headers.Get(),
225 E : &layout->sections);
226 E : return true;
227 E : }
228 :
229 : ImageLayout::ImageLayout(BlockGraph* block_graph)
230 E : : blocks(block_graph) {
231 E : }
232 :
233 E : bool BuildCanonicalImageLayout(ImageLayout* image_layout) {
234 E : DCHECK(image_layout != NULL);
235 :
236 E : BlockGraph* block_graph = image_layout->blocks.graph();
237 :
238 : // First, create an ordering for the sections. This will be the same as
239 : // the ordering in the underlying block graph, but we enforce that
240 : // .rsrc and .reloc come second last and last.
241 E : Sections sections;
242 E : if (!GetOrderedSections(block_graph, §ions))
243 i : return false;
244 :
245 : // Get a list of all of the blocks for each section, and find the header
246 : // blocks.
247 E : SectionBlocks section_blocks;
248 E : BlockGraph::Block* dos_header_block = NULL;
249 E : BlockGraph::Block* nt_headers_block = NULL;
250 : BlockGraph::BlockMap::iterator block_it =
251 E : block_graph->blocks_mutable().begin();
252 E : for (; block_it != block_graph->blocks_mutable().end(); ++block_it) {
253 E : BlockGraph::Block* block = &block_it->second;
254 :
255 : // Block has no section? Identify it as a header block.
256 E : if (block->section() == BlockGraph::kInvalidSectionId) {
257 E : if (dos_header_block == NULL && IsValidDosHeaderBlock(block)) {
258 E : dos_header_block = block;
259 E : } else if (nt_headers_block == NULL && IsValidNtHeadersBlock(block)) {
260 E : nt_headers_block = block;
261 E : } else {
262 i : LOG(ERROR) << "Found invalid header block.";
263 i : return false;
264 : }
265 E : } else {
266 E : section_blocks[block->section()].push_back(block);
267 : }
268 E : }
269 :
270 : // Ensure we found both header blocks.
271 E : if (dos_header_block == NULL || nt_headers_block == NULL) {
272 i : LOG(ERROR) << "Missing one or both header blocks.";
273 i : return false;
274 : }
275 :
276 : // Output the header blocks.
277 E : RelativeAddress insert_at(0);
278 : if (!LayoutBlock(dos_header_block, &image_layout->blocks, &insert_at) ||
279 E : !LayoutBlock(nt_headers_block, &image_layout->blocks, &insert_at)) {
280 i : return false;
281 : }
282 :
283 : // Get the section alignment from the headers.
284 E : ConstTypedBlock<IMAGE_DOS_HEADER> dos_header;
285 E : if (!dos_header.Init(0, dos_header_block)) {
286 i : LOG(ERROR) << "Failed to cast dos_header_block to IMAGE_DOS_HEADER.";
287 i : return false;
288 : }
289 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_header;
290 E : if (!dos_header.Dereference(dos_header->e_lfanew, &nt_header)) {
291 i : LOG(ERROR) << "Failed to cast nt_headers_block to IMAGE_NT_HEADERS.";
292 i : return false;
293 : }
294 E : size_t section_alignment = nt_header->OptionalHeader.SectionAlignment;
295 E : size_t file_alignment = nt_header->OptionalHeader.FileAlignment;
296 :
297 E : image_layout->sections.clear();
298 E : image_layout->sections.reserve(sections.size());
299 :
300 : // Output the sections one at a time.
301 E : for (size_t i = 0; i < sections.size(); ++i) {
302 E : BlockVector& blocks = section_blocks[sections[i]->id()];
303 :
304 : // NOTE: BlockCompare uses BlockIsZeros, which is a slightly simpler
305 : // version of GetExplicitLength (checks for explicit length == 0).
306 : // This called for both blocks in each block comparison while sorting,
307 : // and then explicitly again during layout. It may be worthwhile to
308 : // precalculate and cache the explicit length values.
309 :
310 : // Sort the blocks for the section. This sorts based on the source address
311 : // if the block has a single source, then based on content (empty blocks
312 : // pushed to the end of the section), and finally by block id.
313 E : std::sort(blocks.begin(), blocks.end(), &BlockCompare);
314 :
315 : // Align up for the section.
316 E : insert_at = insert_at.AlignUp(section_alignment);
317 E : RelativeAddress section_start = insert_at;
318 E : RelativeAddress section_data_end = insert_at;
319 :
320 : // Layout the blocks. Keep track of the end of any blocks that aren't
321 : // strictly full of zeroes, in order to determine the data size.
322 E : for (size_t j = 0; j < blocks.size(); ++j) {
323 E : BlockGraph::Block* block = blocks[j];
324 E : if (!LayoutBlock(block, &image_layout->blocks, &insert_at))
325 i : return false;
326 :
327 : // Get the explicit length of this block. If it is non-zero update the
328 : // end of the explicit data for this section.
329 E : size_t explicit_length = GetExplicitLength(block);
330 E : if (explicit_length > 0)
331 E : section_data_end = insert_at - block->size() + explicit_length;
332 E : }
333 :
334 : // Add this section to the image layout.
335 E : ImageLayout::SectionInfo section_info;
336 E : section_info.name = sections[i]->name();
337 E : section_info.addr = section_start;
338 E : section_info.size = insert_at - section_start;
339 : section_info.data_size = common::AlignUp(section_data_end - section_start,
340 E : file_alignment);
341 E : section_info.characteristics = sections[i]->characteristics();
342 E : image_layout->sections.push_back(section_info);
343 E : }
344 :
345 E : return true;
346 E : }
347 :
348 : bool CopyImageLayoutWithoutPadding(const ImageLayout& input_image_layout,
349 E : ImageLayout* output_image_layout) {
350 E : DCHECK(output_image_layout != NULL);
351 : DCHECK_EQ(input_image_layout.blocks.graph(),
352 E : output_image_layout->blocks.graph());
353 E : DCHECK_EQ(0u, output_image_layout->blocks.size());
354 :
355 E : output_image_layout->sections = input_image_layout.sections;
356 E : BlockGraph* block_graph = output_image_layout->blocks.graph();
357 :
358 : // Remove the padding blocks from the decomposition. We also need to create
359 : // a new version of the image layout not containing those blocks.
360 : BlockGraph::AddressSpace::RangeMapConstIter block_it =
361 E : input_image_layout.blocks.begin();
362 E : for (; block_it != input_image_layout.blocks.end(); ++block_it) {
363 E : BlockGraph::Block* block = block_it->second;
364 :
365 : // If it's a padding block, remove it from the block-graph and leave it
366 : // out of the new image layout.
367 E : if ((block->attributes() & BlockGraph::PADDING_BLOCK)) {
368 E : if (!block_graph->RemoveBlock(block)) {
369 i : return false;
370 : }
371 E : } else {
372 : // If it's not a padding block, copy it over to the new image layout.
373 : if (!output_image_layout->blocks.InsertBlock(block_it->first.start(),
374 E : block)) {
375 i : return false;
376 : }
377 : }
378 E : }
379 :
380 E : return true;
381 E : }
382 :
383 : } // namespace pe
|