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/pe_image_layout_builder.h"
16 :
17 : #include <algorithm>
18 : #include <ctime>
19 :
20 : #include "base/strings/string_util.h"
21 : #include "syzygy/block_graph/typed_block.h"
22 : #include "syzygy/common/align.h"
23 : #include "syzygy/pe/pe_utils.h"
24 :
25 : namespace {
26 :
27 : using block_graph::BlockGraph;
28 : using block_graph::ConstTypedBlock;
29 : using block_graph::TypedBlock;
30 : using core::RelativeAddress;
31 :
32 : typedef std::vector<uint8> ByteVector;
33 :
34 : // A utility class to help with formatting the relocations section.
35 : class RelocWriter {
36 : public:
37 E : RelocWriter() : curr_page_(0), curr_header_offset_(0) {
38 E : }
39 :
40 E : void WriteReloc(RelativeAddress addr) {
41 E : DWORD page = PageFromAddr(addr);
42 :
43 : // Initialization case, open the first page.
44 E : if (buf_.size() == 0)
45 E : OpenPage(addr);
46 :
47 : // Close the current page, and open the next if we're outside it.
48 E : if (page != curr_page_) {
49 E : ClosePage();
50 E : OpenPage(addr);
51 : }
52 :
53 E : DCHECK_EQ(curr_page_, page);
54 E : WORD type_offset = (IMAGE_REL_BASED_HIGHLOW << 12) | OffsetFromAddr(addr);
55 E : Append(&type_offset, sizeof(type_offset));
56 E : }
57 :
58 E : void Close(ByteVector* relocs_out) {
59 E : DCHECK(relocs_out != NULL);
60 :
61 : // Close the page in progress.
62 E : if (buf_.size() != 0)
63 E : ClosePage();
64 :
65 E : relocs_out->swap(buf_);
66 E : }
67 :
68 : private:
69 : static const DWORD kPageMask = 0x00000FFF;
70 E : DWORD PageFromAddr(RelativeAddress addr) {
71 E : return addr.value() & ~kPageMask;
72 E : }
73 :
74 E : WORD OffsetFromAddr(RelativeAddress addr) {
75 E : return static_cast<WORD>(addr.value() & kPageMask);
76 E : }
77 :
78 E : void ClosePage() {
79 E : size_t block_len = buf_.size() - curr_header_offset_;
80 E : if (block_len % 4 != 0) {
81 E : DCHECK_EQ(0U, block_len % 2);
82 E : WORD filler = IMAGE_REL_BASED_ABSOLUTE << 12;
83 E : Append(&filler, sizeof(filler));
84 E : block_len += sizeof(filler);
85 : }
86 E : DCHECK_EQ(0U, block_len % 4);
87 :
88 : IMAGE_BASE_RELOCATION* header =
89 E : reinterpret_cast<IMAGE_BASE_RELOCATION*>(&buf_.at(curr_header_offset_));
90 :
91 E : header->SizeOfBlock = block_len;
92 E : }
93 :
94 E : void OpenPage(RelativeAddress addr) {
95 E : curr_page_ = PageFromAddr(addr);
96 E : curr_header_offset_ = buf_.size();
97 :
98 E : IMAGE_BASE_RELOCATION header = { curr_page_, sizeof(header) };
99 E : Append(&header, sizeof(header));
100 E : }
101 :
102 E : void Append(const void* data, size_t size) {
103 E : const uint8* buf = reinterpret_cast<const uint8*>(data);
104 E : buf_.insert(buf_.end(), buf, buf + size);
105 E : }
106 :
107 : // The buffer where we write the data.
108 : ByteVector buf_;
109 :
110 : // The current page our header is for.
111 : DWORD curr_page_;
112 :
113 : // The offset of the last IMAGE_BASE_RELOCATION header we wrote.
114 : size_t curr_header_offset_;
115 : };
116 :
117 : // Returns true iff ref is a valid reference in addr_space.
118 : bool IsValidReference(const BlockGraph::AddressSpace& addr_space,
119 : const BlockGraph::Reference& ref) {
120 : // Check that there is a referenced block.
121 : if (ref.referenced() == NULL)
122 : return false;
123 :
124 : // Check that the block is in the image.
125 : RelativeAddress addr;
126 : if (!addr_space.GetAddressOf(ref.referenced(), &addr))
127 : return false;
128 :
129 : return true;
130 : }
131 :
132 : // Functor to order references by the address of their referred block.
133 : class RefAddrLess {
134 : public:
135 : explicit RefAddrLess(const BlockGraph::AddressSpace* addr_space)
136 : : addr_space_(addr_space),
137 E : failed_(false) {
138 E : DCHECK(addr_space_ != NULL);
139 E : }
140 :
141 : bool operator()(const BlockGraph::Reference& lhs,
142 E : const BlockGraph::Reference& rhs) {
143 E : RelativeAddress lhs_addr;
144 E : RelativeAddress rhs_addr;
145 : if (!addr_space_->GetAddressOf(lhs.referenced(), &lhs_addr) ||
146 E : !addr_space_->GetAddressOf(rhs.referenced(), &rhs_addr)) {
147 i : failed_ = true;
148 : }
149 E : return lhs_addr < rhs_addr;
150 E : }
151 :
152 E : bool failed() const {
153 E : return failed_;
154 E : }
155 :
156 : private:
157 : const BlockGraph::AddressSpace* const addr_space_;
158 : bool failed_;
159 : };
160 :
161 : } // namespace
162 :
163 : namespace pe {
164 :
165 : PEImageLayoutBuilder::PEImageLayoutBuilder(ImageLayout* image_layout)
166 : : PECoffImageLayoutBuilder(image_layout),
167 : dos_header_block_(NULL),
168 E : nt_headers_block_(NULL) {
169 E : }
170 :
171 : bool PEImageLayoutBuilder::LayoutImageHeaders(
172 E : BlockGraph::Block* dos_header_block) {
173 E : DCHECK(dos_header_block != NULL);
174 E : DCHECK(dos_header_block_ == NULL);
175 E : DCHECK_EQ(0u, image_layout_->blocks.address_space_impl().size());
176 E : DCHECK_EQ(0u, image_layout_->sections.size());
177 :
178 E : if (!IsValidDosHeaderBlock(dos_header_block)) {
179 i : LOG(ERROR) << "Invalid DOS header.";
180 i : return false;
181 : }
182 :
183 : BlockGraph::Block* nt_headers_block =
184 E : GetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
185 E : if (nt_headers_block == NULL) {
186 i : LOG(ERROR) << "Invalid NT headers.";
187 i : return false;
188 : }
189 :
190 : // We keep these around for later.
191 E : dos_header_block_ = dos_header_block;
192 E : nt_headers_block_ = nt_headers_block;
193 :
194 : // Initialize alignments.
195 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
196 E : if (!nt_headers.Init(0, nt_headers_block)) {
197 i : LOG(ERROR) << "Unable to cast NT headers.";
198 i : return false;
199 : }
200 : PECoffImageLayoutBuilder::Init(nt_headers->OptionalHeader.SectionAlignment,
201 E : nt_headers->OptionalHeader.FileAlignment);
202 :
203 : // Layout the two blocks in the image layout.
204 E : if (!LayoutBlockImpl(dos_header_block))
205 i : return false;
206 E : if (!LayoutBlockImpl(nt_headers_block))
207 i : return false;
208 :
209 E : return true;
210 E : }
211 :
212 : bool PEImageLayoutBuilder::LayoutOrderedBlockGraph(
213 E : const OrderedBlockGraph& obg) {
214 : // The ordered block graph has to refer to the same underlying block graph,
215 : // and the headers must be laid out. However, nothing else should yet have
216 : // been laid out.
217 E : DCHECK_EQ(obg.block_graph(), image_layout_->blocks.graph());
218 E : DCHECK(nt_headers_block_ != NULL);
219 E : DCHECK_EQ(2u, image_layout_->blocks.address_space_impl().size());
220 E : DCHECK_EQ(0u, image_layout_->sections.size());
221 :
222 : OrderedBlockGraph::SectionList::const_iterator section_it =
223 E : obg.ordered_sections().begin();
224 : OrderedBlockGraph::SectionList::const_iterator section_end =
225 E : obg.ordered_sections().end();
226 :
227 : // Iterate through the sections.
228 E : for (; section_it != section_end; ++section_it) {
229 E : BlockGraph::Section* section = (*section_it)->section();
230 E : DCHECK(section != NULL);
231 :
232 : // Stop iterating when we see the relocs.
233 E : if (section->name() == kRelocSectionName) {
234 E : ++section_it;
235 E : break;
236 : }
237 :
238 E : if (!OpenSection(*section))
239 i : return false;
240 :
241 : // Iterate over the blocks.
242 : OrderedBlockGraph::BlockList::const_iterator block_it =
243 E : (*section_it)->ordered_blocks().begin();
244 : OrderedBlockGraph::BlockList::const_iterator block_end =
245 E : (*section_it)->ordered_blocks().end();
246 E : for (; block_it != block_end; ++block_it) {
247 E : BlockGraph::Block* block = *block_it;
248 E : if (!LayoutBlock(block))
249 i : return false;
250 E : }
251 :
252 E : if (!CloseSection())
253 i : return false;
254 E : }
255 :
256 : // There should be nothing beyond the relocs, if it was present.
257 E : if (section_it != section_end) {
258 i : LOG(ERROR) << kRelocSectionName << " not the last section.";
259 : }
260 :
261 E : return true;
262 E : }
263 :
264 E : bool PEImageLayoutBuilder::Finalize() {
265 E : if (!CreateRelocsSection())
266 i : return false;
267 :
268 E : if (!ReconcileBlockGraphAndImageLayout())
269 i : return false;
270 :
271 E : if (!SortSafeSehTable())
272 i : return false;
273 :
274 E : if (!FinalizeHeaders())
275 i : return false;
276 :
277 E : return true;
278 E : }
279 :
280 E : bool PEImageLayoutBuilder::SortSafeSehTable() {
281 E : DCHECK(nt_headers_block_ != NULL);
282 :
283 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
284 E : if (!nt_headers.Init(0, nt_headers_block_)) {
285 i : LOG(ERROR) << "Unable to cast NT headers.";
286 i : return false;
287 : }
288 :
289 : // If there is no load config directory then we can exit early.
290 : IMAGE_DATA_DIRECTORY* load_config =
291 : nt_headers->OptionalHeader.DataDirectory +
292 E : IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG;
293 : if (load_config->VirtualAddress == 0 && load_config->Size == 0 &&
294 E : !nt_headers.HasReference(load_config->VirtualAddress)) {
295 i : return true;
296 : }
297 :
298 E : TypedBlock<IMAGE_LOAD_CONFIG_DIRECTORY> load_config_directory;
299 : if (!nt_headers.Dereference(load_config->VirtualAddress,
300 E : &load_config_directory)) {
301 i : LOG(ERROR) << "Failed to dereference Load Config Directory.";
302 i : return false;
303 : }
304 :
305 E : TypedBlock<DWORD> safe_seh_table;
306 : if (!load_config_directory.Dereference(
307 E : load_config_directory->SEHandlerTable, &safe_seh_table)) {
308 : // There's no SEHandlerTable.
309 i : return true;
310 : }
311 :
312 : // Grab the references to the safe SEH code blocks.
313 : typedef BlockGraph::Block::ReferenceMap ReferenceMap;
314 E : const ReferenceMap& orig_references = safe_seh_table.block()->references();
315 :
316 : // We should have as many references as there are handlers and we expect the
317 : // safe seh block to be zero offset and exactly the right size.
318 E : size_t num_references = orig_references.size();
319 : if (num_references != load_config_directory->SEHandlerCount ||
320 : safe_seh_table.offset() != 0 ||
321 E : safe_seh_table.block()->size() != num_references * sizeof(DWORD)) {
322 i : LOG(ERROR) << "Safe SEH Table block does not conform to expectations.";
323 i : return false;
324 : }
325 :
326 : // Create a secondary vector large enough to hold the sorted references.
327 : typedef std::vector<BlockGraph::Reference> ReferenceVector;
328 E : ReferenceVector sorted_references;
329 E : sorted_references.reserve(orig_references.size());
330 :
331 : // Copy the references into a secondary vector.
332 E : for (ReferenceMap::const_iterator iter = orig_references.begin();
333 E : iter != orig_references.end();
334 E : ++iter) {
335 E : sorted_references.push_back(iter->second);
336 E : }
337 :
338 : // Sort the secondary vector in the order their referred blocks appear
339 : // in the image layout.
340 E : RefAddrLess comparator(&image_layout_->blocks);
341 E : std::sort(sorted_references.begin(), sorted_references.end(), comparator);
342 E : if (comparator.failed()) {
343 i : LOG(ERROR) << "One or more exception handler blocks is invalid.";
344 i : return false;
345 : }
346 :
347 : // Reset the references in the Safe SEH Table in sorted order.
348 E : size_t offset = 0;
349 E : for (ReferenceVector::iterator iter = sorted_references.begin();
350 E : iter != sorted_references.end();
351 E : offset += sizeof(DWORD), ++iter) {
352 E : DCHECK(iter->size() == sizeof(DWORD));
353 E : DCHECK(iter->referenced()->type() == BlockGraph::CODE_BLOCK);
354 E : safe_seh_table.block()->SetReference(offset, *iter);
355 E : }
356 :
357 E : return true;
358 E : }
359 :
360 E : bool PEImageLayoutBuilder::CreateRelocsSection() {
361 E : RelocWriter writer;
362 :
363 E : DCHECK(nt_headers_block_ != NULL);
364 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
365 E : if (!nt_headers.Init(0, nt_headers_block_)) {
366 i : LOG(ERROR) << "Unable to cast NT headers.";
367 i : return false;
368 : }
369 :
370 : // Get the existing relocs block so we can reuse it.
371 E : TypedBlock<unsigned char> reloc_data;
372 : if (!nt_headers.Dereference(
373 : nt_headers->OptionalHeader.DataDirectory[
374 E : IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress, &reloc_data)) {
375 i : LOG(ERROR) << "Unable to dereference relocs block.";
376 i : return false;
377 : }
378 E : BlockGraph::Block* relocs_block = reloc_data.block();
379 E : CHECK_EQ(0, reloc_data.offset());
380 :
381 : // Iterate over all blocks in the address space, in the order of increasing
382 : // addresses.
383 : BlockGraph::AddressSpace::RangeMap::const_iterator it(
384 E : image_layout_->blocks.address_space_impl().ranges().begin());
385 : BlockGraph::AddressSpace::RangeMap::const_iterator end(
386 E : image_layout_->blocks.address_space_impl().ranges().end());
387 :
388 E : for (; it != end; ++it) {
389 E : const BlockGraph::Block* block = it->second;
390 E : RelativeAddress block_addr;
391 E : CHECK(image_layout_->blocks.GetAddressOf(block, &block_addr));
392 :
393 : // Iterate over all outgoing references in this block in
394 : // order of increasing offset.
395 : BlockGraph::Block::ReferenceMap::const_iterator ref_it(
396 E : block->references().begin());
397 : BlockGraph::Block::ReferenceMap::const_iterator ref_end(
398 E : block->references().end());
399 E : for (; ref_it != ref_end; ++ref_it) {
400 : // Add each absolute reference to the relocs.
401 E : if (ref_it->second.type() == BlockGraph::ABSOLUTE_REF) {
402 E : writer.WriteReloc(block_addr + ref_it->first);
403 : }
404 E : }
405 E : }
406 :
407 : // Get the relocations data from the writer.
408 E : ByteVector relocs;
409 E : writer.Close(&relocs);
410 :
411 : // Update the block and the data directory.
412 E : relocs_block->source_ranges().clear();
413 E : relocs_block->SetData(NULL, 0);
414 E : relocs_block->set_size(relocs.size());
415 E : if (!relocs_block->CopyData(relocs.size(), &relocs.at(0))) {
416 i : LOG(ERROR) << "Unable to copy relocs data.";
417 i : return false;
418 : }
419 : nt_headers->OptionalHeader.DataDirectory[
420 E : IMAGE_DIRECTORY_ENTRY_BASERELOC].Size = relocs.size();
421 :
422 : // Layout the relocs.
423 E : if (!OpenSection(kRelocSectionName, kRelocCharacteristics))
424 i : return false;
425 E : if (!LayoutBlock(relocs_block))
426 i : return false;
427 E : if (!CloseSection())
428 i : return false;
429 :
430 E : return true;
431 E : }
432 :
433 E : bool PEImageLayoutBuilder::ReconcileBlockGraphAndImageLayout() {
434 : // Get the reloc section ID from the block-graph.
435 : BlockGraph::Section* reloc_section =
436 E : image_layout_->blocks.graph()->FindSection(kRelocSectionName);
437 E : if (reloc_section == NULL) {
438 i : LOG(ERROR) << "Unable to find the reloc section in the block-graph.";
439 i : return false;
440 : }
441 E : BlockGraph::SectionId reloc_section_id = reloc_section->id();
442 :
443 : // Iterate over the blocks of the block-graph to see if some of them are not
444 : // in the image layout. If we find one we check if it belongs to the reloc
445 : // section, in this case we put it in a list of blocks that we should remove
446 : // from the graph, otherwise we return an error.
447 : BlockGraph::BlockMap::iterator it_block_graph =
448 E : image_layout_->blocks.graph()->blocks_mutable().begin();
449 E : std::list<BlockGraph::Block*> blocks_to_remove;
450 :
451 E : for (; it_block_graph != image_layout_->blocks.graph()->blocks().end();
452 E : ++it_block_graph) {
453 : // Determine if the current block exist in the image layout.
454 E : if (!image_layout_->blocks.ContainsBlock(&it_block_graph->second)) {
455 : // If it doesn't we check to see if this block belongs to the reloc
456 : // section.
457 i : if (it_block_graph->second.section() != reloc_section_id) {
458 i : LOG(ERROR) << "There is a block in the block-graph that is not in the "
459 : << "image layout (id=" << it_block_graph->second.id()
460 : << ", name=\"" << it_block_graph->second.name() << "\", "
461 : << "original address=" << it_block_graph->second.addr()
462 : << ").";
463 i : return false;
464 i : } else {
465 : // The block is added to the list of blocks to remove from the graph.
466 i : blocks_to_remove.push_back(&it_block_graph->second);
467 : }
468 : }
469 E : }
470 :
471 : // The useless blocks are removed from the block-graph.
472 : std::list<BlockGraph::Block*>::iterator iter_blocks =
473 E : blocks_to_remove.begin();
474 E : for (; iter_blocks != blocks_to_remove.end(); ++iter_blocks) {
475 i : if (!image_layout_->blocks.graph()->RemoveBlock(*iter_blocks)) {
476 i : LOG(ERROR) << "Unable to remove block with ID " << (*iter_blocks)->id()
477 : << " from the block-graph.";
478 : }
479 i : }
480 :
481 : DCHECK_EQ(image_layout_->blocks.size(),
482 E : image_layout_->blocks.graph()->blocks().size());
483 :
484 E : return true;
485 E : }
486 :
487 E : bool PEImageLayoutBuilder::FinalizeHeaders() {
488 : // The DOS and NT headers must be set at this point.
489 E : DCHECK(dos_header_block_ != NULL);
490 E : DCHECK(nt_headers_block_ != NULL);
491 :
492 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
493 E : if (!nt_headers.Init(0, nt_headers_block_)) {
494 i : LOG(ERROR) << "Unable to cast NT headers.";
495 i : return false;
496 : }
497 :
498 E : TypedBlock<IMAGE_SECTION_HEADER> section_headers;
499 E : if (!section_headers.Init(sizeof(IMAGE_NT_HEADERS), nt_headers_block_)) {
500 i : LOG(ERROR) << "Unable to cast section headers.";
501 i : return false;
502 : }
503 :
504 : // Ensure the section headers have the expected size. If they don't we bail,
505 : // as this should have been done prior to layout (PrepareHeadersTransform).
506 E : if (section_headers.ElementCount() != image_layout_->sections.size()) {
507 i : LOG(ERROR) << "Section header count does not agree with layout section "
508 : << "count (" << section_headers.ElementCount() << " != "
509 : << image_layout_->sections.size() << ").";
510 i : return false;
511 : }
512 :
513 : core::FileOffsetAddress section_file_start(
514 E : nt_headers->OptionalHeader.SizeOfHeaders);
515 :
516 : // Iterate through our sections to initialize the code/data fields in the NT
517 : // headers.
518 E : nt_headers->OptionalHeader.SizeOfCode = 0;
519 E : nt_headers->OptionalHeader.SizeOfInitializedData = 0;
520 E : nt_headers->OptionalHeader.SizeOfUninitializedData = 0;
521 E : nt_headers->OptionalHeader.BaseOfCode = 0;
522 E : nt_headers->OptionalHeader.BaseOfData = 0;
523 E : for (size_t i = 0; i < image_layout_->sections.size(); ++i) {
524 E : const ImageLayout::SectionInfo& section = image_layout_->sections[i];
525 E : IMAGE_SECTION_HEADER& hdr = section_headers[i];
526 :
527 E : if (section.characteristics& IMAGE_SCN_CNT_CODE) {
528 E : nt_headers->OptionalHeader.SizeOfCode += section.data_size;
529 E : if (nt_headers->OptionalHeader.BaseOfCode == 0) {
530 E : nt_headers->OptionalHeader.BaseOfCode = section.addr.value();
531 : }
532 : }
533 E : if (section.characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA) {
534 E : nt_headers->OptionalHeader.SizeOfInitializedData += section.data_size;
535 :
536 E : if (nt_headers->OptionalHeader.BaseOfData == 0)
537 E : nt_headers->OptionalHeader.BaseOfData = section.addr.value();
538 : }
539 E : if (section.characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA) {
540 : nt_headers->OptionalHeader.SizeOfUninitializedData +=
541 i : section.data_size;
542 i : if (nt_headers->OptionalHeader.BaseOfData == 0)
543 i : nt_headers->OptionalHeader.BaseOfData = section.addr.value();
544 : }
545 :
546 : // Zero the header to get rid of any old crud in it.
547 E : memset(&hdr, 0, sizeof(hdr));
548 :
549 : strncpy(reinterpret_cast<char*>(hdr.Name),
550 : section.name.c_str(),
551 E : arraysize(hdr.Name));
552 E : hdr.Misc.VirtualSize = section.size;
553 E : hdr.VirtualAddress = section.addr.value();
554 E : hdr.SizeOfRawData = section.data_size;
555 E : hdr.PointerToRawData = section_file_start.value();
556 E : hdr.Characteristics = section.characteristics;
557 :
558 E : section_file_start += section.data_size;
559 E : }
560 :
561 : nt_headers->OptionalHeader.SizeOfImage =
562 E : cursor_.AlignUp(nt_headers->OptionalHeader.SectionAlignment).value();
563 :
564 E : return true;
565 E : }
566 :
567 : } // namespace pe
|