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/pe_file_writer.h"
16 :
17 : #include <windows.h>
18 : #include <winnt.h>
19 : #include <imagehlp.h>
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/common/align.h"
27 : #include "syzygy/pe/pe_utils.h"
28 :
29 : namespace pe {
30 :
31 : using block_graph::BlockGraph;
32 : using core::AbsoluteAddress;
33 : using core::FileOffsetAddress;
34 : using core::RelativeAddress;
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 : } // namespace
53 :
54 : PEFileWriter::PEFileWriter(const ImageLayout& image_layout)
55 E : : image_layout_(image_layout), nt_headers_(NULL) {
56 E : }
57 :
58 E : bool PEFileWriter::WriteImage(const FilePath& path) {
59 : // Start by attempting to open the destination file.
60 E : file_util::ScopedFILE file(file_util::OpenFile(path, "wb"));
61 E : if (file.get() == NULL) {
62 i : LOG(ERROR) << "Unable to open " << path.value();
63 i : return false;
64 : }
65 :
66 E : if (!ValidateHeaders())
67 i : return false;
68 :
69 E : DCHECK(nt_headers_ != NULL);
70 :
71 E : bool success = InitializeSectionFileAddressSpace();
72 E : if (success)
73 E : success = WriteBlocks(file.get());
74 :
75 E : nt_headers_ = NULL;
76 :
77 : // Close the file.
78 E : file.reset();
79 :
80 E : if (success)
81 E : success = UpdateFileChecksum(path);
82 :
83 E : return success;
84 E : }
85 :
86 E : bool PEFileWriter::UpdateFileChecksum(const FilePath& path) {
87 : // Open the image file for exclusive write.
88 : base::win::ScopedHandle image_handle(
89 : ::CreateFile(path.value().c_str(), GENERIC_READ | GENERIC_WRITE, 0,
90 E : NULL, OPEN_EXISTING, 0, NULL));
91 :
92 E : if (!image_handle.IsValid()) {
93 E : LOG(ERROR) << "Failed to open file " << path.value();
94 E : return false;
95 : }
96 :
97 E : size_t file_size = ::GetFileSize(image_handle.Get(), NULL);
98 :
99 : // Create an anonymous read/write mapping on the file.
100 : base::win::ScopedHandle image_mapping(::CreateFileMapping(image_handle.Get(),
101 : NULL,
102 : PAGE_READWRITE,
103 : 0,
104 : 0,
105 E : NULL));
106 : // Map the entire file read/write to memory.
107 E : void* image_ptr = NULL;
108 :
109 E : if (image_mapping.IsValid()) {
110 : image_ptr = ::MapViewOfFile(image_mapping.Get(),
111 : FILE_MAP_WRITE,
112 : 0,
113 : 0,
114 E : file_size);
115 : }
116 :
117 E : if (image_ptr == NULL) {
118 E : LOG(ERROR) << "Failed to create image mapping.";
119 E : return false;
120 : }
121 :
122 : // Calculate the image checksum.
123 E : DWORD original_checksum = 0;
124 E : DWORD new_checksum = 0;
125 : IMAGE_NT_HEADERS* nt_headers = ::CheckSumMappedFile(image_ptr,
126 : file_size,
127 : &original_checksum,
128 E : &new_checksum);
129 :
130 E : if (nt_headers == NULL) {
131 i : DWORD error = ::GetLastError();
132 i : LOG(ERROR) << "CheckSumMappedFile failed: " << com::LogWe(error);
133 : }
134 :
135 : // On success, we write the checksum back to the file header.
136 E : if (nt_headers != NULL) {
137 E : nt_headers->OptionalHeader.CheckSum = new_checksum;
138 : }
139 E : CHECK(::UnmapViewOfFile(image_ptr));
140 :
141 E : return nt_headers != NULL;
142 E : }
143 :
144 E : bool PEFileWriter::ValidateHeaders() {
145 E : DCHECK(nt_headers_ == NULL);
146 :
147 : // Get the DOS header block.
148 : BlockGraph::Block* dos_header_block =
149 E : image_layout_.blocks.GetBlockByAddress(RelativeAddress(0));
150 E : if (dos_header_block == NULL) {
151 i : LOG(ERROR) << "No DOS header in image.";
152 i : return false;
153 : }
154 E : if (!IsValidDosHeaderBlock(dos_header_block)) {
155 i : LOG(ERROR) << "Invalid DOS header in image.";
156 i : return false;
157 : }
158 : BlockGraph::Block* nt_headers_block =
159 E : GetNtHeadersBlockFromDosHeaderBlock(dos_header_block);
160 E : DCHECK(nt_headers_block != NULL);
161 :
162 : const IMAGE_NT_HEADERS* nt_headers =
163 E : reinterpret_cast<const IMAGE_NT_HEADERS*>(nt_headers_block->data());
164 E : DCHECK(nt_headers != NULL);
165 :
166 : // TODO(siggi): Validate that NT headers and ImageLayout match, or else
167 : // from an AddressSpace, and forget the ImageLayout.
168 E : nt_headers_ = nt_headers;
169 :
170 E : return true;
171 E : }
172 :
173 E : bool PEFileWriter::InitializeSectionFileAddressSpace() {
174 E : DCHECK(nt_headers_ != NULL);
175 :
176 : // Now set up the address mappings from RVA to disk offset for the entire
177 : // image. The first mapping starts at zero, and covers the header(s).
178 : SectionFileAddressSpace::Range header_range(
179 E : RelativeAddress(0), nt_headers_->OptionalHeader.SizeOfHeaders);
180 E : section_file_offsets_.Insert(header_range, FileOffsetAddress(0));
181 :
182 : // The remainder of the mappings are for the sections. While we run through
183 : // and set up the section mappings, we also make sure they're sane by
184 : // checking that:
185 : // - they're arranged sequentially,
186 : // - there are no gaps between sections,
187 : // - that they don't run into one another.
188 : RelativeAddress previous_section_end(
189 E : header_range.start() + header_range.size());
190 E : FileOffsetAddress previous_section_file_end(previous_section_end.value());
191 :
192 E : for (size_t i = 0; i < image_layout_.sections.size(); ++i) {
193 E : const ImageLayout::SectionInfo& section = image_layout_.sections[i];
194 E : RelativeAddress section_start(section.addr);
195 E : size_t section_size = section.size;
196 : // Calculate the file offset start for this section.
197 E : FileOffsetAddress section_file_start(previous_section_file_end);
198 E : size_t section_file_size = section.data_size;
199 :
200 : if (section_start < previous_section_end ||
201 E : section_file_start < previous_section_file_end) {
202 i : LOG(ERROR) << "Section " << section.name <<
203 : " runs into previous section (or header).";
204 i : return false;
205 : }
206 :
207 : if ((section_start.value() %
208 : nt_headers_->OptionalHeader.SectionAlignment) != 0 ||
209 : (section_file_start.value() %
210 E : nt_headers_->OptionalHeader.FileAlignment) != 0) {
211 i : LOG(ERROR) << "Section " << section.name <<
212 : " has incorrect alignment.";
213 i : return false;
214 : }
215 :
216 : if ((section_start - previous_section_end > static_cast<ptrdiff_t>(
217 : nt_headers_->OptionalHeader.SectionAlignment)) ||
218 : (section_file_start - previous_section_file_end >
219 : static_cast<ptrdiff_t>(
220 E : nt_headers_->OptionalHeader.FileAlignment))) {
221 i : LOG(ERROR) << "Section " << section.name <<
222 : " leaves a gap from previous section.";
223 i : return false;
224 : }
225 :
226 : // Ok, it all passes inspection so far. If the file size is non-zero,
227 : // go ahead and record the mapping.
228 E : size_t file_size = section.data_size;
229 E : if (file_size != 0) {
230 E : SectionFileAddressSpace::Range file_range(section_start, file_size);
231 E : section_file_offsets_.Insert(file_range, section_file_start);
232 : }
233 :
234 E : previous_section_end = section_start + section_size;
235 E : previous_section_file_end = section_file_start + section_file_size;
236 E : }
237 :
238 E : return true;
239 E : }
240 :
241 E : bool PEFileWriter::WriteBlocks(FILE* file) {
242 E : AbsoluteAddress image_base(nt_headers_->OptionalHeader.ImageBase);
243 :
244 : // Iterate through all blocks in the address space.
245 : BlockGraph::AddressSpace::RangeMap::const_iterator it(
246 E : image_layout_.blocks.address_space_impl().ranges().begin());
247 : BlockGraph::AddressSpace::RangeMap::const_iterator end(
248 E : image_layout_.blocks.address_space_impl().ranges().end());
249 :
250 E : for (; it != end; ++it) {
251 E : BlockGraph::Block* block = const_cast<BlockGraph::Block*>(it->second);
252 :
253 E : if (!WriteOneBlock(image_base, block, file)) {
254 i : LOG(ERROR) << "Failed to write block " << block->name();
255 i : return false;
256 : }
257 E : }
258 :
259 : // Now round the file to the required size.
260 E : if (image_layout_.sections.size() == 0) {
261 i : LOG(ERROR) << "Missing or corrupt image section headers";
262 i : return false;
263 : }
264 :
265 E : const ImageLayout::SectionInfo& last_section = image_layout_.sections.back();
266 E : if (last_section.data_size > last_section.size) {
267 E : DCHECK_LT(0u, section_file_offsets_.size());
268 : SectionFileAddressSpace::const_iterator section_it =
269 E : --section_file_offsets_.end();
270 E : size_t file_size = section_it->second.value() + section_it->first.size();
271 E : DCHECK_EQ(0U, file_size % nt_headers_->OptionalHeader.FileAlignment);
272 :
273 : if (fseek(file, file_size - 1, SEEK_SET) != 0 ||
274 E : fwrite("\0", 1, 1, file) != 1) {
275 i : LOG(ERROR) << "Unable to round out file size.";
276 i : return false;
277 : }
278 E : }
279 :
280 E : return true;
281 E : }
282 :
283 : bool PEFileWriter::WriteOneBlock(AbsoluteAddress image_base,
284 : const BlockGraph::Block* block,
285 E : FILE* file) {
286 : // This function walks through the data referred by the input block, and
287 : // patches it to reflect the addresses and offsets of the blocks
288 : // referenced before writing the block's data to the file.
289 E : DCHECK(block != NULL);
290 E : DCHECK(file != NULL);
291 : // If the block has no data, there's nothing to write.
292 E : if (block->data() == NULL) {
293 : // A block with no data can't have references to anything else.
294 E : DCHECK(block->references().empty());
295 E : return true;
296 : }
297 :
298 E : RelativeAddress addr;
299 E : if (!image_layout_.blocks.GetAddressOf(block, &addr)) {
300 i : LOG(ERROR) << "All blocks must have an address.";
301 i : return false;
302 : }
303 :
304 : // Find the section that contains this block.
305 : SectionFileAddressSpace::RangeMap::const_iterator it(
306 : section_file_offsets_.FindContaining(
307 E : SectionFileAddressSpace::Range(addr, block->data_size())));
308 E : if (it == section_file_offsets_.ranges().end()) {
309 i : LOG(ERROR) << "Block with data outside defined sections at: " << addr;
310 i : return false;
311 : }
312 :
313 : // Calculate the offset from the start of the section to
314 : // the start of the block, and the block's file offset.
315 E : BlockGraph::Offset offs = addr - it->first.start();
316 E : DCHECK_GE(offs, 0);
317 E : FileOffsetAddress file_offs = it->second + offs;
318 :
319 : // Copy the block data.
320 E : std::vector<uint8> data(block->data_size());
321 E : std::copy(block->data(), block->data() + block->data_size(), data.begin());
322 :
323 : // Patch up all the references.
324 : BlockGraph::Block::ReferenceMap::const_iterator ref_it(
325 E : block->references().begin());
326 : BlockGraph::Block::ReferenceMap::const_iterator ref_end(
327 E : block->references().end());
328 E : for (; ref_it != ref_end; ++ref_it) {
329 E : BlockGraph::Offset start = ref_it->first;
330 E : BlockGraph::Reference ref = ref_it->second;
331 E : BlockGraph::Block* dst = ref.referenced();
332 :
333 E : RelativeAddress src_addr(addr + start);
334 E : RelativeAddress dst_addr;
335 E : if (!image_layout_.blocks.GetAddressOf(dst, &dst_addr)) {
336 i : LOG(ERROR) << "All blocks must have an address";
337 i : return false;
338 : }
339 E : dst_addr += ref.offset();
340 :
341 : // Compute the new value of the reference.
342 E : uint32 value = 0;
343 E : switch (ref.type()) {
344 : case BlockGraph::ABSOLUTE_REF: {
345 E : value = image_base.value() + dst_addr.value();
346 : #ifndef NDEBUG
347 E : DCHECK(value >= nt_headers_->OptionalHeader.ImageBase);
348 : DCHECK(value < nt_headers_->OptionalHeader.ImageBase +
349 E : nt_headers_->OptionalHeader.SizeOfImage);
350 : #endif
351 : }
352 E : break;
353 :
354 : case BlockGraph::PC_RELATIVE_REF:
355 E : value = dst_addr - (src_addr + ref.size());
356 E : break;
357 :
358 : case BlockGraph::RELATIVE_REF:
359 E : value = dst_addr.value();
360 E : break;
361 :
362 : case BlockGraph::FILE_OFFSET_REF: {
363 : // Find the section that contains the destination address.
364 : SectionFileAddressSpace::RangeMap::const_iterator it(
365 : section_file_offsets_.FindContaining(
366 E : SectionFileAddressSpace::Range(dst_addr, 1)));
367 E : DCHECK(it != section_file_offsets_.ranges().end());
368 :
369 E : value = it->second.value() + (dst_addr - it->first.start());
370 E : }
371 E : break;
372 :
373 : default:
374 i : LOG(ERROR) << "Impossible reference type";
375 i : return false;
376 : break;
377 : }
378 :
379 : // Now store the new value.
380 E : switch (ref.size()) {
381 : case sizeof(uint8):
382 E : if (!UpdateReference(start, static_cast<uint8>(value), &data))
383 i : return false;
384 E : break;
385 :
386 : case sizeof(uint16):
387 i : if (!UpdateReference(start, static_cast<uint16>(value), &data))
388 i : return false;
389 i : break;
390 :
391 : case sizeof(uint32):
392 E : if (!UpdateReference(start, static_cast<uint32>(value), &data))
393 i : return false;
394 E : break;
395 :
396 : default:
397 i : LOG(ERROR) << "Unsupported reference size.";
398 i : return false;
399 : }
400 E : }
401 :
402 E : if (fseek(file, file_offs.value(), SEEK_SET) != 0) {
403 i : LOG(ERROR) << "Unable to seek file";
404 i : return false;
405 : }
406 E : if (fwrite(&data[0], sizeof(data[0]), data.size(), file) != data.size()) {
407 i : LOG(ERROR) << "Unable to write block";
408 i : return false;
409 : }
410 E : return true;
411 E : }
412 :
413 : } // namespace pe
|