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