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