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