1 : // Copyright 2012 Google Inc.
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_builder.h"
16 :
17 : #include <delayimp.h>
18 : #include <algorithm>
19 : #include <ctime>
20 :
21 : #include "base/string_util.h"
22 : #include "syzygy/block_graph/typed_block.h"
23 : #include "syzygy/common/align.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 : ImageLayoutBuilder::ImageLayoutBuilder(ImageLayout* image_layout)
167 : : image_layout_(image_layout),
168 : padding_(0),
169 : dos_header_block_(NULL),
170 E : nt_headers_block_(NULL) {
171 E : DCHECK(image_layout != NULL);
172 E : DCHECK_EQ(0u, image_layout->blocks.address_space_impl().size());
173 E : DCHECK_EQ(0u, image_layout->sections.size());
174 E : }
175 :
176 : bool ImageLayoutBuilder::LayoutImageHeaders(
177 E : BlockGraph::Block* dos_header_block) {
178 E : DCHECK(dos_header_block != NULL);
179 E : DCHECK(dos_header_block_ == NULL);
180 E : DCHECK_EQ(0u, image_layout_->blocks.address_space_impl().size());
181 E : DCHECK_EQ(0u, image_layout_->sections.size());
182 :
183 E : if (!IsValidDosHeaderBlock(dos_header_block)) {
184 i : LOG(ERROR) << "Invalid DOS header.";
185 i : return false;
186 : }
187 :
188 : BlockGraph::Block* nt_headers_block =
189 E : GetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
190 E : if (nt_headers_block == NULL) {
191 i : LOG(ERROR) << "Invalid NT headers.";
192 i : return false;
193 : }
194 :
195 : // We keep these around for later.
196 E : dos_header_block_ = dos_header_block;
197 E : nt_headers_block_ = nt_headers_block;
198 :
199 : // Layout the two blocks in the image layout.
200 E : if (!LayoutBlockImpl(dos_header_block))
201 i : return false;
202 E : if (!LayoutBlockImpl(nt_headers_block))
203 i : return false;
204 :
205 E : return true;
206 E : }
207 :
208 E : bool ImageLayoutBuilder::OpenSection(const char* name, uint32 characteristics) {
209 E : DCHECK(name != NULL);
210 E : DCHECK(nt_headers_block_ != NULL);
211 :
212 : // If we're already in a section, close it.
213 E : if (section_start_.value() != 0)
214 i : CloseSection();
215 :
216 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
217 E : if (!nt_headers.Init(0, nt_headers_block_)) {
218 i : LOG(ERROR) << "Unable to cast NT headers.";
219 i : return false;
220 : }
221 :
222 : // Align to the start of the next section.
223 E : DCHECK_GT(cursor_.value(), 0u);
224 E : cursor_ = cursor_.AlignUp(nt_headers->OptionalHeader.SectionAlignment);
225 :
226 : // Remember the start of the section and reset the initialized data cursors.
227 E : DCHECK_EQ(0u, section_start_.value());
228 E : DCHECK_EQ(0u, section_auto_init_end_.value());
229 E : DCHECK_EQ(0u, section_init_end_.value());
230 E : section_start_ = cursor_;
231 E : section_auto_init_end_ = cursor_;
232 E : section_init_end_ = cursor_;
233 :
234 : // Create a section.
235 E : ImageLayout::SectionInfo section_info;
236 E : section_info.name = name;
237 E : section_info.addr = section_start_;
238 E : section_info.size = 0;
239 E : section_info.data_size = 0;
240 E : section_info.characteristics = characteristics;
241 E : image_layout_->sections.push_back(section_info);
242 :
243 E : return true;
244 E : }
245 :
246 E : bool ImageLayoutBuilder::OpenSection(const BlockGraph::Section* section) {
247 E : return OpenSection(section->name().c_str(), section->characteristics());
248 E : }
249 :
250 : void ImageLayoutBuilder::AdvanceCursor(size_t distance) {
251 : cursor_ += distance;
252 : }
253 :
254 : size_t ImageLayoutBuilder::AlignCursor(size_t alignment) {
255 : RelativeAddress old_cursor(cursor_);
256 : cursor_ = cursor_.AlignUp(alignment);
257 : return cursor_ - old_cursor;
258 : }
259 :
260 E : bool ImageLayoutBuilder::LayoutBlock(BlockGraph::Block* block) {
261 E : DCHECK(block != NULL);
262 E : return LayoutBlock(block->alignment(), block);
263 E : }
264 :
265 : bool ImageLayoutBuilder::LayoutBlock(size_t alignment,
266 E : BlockGraph::Block* block) {
267 E : DCHECK(block != NULL);
268 E : DCHECK_NE(0u, section_start_.value());
269 :
270 : // If this is not the first block of the section and we have padding, then
271 : // output the padding.
272 E : if (padding_ > 0 && cursor_ > section_start_)
273 E : cursor_ += padding_;
274 :
275 E : cursor_ = cursor_.AlignUp(alignment);
276 :
277 : // If we have explicit data, advance the explicit data cursor.
278 E : if (block->data_size() > 0)
279 E : section_auto_init_end_ = cursor_ + block->data_size();
280 :
281 : // This advances the cursor for us.
282 E : if (!LayoutBlockImpl(block))
283 i : return false;
284 :
285 E : return true;
286 E : }
287 :
288 : void ImageLayoutBuilder::CloseExplicitSectionData() {
289 : DCHECK_NE(0u, section_start_.value());
290 : section_init_end_ = cursor_;
291 : }
292 :
293 E : bool ImageLayoutBuilder::CloseSection() {
294 E : DCHECK_NE(0u, section_start_.value());
295 E : DCHECK_LT(0u, image_layout_->sections.size());
296 :
297 E : DCHECK(nt_headers_block_ != NULL);
298 E : ConstTypedBlock<IMAGE_NT_HEADERS> nt_headers;
299 E : if (!nt_headers.Init(0, nt_headers_block_)) {
300 i : LOG(ERROR) << "Unable to cast NT headers.";
301 i : return false;
302 : }
303 :
304 E : size_t section_size = cursor_ - section_start_;
305 :
306 : // If provided use the explicit initialized data size, otherwise use the
307 : // automatic one.
308 E : size_t init_size = 0;
309 E : if (section_init_end_ > cursor_) {
310 i : if (section_auto_init_end_ > section_init_end_) {
311 i : LOG(ERROR) << "Blocks with initialized data lay beyond explicitly "
312 : "specified end of initialized data.";
313 i : return false;
314 : }
315 i : init_size = section_init_end_ - section_start_;
316 i : } else {
317 E : init_size = section_auto_init_end_ - section_start_;
318 : }
319 :
320 : // A section must have *some* presence in the file.
321 E : if (init_size == 0)
322 i : init_size = 1;
323 :
324 : init_size = common::AlignUp(init_size,
325 E : nt_headers->OptionalHeader.FileAlignment);
326 :
327 E : ImageLayout::SectionInfo& section_info = image_layout_->sections.back();
328 E : section_info.size = section_size;
329 E : section_info.data_size = init_size;
330 :
331 E : section_start_.set_value(0);
332 E : section_auto_init_end_.set_value(0);
333 E : section_init_end_.set_value(0);
334 :
335 E : return true;
336 E : }
337 :
338 : bool ImageLayoutBuilder::LayoutOrderedBlockGraph(
339 E : const OrderedBlockGraph& obg) {
340 : // The ordered block graph has to refer to the same underlying block graph,
341 : // and the headers must be laid out. However, nothing else should yet have
342 : // been laid out.
343 E : DCHECK_EQ(obg.block_graph(), image_layout_->blocks.graph());
344 E : DCHECK(nt_headers_block_ != NULL);
345 E : DCHECK_EQ(2u, image_layout_->blocks.address_space_impl().size());
346 E : DCHECK_EQ(0u, image_layout_->sections.size());
347 :
348 : OrderedBlockGraph::SectionList::const_iterator section_it =
349 E : obg.ordered_sections().begin();
350 : OrderedBlockGraph::SectionList::const_iterator section_end =
351 E : obg.ordered_sections().end();
352 :
353 : // Iterate through the sections.
354 E : for (; section_it != section_end; ++section_it) {
355 E : BlockGraph::Section* section = (*section_it)->section();
356 :
357 : // Stop iterating when we see the relocs.
358 E : if (section->name() == kRelocSectionName) {
359 E : ++section_it;
360 E : break;
361 : }
362 :
363 E : if (!OpenSection(section))
364 i : return false;
365 :
366 : // Iterate over the blocks.
367 : OrderedBlockGraph::BlockList::const_iterator block_it =
368 E : (*section_it)->ordered_blocks().begin();
369 : OrderedBlockGraph::BlockList::const_iterator block_end =
370 E : (*section_it)->ordered_blocks().end();
371 E : for (; block_it != block_end; ++block_it) {
372 E : BlockGraph::Block* block = *block_it;
373 E : if (!LayoutBlock(block))
374 i : return false;
375 E : }
376 :
377 E : if (!CloseSection())
378 i : return false;
379 E : }
380 :
381 : // There should be nothing beyond the relocs, if it was present.
382 E : if (section_it != section_end) {
383 i : LOG(ERROR) << kRelocSectionName << " not the last section.";
384 : }
385 :
386 E : return true;
387 E : }
388 :
389 E : bool ImageLayoutBuilder::Finalize() {
390 E : if (!CreateRelocsSection())
391 i : return false;
392 :
393 E : if (!ReconcileBlockGraphAndImageLayout())
394 i : return false;
395 :
396 E : if (!SortSafeSehTable())
397 i : return false;
398 :
399 E : if (!FinalizeHeaders())
400 i : return false;
401 :
402 E : return true;
403 E : }
404 :
405 : // Lays out a block at the current cursor location.
406 E : bool ImageLayoutBuilder::LayoutBlockImpl(BlockGraph::Block* block) {
407 E : DCHECK(block != NULL);
408 E : if (!image_layout_->blocks.InsertBlock(cursor_, block)) {
409 i : LOG(ERROR) << "InsertBlock failed for block (id=" << block->id()
410 : << ", name=\"" << block->name() << "\").";
411 i : return false;
412 : }
413 E : cursor_ += block->size();
414 E : return true;
415 E : }
416 :
417 E : bool ImageLayoutBuilder::SortSafeSehTable() {
418 E : DCHECK(nt_headers_block_ != NULL);
419 :
420 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
421 E : if (!nt_headers.Init(0, nt_headers_block_)) {
422 i : LOG(ERROR) << "Unable to cast NT headers.";
423 i : return false;
424 : }
425 :
426 E : TypedBlock<IMAGE_LOAD_CONFIG_DIRECTORY> load_config_directory;
427 : if (!nt_headers.Dereference(
428 : nt_headers->OptionalHeader.DataDirectory[
429 : IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG].VirtualAddress,
430 E : &load_config_directory)) {
431 i : LOG(ERROR) << "Failed to dereference Load Config Directory.";
432 i : return false;
433 : }
434 :
435 E : TypedBlock<DWORD> safe_seh_table;
436 : if (!load_config_directory.Dereference(
437 E : load_config_directory->SEHandlerTable, &safe_seh_table)) {
438 : // There's no SEHandlerTable.
439 i : return true;
440 : }
441 :
442 : // Grab the references to the safe SEH code blocks.
443 : typedef BlockGraph::Block::ReferenceMap ReferenceMap;
444 E : const ReferenceMap& orig_references = safe_seh_table.block()->references();
445 :
446 : // We should have as many references as there are handlers and we expect the
447 : // safe seh block to be zero offset and exactly the right size.
448 E : size_t num_references = orig_references.size();
449 : if (num_references != load_config_directory->SEHandlerCount ||
450 : safe_seh_table.offset() != 0 ||
451 E : safe_seh_table.block()->size() != num_references * sizeof(DWORD)) {
452 i : LOG(ERROR) << "Safe SEH Table block does not conform to expectations.";
453 i : return false;
454 : }
455 :
456 : // Create a secondary vector large enough to hold the sorted references.
457 : typedef std::vector<BlockGraph::Reference> ReferenceVector;
458 E : ReferenceVector sorted_references;
459 E : sorted_references.reserve(orig_references.size());
460 :
461 : // Copy the references into a secondary vector.
462 E : for (ReferenceMap::const_iterator iter = orig_references.begin();
463 E : iter != orig_references.end();
464 E : ++iter) {
465 E : sorted_references.push_back(iter->second);
466 E : }
467 :
468 : // Sort the secondary vector in the order their referred blocks appear
469 : // in the image layout.
470 E : RefAddrLess comparator(&image_layout_->blocks);
471 E : std::sort(sorted_references.begin(), sorted_references.end(), comparator);
472 E : if (comparator.failed()) {
473 i : LOG(ERROR) << "One or more exception handler blocks is invalid.";
474 i : return false;
475 : }
476 :
477 : // Reset the references in the Safe SEH Table in sorted order.
478 E : size_t offset = 0;
479 E : for (ReferenceVector::iterator iter = sorted_references.begin();
480 E : iter != sorted_references.end();
481 E : offset += sizeof(DWORD), ++iter) {
482 E : DCHECK(iter->size() == sizeof(DWORD));
483 E : DCHECK(iter->referenced()->type() == BlockGraph::CODE_BLOCK);
484 E : safe_seh_table.block()->SetReference(offset, *iter);
485 E : }
486 :
487 E : return true;
488 E : }
489 :
490 E : bool ImageLayoutBuilder::CreateRelocsSection() {
491 E : RelocWriter writer;
492 :
493 E : DCHECK(nt_headers_block_ != NULL);
494 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
495 E : if (!nt_headers.Init(0, nt_headers_block_)) {
496 i : LOG(ERROR) << "Unable to cast NT headers.";
497 i : return false;
498 : }
499 :
500 : // Get the existing relocs block so we can reuse it.
501 E : TypedBlock<unsigned char> reloc_data;
502 : if (!nt_headers.Dereference(
503 : nt_headers->OptionalHeader.DataDirectory[
504 E : IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress, &reloc_data)) {
505 i : LOG(ERROR) << "Unable to dereference relocs block.";
506 i : return false;
507 : }
508 E : BlockGraph::Block* relocs_block = reloc_data.block();
509 E : CHECK_EQ(0, reloc_data.offset());
510 :
511 : // Iterate over all blocks in the address space, in the
512 : // order of increasing addresses.
513 : BlockGraph::AddressSpace::RangeMap::const_iterator it(
514 E : image_layout_->blocks.address_space_impl().ranges().begin());
515 : BlockGraph::AddressSpace::RangeMap::const_iterator end(
516 E : image_layout_->blocks.address_space_impl().ranges().end());
517 :
518 E : for (; it != end; ++it) {
519 E : const BlockGraph::Block* block = it->second;
520 E : RelativeAddress block_addr;
521 E : CHECK(image_layout_->blocks.GetAddressOf(block, &block_addr));
522 :
523 : // Iterate over all outgoing references in this block in
524 : // order of increasing offset.
525 : BlockGraph::Block::ReferenceMap::const_iterator ref_it(
526 E : block->references().begin());
527 : BlockGraph::Block::ReferenceMap::const_iterator ref_end(
528 E : block->references().end());
529 E : for (; ref_it != ref_end; ++ref_it) {
530 : // Add each absolute reference to the relocs.
531 E : if (ref_it->second.type() == BlockGraph::ABSOLUTE_REF) {
532 E : writer.WriteReloc(block_addr + ref_it->first);
533 : }
534 E : }
535 E : }
536 :
537 : // Get the relocations data from the writer.
538 E : ByteVector relocs;
539 E : writer.Close(&relocs);
540 :
541 : // Update the block and the data directory.
542 E : relocs_block->source_ranges().clear();
543 E : relocs_block->SetData(NULL, 0);
544 E : relocs_block->set_size(relocs.size());
545 E : if (!relocs_block->CopyData(relocs.size(), &relocs.at(0))) {
546 i : LOG(ERROR) << "Unable to copy relocs data.";
547 i : return false;
548 : }
549 : nt_headers->OptionalHeader.DataDirectory[
550 E : IMAGE_DIRECTORY_ENTRY_BASERELOC].Size = relocs.size();
551 :
552 : // Layout the relocs.
553 E : if (!OpenSection(kRelocSectionName, kRelocCharacteristics))
554 i : return false;
555 E : if (!LayoutBlock(relocs_block))
556 i : return false;
557 E : if (!CloseSection())
558 i : return false;
559 :
560 E : return true;
561 E : }
562 :
563 E : bool ImageLayoutBuilder::ReconcileBlockGraphAndImageLayout() {
564 : // Get the reloc section ID from the block-graph.
565 : BlockGraph::Section* reloc_section =
566 E : image_layout_->blocks.graph()->FindSection(kRelocSectionName);
567 E : if (reloc_section == NULL) {
568 i : LOG(ERROR) << "Unable to find the reloc section in the block-graph.";
569 i : return false;
570 : }
571 E : BlockGraph::SectionId reloc_section_id = reloc_section->id();
572 :
573 : // Iterate over the blocks of the block-graph to see if some of them are not
574 : // in the image layout. If we find one we check if it belongs to the reloc
575 : // section, in this case we put it in a list of blocks that we should remove
576 : // from the graph, otherwise we return an error.
577 : BlockGraph::BlockMap::iterator it_block_graph =
578 E : image_layout_->blocks.graph()->blocks_mutable().begin();
579 E : std::list<BlockGraph::Block*> blocks_to_remove;
580 :
581 E : for (; it_block_graph != image_layout_->blocks.graph()->blocks().end();
582 E : it_block_graph++) {
583 : // Determine if the current block exist in the image layout.
584 E : if (!image_layout_->blocks.ContainsBlock(&it_block_graph->second)) {
585 : // If it doesn't we check to see if this block belongs to the reloc
586 : // section.
587 E : if (it_block_graph->second.section() != reloc_section_id) {
588 i : LOG(ERROR) << "There is a block in the block-graph that is not in the "
589 : << "image layout (id=" << it_block_graph->second.id()
590 : << ", name=\"" << it_block_graph->second.name() << "\", "
591 : << "original address=" << it_block_graph->second.addr()
592 : << ").";
593 i : return false;
594 : } else {
595 : // The block is added to the list of blocks to remove from the graph.
596 E : blocks_to_remove.push_back(&it_block_graph->second);
597 : }
598 : }
599 E : }
600 :
601 : // The useless blocks are removed from the block-graph.
602 : std::list<BlockGraph::Block*>::iterator iter_blocks =
603 E : blocks_to_remove.begin();
604 E : for (; iter_blocks != blocks_to_remove.end(); iter_blocks++) {
605 E : if (!image_layout_->blocks.graph()->RemoveBlock(*iter_blocks)) {
606 i : LOG(ERROR) << "Unable to remove block with ID " << (*iter_blocks)->id()
607 : << " from the block-graph.";
608 : }
609 E : }
610 :
611 : DCHECK_EQ(image_layout_->blocks.size(),
612 E : image_layout_->blocks.graph()->blocks().size());
613 :
614 E : return true;
615 E : }
616 :
617 E : bool ImageLayoutBuilder::FinalizeHeaders() {
618 : // The DOS and NT headers must be set at this point.
619 E : DCHECK(dos_header_block_ != NULL);
620 E : DCHECK(nt_headers_block_ != NULL);
621 :
622 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
623 E : if (!nt_headers.Init(0, nt_headers_block_)) {
624 i : LOG(ERROR) << "Unable to cast NT headers.";
625 i : return false;
626 : }
627 :
628 E : TypedBlock<IMAGE_SECTION_HEADER> section_headers;
629 E : if (!section_headers.Init(sizeof(IMAGE_NT_HEADERS), nt_headers_block_)) {
630 i : LOG(ERROR) << "Unable to cast section headers.";
631 i : return false;
632 : }
633 :
634 : // Ensure the section headers have the expected size. If they don't we bail,
635 : // as this should have been done prior to layout (PrepareHeadersTransform).
636 E : if (section_headers.ElementCount() != image_layout_->sections.size()) {
637 i : LOG(ERROR) << "Section header count does not agree with layout section "
638 : << "count (" << section_headers.ElementCount() << " != "
639 : << image_layout_->sections.size() << ").";
640 i : return false;
641 : }
642 :
643 : core::FileOffsetAddress section_file_start(
644 E : nt_headers->OptionalHeader.SizeOfHeaders);
645 :
646 : // Iterate through our sections to initialize the code/data fields in the NT
647 : // headers.
648 E : for (size_t i = 0; i < image_layout_->sections.size(); ++i) {
649 E : const ImageLayout::SectionInfo& section = image_layout_->sections[i];
650 E : IMAGE_SECTION_HEADER& hdr = section_headers[i];
651 :
652 E : if (section.characteristics& IMAGE_SCN_CNT_CODE) {
653 E : nt_headers->OptionalHeader.SizeOfCode += section.data_size;
654 E : if (nt_headers->OptionalHeader.BaseOfCode == 0) {
655 i : nt_headers->OptionalHeader.BaseOfCode = section.addr.value();
656 : }
657 : }
658 E : if (section.characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA) {
659 E : nt_headers->OptionalHeader.SizeOfInitializedData += section.data_size;
660 :
661 E : if (nt_headers->OptionalHeader.BaseOfData == 0)
662 i : nt_headers->OptionalHeader.BaseOfData = section.addr.value();
663 : }
664 E : if (section.characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA) {
665 : nt_headers->OptionalHeader.SizeOfUninitializedData +=
666 i : section.data_size;
667 i : if (nt_headers->OptionalHeader.BaseOfData == 0)
668 i : nt_headers->OptionalHeader.BaseOfData = section.addr.value();
669 : }
670 :
671 : // Zero the header to get rid of any old crud in it.
672 E : memset(&hdr, 0, sizeof(hdr));
673 :
674 : strncpy(reinterpret_cast<char*>(hdr.Name),
675 : section.name.c_str(),
676 E : arraysize(hdr.Name));
677 E : hdr.Misc.VirtualSize = section.size;
678 E : hdr.VirtualAddress = section.addr.value();
679 E : hdr.SizeOfRawData = section.data_size;
680 E : hdr.PointerToRawData = section_file_start.value();
681 E : hdr.Characteristics = section.characteristics;
682 :
683 E : section_file_start += section.data_size;
684 E : }
685 :
686 : nt_headers->OptionalHeader.SizeOfImage =
687 E : cursor_.AlignUp(nt_headers->OptionalHeader.SectionAlignment).value();
688 :
689 E : return true;
690 E : }
691 :
692 : } // namespace pe
|