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_file_writer.h"
16 :
17 : #include <windows.h>
18 : #include <winnt.h>
19 : #include <imagehlp.h> // NOLINT
20 :
21 : #include "base/file_util.h"
22 : #include "base/logging.h"
23 : #include "base/win/scoped_handle.h"
24 : #include "syzygy/common/buffer_parser.h"
25 : #include "syzygy/common/com_utils.h"
26 : #include "syzygy/pe/pe_utils.h"
27 :
28 : namespace pe {
29 :
30 : using block_graph::BlockGraph;
31 : using common::BinaryBufferParser;
32 : using core::AbsoluteAddress;
33 : using core::FileOffsetAddress;
34 : using core::RelativeAddress;
35 : using pe::ImageLayout;
36 :
37 : namespace {
38 :
39 : template <class Type>
40 E : bool UpdateReference(size_t start, Type new_value, std::vector<uint8>* data) {
41 E : BinaryBufferParser parser(&data->at(0), data->size());
42 :
43 E : Type* ref_ptr = NULL;
44 E : if (!parser.GetAt(start, const_cast<const Type**>(&ref_ptr))) {
45 i : LOG(ERROR) << "Reference data not in block";
46 i : return false;
47 : }
48 E : *ref_ptr = new_value;
49 :
50 E : return true;
51 E : }
52 :
53 : // Returns the type of padding byte to use for a given section. Int3s will be
54 : // used for executable sections, nulls for everything else.
55 : uint8 GetSectionPaddingByte(const ImageLayout& image_layout,
56 E : size_t section_index) {
57 E : const uint8 kZero = 0;
58 E : const uint8 kInt3 = 0xCC;
59 :
60 E : if (section_index == BlockGraph::kInvalidSectionId)
61 E : return kZero;
62 E : DCHECK_GT(image_layout.sections.size(), section_index);
63 :
64 : const ImageLayout::SectionInfo& section_info =
65 E : image_layout.sections[section_index];
66 : bool is_executable =
67 E : (section_info.characteristics & IMAGE_SCN_MEM_EXECUTE) != 0;
68 E : if (is_executable)
69 E : return kInt3;
70 E : return kZero;
71 E : }
72 :
73 : // Returns the length of explicitly initialized data in a block.
74 : // TODO(chrisha): Move this to block_util and unittest it.
75 E : size_t GetBlockInitializedDataSize(const BlockGraph::Block* block) {
76 E : DCHECK(block != NULL);
77 :
78 : // All references contain initialized data so must be explicitly written.
79 : // Use the position and the size of the last offset as the initialized length.
80 E : size_t length = 0;
81 E : if (!block->references().empty()) {
82 : BlockGraph::Block::ReferenceMap::const_reverse_iterator ref_it =
83 E : block->references().rbegin();
84 E : length = ref_it->first + ref_it->second.size();
85 : }
86 :
87 : // Otherwise, we use the block data size.
88 : // TODO(chrisha): If we really wanted to, we could strip off trailing zeros
89 : // from the block data, but that's maybe a little overkill.
90 E : length = std::max(length, block->data_size());
91 :
92 E : return length;
93 E : }
94 :
95 : size_t GetSectionOffset(const ImageLayout& image_layout,
96 : const RelativeAddress rel_addr,
97 E : size_t section_index) {
98 E : if (section_index == BlockGraph::kInvalidSectionId)
99 i : return rel_addr.value();
100 :
101 E : DCHECK_GT(image_layout.sections.size(), section_index);
102 : const ImageLayout::SectionInfo& section_info =
103 E : image_layout.sections[section_index];
104 :
105 E : DCHECK_GE(rel_addr, section_info.addr);
106 E : return rel_addr - section_info.addr;
107 E : }
108 :
109 : } // namespace
110 :
111 : PEFileWriter::PEFileWriter(const ImageLayout& image_layout)
112 E : : image_layout_(image_layout), nt_headers_(NULL) {
113 E : }
114 :
115 E : bool PEFileWriter::WriteImage(const base::FilePath& path) {
116 : // Start by attempting to open the destination file.
117 E : base::ScopedFILE file(base::OpenFile(path, "wb"));
118 E : if (file.get() == NULL) {
119 i : LOG(ERROR) << "Unable to open " << path.value();
120 i : return false;
121 : }
122 :
123 E : if (!ValidateHeaders())
124 E : return false;
125 :
126 E : DCHECK(nt_headers_ != NULL);
127 :
128 E : bool success = CalculateSectionRanges();
129 E : if (success)
130 E : success = WriteBlocks(file.get());
131 :
132 E : nt_headers_ = NULL;
133 :
134 : // Close the file.
135 E : file.reset();
136 :
137 E : if (success)
138 E : success = UpdateFileChecksum(path);
139 :
140 E : return success;
141 E : }
142 :
143 E : bool PEFileWriter::UpdateFileChecksum(const base::FilePath& path) {
144 : // Open the image file for exclusive write.
145 : base::win::ScopedHandle image_handle(
146 : ::CreateFile(path.value().c_str(), GENERIC_READ | GENERIC_WRITE, 0,
147 E : NULL, OPEN_EXISTING, 0, NULL));
148 :
149 E : if (!image_handle.IsValid()) {
150 E : LOG(ERROR) << "Failed to open file " << path.value();
151 E : return false;
152 : }
153 :
154 E : size_t file_size = ::GetFileSize(image_handle.Get(), NULL);
155 :
156 : // Create an anonymous read/write mapping on the file.
157 : base::win::ScopedHandle image_mapping(::CreateFileMapping(image_handle.Get(),
158 : NULL,
159 : PAGE_READWRITE,
160 : 0,
161 : 0,
162 E : NULL));
163 : // Map the entire file read/write to memory.
164 E : void* image_ptr = NULL;
165 :
166 E : if (image_mapping.IsValid()) {
167 : image_ptr = ::MapViewOfFile(image_mapping.Get(),
168 : FILE_MAP_WRITE,
169 : 0,
170 : 0,
171 E : file_size);
172 : }
173 :
174 E : if (image_ptr == NULL) {
175 E : LOG(ERROR) << "Failed to create image mapping.";
176 E : return false;
177 : }
178 :
179 : // Calculate the image checksum.
180 E : DWORD original_checksum = 0;
181 E : DWORD new_checksum = 0;
182 : IMAGE_NT_HEADERS* nt_headers = ::CheckSumMappedFile(image_ptr,
183 : file_size,
184 : &original_checksum,
185 E : &new_checksum);
186 :
187 E : if (nt_headers == NULL) {
188 i : DWORD error = ::GetLastError();
189 i : LOG(ERROR) << "CheckSumMappedFile failed: " << common::LogWe(error);
190 : }
191 :
192 : // On success, we write the checksum back to the file header.
193 E : if (nt_headers != NULL) {
194 E : nt_headers->OptionalHeader.CheckSum = new_checksum;
195 : }
196 E : CHECK(::UnmapViewOfFile(image_ptr));
197 :
198 E : return nt_headers != NULL;
199 E : }
200 :
201 E : bool PEFileWriter::ValidateHeaders() {
202 E : DCHECK(nt_headers_ == NULL);
203 :
204 : // Get the DOS header block.
205 : BlockGraph::Block* dos_header_block =
206 E : image_layout_.blocks.GetBlockByAddress(RelativeAddress(0));
207 E : if (dos_header_block == NULL) {
208 i : LOG(ERROR) << "No DOS header in image.";
209 i : return false;
210 : }
211 E : if (!IsValidDosHeaderBlock(dos_header_block)) {
212 E : LOG(ERROR) << "Invalid DOS header in image.";
213 E : return false;
214 : }
215 : BlockGraph::Block* nt_headers_block =
216 E : GetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
217 E : DCHECK(nt_headers_block != NULL);
218 :
219 : const IMAGE_NT_HEADERS* nt_headers =
220 E : reinterpret_cast<const IMAGE_NT_HEADERS*>(nt_headers_block->data());
221 E : DCHECK(nt_headers != NULL);
222 :
223 E : nt_headers_ = nt_headers;
224 :
225 E : return true;
226 E : }
227 :
228 E : bool PEFileWriter::CalculateSectionRanges() {
229 E : DCHECK(nt_headers_ != NULL);
230 E : DCHECK_EQ(0u, section_file_range_map_.size());
231 E : DCHECK_EQ(0u, section_index_space_.size());
232 :
233 E : size_t section_alignment = nt_headers_->OptionalHeader.SectionAlignment;
234 E : size_t file_alignment = nt_headers_->OptionalHeader.FileAlignment;
235 :
236 : // Keep track of the end of each section, both in memory and on disk.
237 : RelativeAddress previous_section_end =
238 : RelativeAddress(nt_headers_->OptionalHeader.SizeOfHeaders).AlignUp(
239 E : section_alignment);
240 : FileOffsetAddress previous_section_file_end =
241 : FileOffsetAddress(nt_headers_->OptionalHeader.SizeOfHeaders).AlignUp(
242 E : file_alignment);
243 :
244 : // This first range is for the header and doesn't correspond to any section.
245 : CHECK(section_file_range_map_.insert(
246 : std::make_pair(BlockGraph::kInvalidSectionId,
247 : FileRange(FileOffsetAddress(0),
248 E : previous_section_file_end.value()))).second);
249 :
250 : // Validate the number of sections in the headers.
251 : if (nt_headers_->FileHeader.NumberOfSections !=
252 E : image_layout_.sections.size()) {
253 E : LOG(ERROR) << "NT headers section count mismatch.";
254 E : return false;
255 : }
256 :
257 : // The remainder of the mappings are for the sections. While we run through
258 : // and calculate the section ranges, we also make sure they're sane by
259 : // checking that:
260 : // - they're arranged sequentially,
261 : // - there are no gaps between sections,
262 : // - that they don't run into one another.
263 :
264 E : IMAGE_SECTION_HEADER* section_header = IMAGE_FIRST_SECTION(nt_headers_);
265 E : for (size_t i = 0; i < image_layout_.sections.size(); ++i, ++section_header) {
266 E : const ImageLayout::SectionInfo& section = image_layout_.sections[i];
267 E : RelativeAddress section_start(section.addr);
268 E : size_t section_size = section.size;
269 :
270 : // Calculate the file offset start for this section.
271 : FileOffsetAddress section_file_start =
272 E : previous_section_file_end.AlignUp(file_alignment);
273 E : size_t section_file_size = section.data_size;
274 :
275 : // Validate that the section doesn't overlap in memory or on disk.
276 : if (section_start < previous_section_end ||
277 E : section_file_start < previous_section_file_end) {
278 E : LOG(ERROR) << "Section " << section.name <<
279 : " runs into previous section (or header).";
280 E : return false;
281 : }
282 :
283 : // Validate the alignment of the section start addresses in memory and on
284 : // disk.
285 : if ((section_start.value() % section_alignment) != 0 ||
286 E : (section_file_start.value() % file_alignment) != 0) {
287 E : LOG(ERROR) << "Section " << section.name <<
288 : " has incorrect alignment.";
289 E : return false;
290 : }
291 :
292 : // Make sure there are no unexpected gaps between sections (the packing
293 : // should be as tight as possible).
294 : if ((section_start - previous_section_end >
295 : static_cast<ptrdiff_t>(section_alignment)) ||
296 : (section_file_start - previous_section_file_end >
297 E : static_cast<ptrdiff_t>(file_alignment))) {
298 i : LOG(ERROR) << "Section " << section.name <<
299 : " leaves a gap from previous section.";
300 i : return false;
301 : }
302 :
303 : // Ok, it all passes inspection so far. Record the mapping.
304 E : FileRange section_file_range(section_file_start, section_file_size);
305 : CHECK(section_file_range_map_.insert(
306 E : std::make_pair(i, section_file_range)).second);
307 :
308 : CHECK(section_index_space_.Insert(
309 E : SectionIndexSpace::Range(section_start, section_size), i, NULL));
310 :
311 : // Validate that the NT section headers match what we calculate.
312 : if (section_header->VirtualAddress != section_start.value() ||
313 : section_header->SizeOfRawData != section_file_size ||
314 : section_header->PointerToRawData != section_file_start.value() ||
315 E : section_header->Misc.VirtualSize != section_size) {
316 E : LOG(ERROR) << "NT section headers are inconsistent with image layout.";
317 E : return false;
318 : }
319 :
320 E : previous_section_end = section_start + section_size;
321 E : previous_section_file_end = section_file_start + section_file_size;
322 E : }
323 :
324 E : return true;
325 E : }
326 :
327 E : bool PEFileWriter::WriteBlocks(FILE* file) {
328 E : DCHECK(file != NULL);
329 :
330 E : AbsoluteAddress image_base(nt_headers_->OptionalHeader.ImageBase);
331 :
332 : // Create the output buffer, reserving enough room for the whole file.
333 E : DCHECK(!image_layout_.sections.empty());
334 E : size_t last_section_index = image_layout_.sections.size() - 1;
335 E : size_t image_size = section_file_range_map_[last_section_index].end().value();
336 E : std::vector<uint8> buffer;
337 E : buffer.reserve(image_size);
338 :
339 : // Iterate through all blocks in the address space writing them as we go.
340 : BlockGraph::AddressSpace::RangeMap::const_iterator block_it(
341 E : image_layout_.blocks.address_space_impl().ranges().begin());
342 : BlockGraph::AddressSpace::RangeMap::const_iterator block_end(
343 E : image_layout_.blocks.address_space_impl().ranges().end());
344 :
345 : // Write all of the blocks. We take care of writing the padding at the
346 : // end of each section. Note that the section index is not the same thing as
347 : // the section_id stored in the block; the section IDs are relative to the
348 : // section data stored in the block-graph, not the ordered section infos
349 : // stored in the image layout.
350 E : BlockGraph::SectionId section_id = BlockGraph::kInvalidSectionId;
351 E : size_t section_index = BlockGraph::kInvalidSectionId;
352 E : for (; block_it != block_end; ++block_it) {
353 E : BlockGraph::Block* block = const_cast<BlockGraph::Block*>(block_it->second);
354 :
355 : // If we're jumping to a new section output the necessary padding.
356 E : if (block->section() != section_id) {
357 E : FlushSection(section_index, &buffer);
358 E : section_id = block->section();
359 E : section_index++;
360 E : DCHECK_GT(image_layout_.sections.size(), section_index);
361 : }
362 :
363 E : if (!WriteOneBlock(image_base, section_index, block, &buffer)) {
364 E : LOG(ERROR) << "Failed to write block \"" << block->name() << "\".";
365 E : return false;
366 : }
367 E : }
368 :
369 E : FlushSection(last_section_index, &buffer);
370 E : DCHECK_EQ(image_size, buffer.size());
371 :
372 : // Write the whole image to disk in one go.
373 : if (::fwrite(&buffer[0], sizeof(buffer[0]), buffer.size(), file) !=
374 E : buffer.size()) {
375 i : LOG(ERROR) << "Failed to write image to file.";
376 i : return false;
377 : }
378 :
379 E : return true;
380 E : }
381 :
382 : void PEFileWriter::FlushSection(size_t section_index,
383 E : std::vector<uint8>* buffer) {
384 E : DCHECK(buffer != NULL);
385 :
386 : size_t section_file_end =
387 E : section_file_range_map_[section_index].end().value();
388 :
389 : // We've already sanity checked this in CalculateSectionFileRanges, so this
390 : // should be true.
391 E : DCHECK_GE(section_file_end, buffer->size());
392 E : if (section_file_end == buffer->size())
393 E : return;
394 :
395 E : uint8 padding_byte = GetSectionPaddingByte(image_layout_, section_index);
396 E : buffer->resize(section_file_end, padding_byte);
397 :
398 : return;
399 E : }
400 :
401 : bool PEFileWriter::WriteOneBlock(AbsoluteAddress image_base,
402 : size_t section_index,
403 : const BlockGraph::Block* block,
404 E : std::vector<uint8>* buffer) {
405 : // This function walks through the data referred by the input block, and
406 : // patches it to reflect the addresses and offsets of the blocks
407 : // referenced before writing the block's data to the file.
408 E : DCHECK(block != NULL);
409 E : DCHECK(buffer != NULL);
410 :
411 E : RelativeAddress addr;
412 E : if (!image_layout_.blocks.GetAddressOf(block, &addr)) {
413 i : LOG(ERROR) << "All blocks must have an address.";
414 i : return false;
415 : }
416 :
417 : // Get the start address of the section containing this block as well as the
418 : // padding byte we need to use.
419 E : RelativeAddress section_start(0);
420 E : RelativeAddress section_end(image_layout_.sections[0].addr);
421 E : uint8 padding_byte = GetSectionPaddingByte(image_layout_, section_index);
422 E : if (section_index != BlockGraph::kInvalidSectionId) {
423 : const ImageLayout::SectionInfo& section_info =
424 E : image_layout_.sections[section_index];
425 E : section_start = RelativeAddress(section_info.addr);
426 E : section_end = section_start + section_info.size;
427 : }
428 :
429 E : const FileRange& section_file_range = section_file_range_map_[section_index];
430 :
431 : // The block should lie entirely within the section.
432 E : if (addr < section_start || addr + block->size() > section_end) {
433 E : LOG(ERROR) << "Block lies outside of section.";
434 E : return false;
435 : }
436 :
437 : // Calculate the offset from the start of the section to
438 : // the start of the block, and the block's file offset.
439 E : BlockGraph::Offset section_offs = addr - section_start;
440 E : FileOffsetAddress file_offs = section_file_range.start() + section_offs;
441 :
442 : // We shouldn't have written anything to the spot where the block belongs.
443 : // This is only a DCHECK because the address space of the image layout and
444 : // the consistency of the sections guarantees this for us.
445 E : DCHECK_LE(buffer->size(), file_offs.value());
446 :
447 E : size_t inited_data_size = GetBlockInitializedDataSize(block);
448 :
449 : // If this block is entirely in the virtual portion of the section, skip it.
450 E : if (file_offs >= section_file_range.end()) {
451 E : if (inited_data_size != 0) {
452 E : LOG(ERROR) << "Block contains explicit data or references but is in "
453 : << "virtual portion of section.";
454 E : return false;
455 : }
456 :
457 E : return true;
458 : }
459 :
460 : // The initialized portion of data for this block must lie entirely within the
461 : // initialized data for this section (this includes references to be filled in
462 : // and the explicit block data).
463 E : if (file_offs + inited_data_size > section_file_range.end()) {
464 i : LOG(ERROR) << "Initialized portion of block data lies outside of section.";
465 i : return false;
466 : }
467 :
468 : // Add any necessary padding to get us to the block offset.
469 E : if (buffer->size() < file_offs.value())
470 E : buffer->resize(file_offs.value(), padding_byte);
471 :
472 : // Copy the block data into the buffer.
473 : buffer->insert(buffer->end(),
474 : block->data(),
475 E : block->data() + block->data_size());
476 :
477 : // We now want to append zeros for the implicit portion of the block data.
478 E : size_t trailing_zeros = block->size() - block->data_size();
479 E : if (trailing_zeros > 0) {
480 : // It is possible for a block to be laid out at the end of a section such
481 : // that part of its data lies within the virtual portion of the section.
482 : // Since padding between blocks can be non-zero we explicitly write out any
483 : // trailing zeros here. So use the section size to determine how much we are
484 : // supposed to write.
485 E : FileOffsetAddress block_file_end = file_offs + block->size();
486 E : if (block_file_end > section_file_range.end()) {
487 : size_t implicit_trailing_zeros =
488 E : block_file_end - section_file_range.end();
489 E : DCHECK_LE(implicit_trailing_zeros, trailing_zeros);
490 E : trailing_zeros -= implicit_trailing_zeros;
491 : }
492 :
493 : // Write the implicit trailing zeros.
494 E : buffer->insert(buffer->end(), trailing_zeros, 0);
495 : }
496 :
497 : // Patch up all the references.
498 : BlockGraph::Block::ReferenceMap::const_iterator ref_it(
499 E : block->references().begin());
500 : BlockGraph::Block::ReferenceMap::const_iterator ref_end(
501 E : block->references().end());
502 E : for (; ref_it != ref_end; ++ref_it) {
503 E : BlockGraph::Offset start = ref_it->first;
504 E : const BlockGraph::Reference& ref = ref_it->second;
505 E : BlockGraph::Block* dst = ref.referenced();
506 :
507 E : RelativeAddress src_addr(addr + start);
508 E : RelativeAddress dst_addr;
509 E : if (!image_layout_.blocks.GetAddressOf(dst, &dst_addr)) {
510 i : LOG(ERROR) << "All blocks must have an address.";
511 i : return false;
512 : }
513 E : dst_addr += ref.offset();
514 :
515 : // Compute the new value of the reference.
516 E : uint32 value = 0;
517 E : switch (ref.type()) {
518 : case BlockGraph::ABSOLUTE_REF:
519 E : value = image_base.value() + dst_addr.value();
520 E : break;
521 :
522 : case BlockGraph::PC_RELATIVE_REF:
523 E : value = dst_addr - (src_addr + ref.size());
524 E : break;
525 :
526 : case BlockGraph::RELATIVE_REF:
527 E : value = dst_addr.value();
528 E : break;
529 :
530 : case BlockGraph::FILE_OFFSET_REF: {
531 : // Get the index of the section containing the destination block.
532 : SectionIndexSpace::const_iterator section_index_space_it =
533 : section_index_space_.FindContaining(
534 E : SectionIndexSpace::Range(dst_addr, 1));
535 E : DCHECK(section_index_space_it != section_index_space_.end());
536 E : size_t dst_section_index = section_index_space_it->second;
537 :
538 : // Get the offset of the block in its section, as well as the range of
539 : // the section on disk. Validate that the referred location is
540 : // actually directly represented on disk (not in implicit virtual data).
541 : const FileRange& file_range =
542 E : section_file_range_map_[dst_section_index];
543 : size_t section_offset = GetSectionOffset(image_layout_,
544 : dst_addr,
545 E : dst_section_index);
546 E : if (section_offset >= file_range.size()) {
547 E : LOG(ERROR) << "Encountered file offset reference that refers to "
548 : << "a location outside of the explicit section data.";
549 E : return false;
550 : }
551 :
552 : // Finally, calculate the value of the file offset.
553 E : value = file_range.start().value() + section_offset;
554 E : break;
555 : }
556 :
557 : default:
558 i : LOG(ERROR) << "Impossible reference type";
559 i : return false;
560 : }
561 :
562 : // Now store the new value.
563 E : BlockGraph::Offset ref_offset = file_offs.value() + start;
564 E : switch (ref.size()) {
565 : case sizeof(uint8):
566 E : if (!UpdateReference(ref_offset, static_cast<uint8>(value), buffer))
567 i : return false;
568 E : break;
569 :
570 : case sizeof(uint16):
571 i : if (!UpdateReference(ref_offset, static_cast<uint16>(value), buffer))
572 i : return false;
573 i : break;
574 :
575 : case sizeof(uint32):
576 E : if (!UpdateReference(ref_offset, static_cast<uint32>(value), buffer))
577 i : return false;
578 E : break;
579 :
580 : default:
581 i : LOG(ERROR) << "Unsupported reference size.";
582 i : return false;
583 : }
584 E : }
585 :
586 E : return true;
587 E : }
588 :
589 : } // namespace pe
|