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