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_relinker.h"
16 :
17 : #include "base/file_util.h"
18 : #include "syzygy/block_graph/orderers/original_orderer.h"
19 : #include "syzygy/core/zstream.h"
20 : #include "syzygy/pdb/pdb_byte_stream.h"
21 : #include "syzygy/pdb/pdb_file.h"
22 : #include "syzygy/pdb/pdb_reader.h"
23 : #include "syzygy/pdb/pdb_util.h"
24 : #include "syzygy/pdb/pdb_writer.h"
25 : #include "syzygy/pe/decomposer.h"
26 : #include "syzygy/pe/find.h"
27 : #include "syzygy/pe/image_layout_builder.h"
28 : #include "syzygy/pe/image_source_map.h"
29 : #include "syzygy/pe/metadata.h"
30 : #include "syzygy/pe/pdb_info.h"
31 : #include "syzygy/pe/pe_file_writer.h"
32 : #include "syzygy/pe/pe_utils.h"
33 : #include "syzygy/pe/serialization.h"
34 : #include "syzygy/pe/orderers/pe_orderer.h"
35 : #include "syzygy/pe/transforms/add_metadata_transform.h"
36 : #include "syzygy/pe/transforms/add_pdb_info_transform.h"
37 : #include "syzygy/pe/transforms/prepare_headers_transform.h"
38 :
39 : namespace pe {
40 :
41 : namespace {
42 :
43 : typedef block_graph::BlockGraphTransformInterface Transform;
44 : typedef block_graph::BlockGraphOrdererInterface Orderer;
45 :
46 : using block_graph::ApplyBlockGraphTransform;
47 : using block_graph::BlockGraph;
48 : using block_graph::OrderedBlockGraph;
49 : using core::RelativeAddress;
50 : using pdb::NameStreamMap;
51 : using pdb::PdbByteStream;
52 : using pdb::PdbFile;
53 : using pdb::PdbInfoHeader70;
54 : using pdb::PdbMutatorInterface;
55 : using pdb::PdbStream;
56 : using pdb::WritablePdbStream;
57 :
58 : // A utility class for wrapping a serialization OutStream around a
59 : // WritablePdbStream.
60 : // TODO(chrisha): We really need to centralize stream/buffer semantics in
61 : // a small set of clean interfaces, and make all input/output/parsing work
62 : // on these interfaces.
63 : class PdbOutStream : public core::OutStream {
64 : public:
65 E : explicit PdbOutStream(WritablePdbStream* pdb_stream)
66 : : pdb_stream_(pdb_stream) {
67 E : DCHECK(pdb_stream != NULL);
68 E : }
69 :
70 E : virtual ~PdbOutStream() { }
71 :
72 E : virtual bool Write(size_t length, const core::Byte* bytes) OVERRIDE {
73 E : return pdb_stream_->Write(length, bytes);
74 E : }
75 :
76 : private:
77 : scoped_refptr<WritablePdbStream> pdb_stream_;
78 : };
79 :
80 : void GetOmapRange(const std::vector<ImageLayout::SectionInfo>& sections,
81 E : RelativeAddressRange* range) {
82 E : DCHECK(range != NULL);
83 :
84 : // There need to be at least two sections, one containing something and the
85 : // other containing the relocs.
86 E : DCHECK_GT(sections.size(), 1u);
87 E : DCHECK_EQ(sections.back().name, std::string(kRelocSectionName));
88 :
89 : // For some reason, if we output OMAP entries for the headers (before the
90 : // first section), everything falls apart. Not outputting these allows the
91 : // unittests to pass. Also, we don't want to output OMAP information for
92 : // the relocs, as these are entirely different from image to image.
93 E : RelativeAddress start_of_image = sections.front().addr;
94 E : RelativeAddress end_of_image = sections.back().addr;
95 E : *range = RelativeAddressRange(start_of_image, end_of_image - start_of_image);
96 E : }
97 :
98 : // TODO(chrisha): Make this utility function part of orderer.h.
99 : bool ApplyOrderer(Orderer* orderer,
100 : OrderedBlockGraph* obg,
101 E : BlockGraph::Block* header_block) {
102 E : DCHECK(orderer != NULL);
103 E : DCHECK(obg != NULL);
104 E : DCHECK(header_block != NULL);
105 :
106 E : if (!orderer->OrderBlockGraph(obg, header_block)) {
107 E : LOG(ERROR) << "Orderer failed: " << orderer->name();
108 E : return false;
109 : }
110 :
111 E : return true;
112 E : }
113 :
114 : bool ApplyPdbMutator(PdbMutatorInterface* pdb_mutator,
115 E : PdbFile* pdb_file) {
116 E : DCHECK(pdb_mutator != NULL);
117 E : DCHECK(pdb_file != NULL);
118 :
119 E : if (!pdb_mutator->MutatePdb(pdb_file)) {
120 E : LOG(ERROR) << "PDB mutator failed: " << pdb_mutator->name();
121 E : return false;
122 : }
123 :
124 E : return true;
125 E : }
126 :
127 : // Initializes input_pdb_path, output_pdb_path, pe_file and guid. If the input
128 : // paths are unable to be found this will return false. If @p allow_overwrite is
129 : // false and output path or output_pdb_path will overwrite an existing file this
130 : // will return false. @p input_pdb_path may be left empty in which case it will
131 : // be automatically determined from the debug information in @p input_path; this
132 : // step may fail causing this to return false. @p output_pdb_path may also be
133 : // left empty in which case it will be inferred from input_pdb_path, being
134 : // placed alongside output_path.
135 : bool InitializePaths(const FilePath& input_path,
136 : const FilePath& output_path,
137 : bool allow_overwrite,
138 : FilePath* input_pdb_path,
139 E : FilePath* output_pdb_path) {
140 E : DCHECK(input_pdb_path != NULL);
141 E : DCHECK(output_pdb_path != NULL);
142 :
143 : // At a very minimum we have to specify input and outputs.
144 E : if (input_path.empty() || output_path.empty()) {
145 E : LOG(ERROR) << "input_path and output_path must be set!";
146 E : return false;
147 : }
148 :
149 E : if (!file_util::PathExists(input_path)) {
150 E : LOG(ERROR) << "Input module not found: " << input_path.value();
151 E : return false;
152 : }
153 :
154 : // No input PDB specified? Find it automagically.
155 E : if (input_pdb_path->empty()) {
156 E : LOG(INFO) << "Input PDB not specified, searching for it.";
157 : if (!FindPdbForModule(input_path, input_pdb_path) ||
158 E : input_pdb_path->empty()) {
159 i : LOG(ERROR) << "Unable to find PDB file for module: "
160 : << input_path.value();
161 i : return false;
162 : }
163 : }
164 :
165 E : if (!file_util::PathExists(*input_pdb_path)) {
166 i : LOG(ERROR) << "Input PDB not found: " << input_pdb_path->value();
167 i : return false;
168 : }
169 :
170 : // If no output PDB path is specified, infer one.
171 E : if (output_pdb_path->empty()) {
172 : // If the input and output DLLs have the same basename, default to writing
173 : // using the same PDB basename, but alongside the new module.
174 E : if (input_path.BaseName() == output_path.BaseName()) {
175 : *output_pdb_path = output_path.DirName().Append(
176 E : input_pdb_path->BaseName());
177 E : } else {
178 : // Otherwise, default to using the output basename with a PDB extension.
179 i : *output_pdb_path = output_path.ReplaceExtension(L"pdb");
180 : }
181 :
182 E : LOG(INFO) << "Using default output PDB path: " << output_pdb_path->value();
183 : }
184 :
185 : // Ensure we aren't about to overwrite anything we don't want to. We do this
186 : // early on so that we abort before decomposition, transformation, etc.
187 E : if (!allow_overwrite) {
188 E : bool terminate = false;
189 E : if (file_util::PathExists(output_path)) {
190 E : terminate = true;
191 E : LOG(ERROR) << "Output module path already exists.";
192 : }
193 E : if (file_util::PathExists(*output_pdb_path)) {
194 i : terminate = true;
195 i : LOG(ERROR) << "Output PDB path already exists.";
196 : }
197 E : if (terminate)
198 E : return false;
199 : }
200 :
201 E : return true;
202 E : }
203 :
204 : // Decomposes the module enclosed by the given PE file.
205 : bool Decompose(const PEFile& pe_file,
206 : const FilePath& pdb_path,
207 : ImageLayout* image_layout,
208 E : BlockGraph::Block** dos_header_block) {
209 E : DCHECK(image_layout != NULL);
210 E : DCHECK(dos_header_block != NULL);
211 :
212 E : LOG(INFO) << "Decomposing module: " << pe_file.path().value();
213 :
214 : // Decompose the input image.
215 E : Decomposer decomposer(pe_file);
216 E : decomposer.set_pdb_path(pdb_path);
217 E : if (!decomposer.Decompose(image_layout)) {
218 i : LOG(ERROR) << "Unable to decompose module: " << pe_file.path().value();
219 i : return false;
220 : }
221 :
222 : // Get the DOS header block.
223 : *dos_header_block =
224 : image_layout->blocks.GetBlockByAddress(
225 E : BlockGraph::RelativeAddress(0));
226 E : if (*dos_header_block == NULL) {
227 i : LOG(ERROR) << "Unable to find the DOS header block.";
228 i : return false;
229 : }
230 :
231 E : return true;
232 E : }
233 :
234 : bool ApplyTransforms(const FilePath& input_path,
235 : const FilePath& output_pdb_path,
236 : const GUID& guid,
237 : bool add_metadata,
238 : std::vector<Transform*>* transforms,
239 : BlockGraph* block_graph,
240 E : BlockGraph::Block* dos_header_block) {
241 E : DCHECK(transforms != NULL);
242 E : DCHECK(block_graph != NULL);
243 E : DCHECK(dos_header_block != NULL);
244 :
245 E : LOG(INFO) << "Transforming block graph.";
246 :
247 E : std::vector<Transform*> local_transforms(*transforms);
248 :
249 E : pe::transforms::AddMetadataTransform add_metadata_tx(input_path);
250 E : pe::transforms::AddPdbInfoTransform add_pdb_info_tx(output_pdb_path, 1, guid);
251 E : pe::transforms::PrepareHeadersTransform prep_headers_tx;
252 :
253 E : if (add_metadata)
254 E : local_transforms.push_back(&add_metadata_tx);
255 E : local_transforms.push_back(&add_pdb_info_tx);
256 E : local_transforms.push_back(&prep_headers_tx);
257 :
258 : // Apply the transforms.
259 E : for (size_t i = 0; i < local_transforms.size(); ++i) {
260 E : LOG(INFO) << "Applying transform: " << local_transforms[i]->name() << ".";
261 : // ApplyBlockGraphTransform takes care of verbosely logging any failures.
262 : if (!ApplyBlockGraphTransform(local_transforms[i],
263 : block_graph,
264 E : dos_header_block)) {
265 E : return false;
266 : }
267 E : }
268 :
269 E : return true;
270 E : }
271 :
272 : bool ApplyOrderers(std::vector<Orderer*>* orderers,
273 : OrderedBlockGraph* obg,
274 E : BlockGraph::Block* dos_header_block) {
275 E : DCHECK(orderers != NULL);
276 E : DCHECK(obg != NULL);
277 E : DCHECK(dos_header_block != NULL);
278 :
279 E : LOG(INFO) << "Ordering block graph.";
280 :
281 E : std::vector<Orderer*> local_orderers(*orderers);
282 :
283 E : block_graph::orderers::OriginalOrderer orig_orderer;
284 E : pe::orderers::PEOrderer pe_orderer;
285 :
286 E : if (local_orderers.size() == 0) {
287 E : LOG(INFO) << "No orderers specified, using original orderer.";
288 E : local_orderers.push_back(&orig_orderer);
289 : }
290 :
291 E : local_orderers.push_back(&pe_orderer);
292 :
293 : // Apply the orderers.
294 E : for (size_t i = 0; i < local_orderers.size(); ++i) {
295 E : LOG(INFO) << "Applying orderer: " << local_orderers[i]->name();
296 E : if (!ApplyOrderer(local_orderers[i], obg, dos_header_block))
297 E : return false;
298 E : }
299 :
300 E : return true;
301 E : }
302 :
303 : bool ApplyPdbMutators(const std::vector<PdbMutatorInterface*>& pdb_mutators,
304 E : PdbFile* pdb_file) {
305 E : DCHECK(pdb_file != NULL);
306 :
307 E : LOG(INFO) << "Mutating PDB.";
308 :
309 : // Apply the orderers.
310 E : for (size_t i = 0; i < pdb_mutators.size(); ++i) {
311 E : LOG(INFO) << "Applying PDB mutator: " << pdb_mutators[i]->name();
312 E : if (!ApplyPdbMutator(pdb_mutators[i], pdb_file))
313 E : return false;
314 E : }
315 :
316 E : return true;
317 E : }
318 :
319 : // Lays out the image.
320 : bool BuildImageLayout(size_t padding,
321 : const OrderedBlockGraph& ordered_block_graph,
322 : BlockGraph::Block* dos_header_block,
323 E : ImageLayout* image_layout) {
324 E : DCHECK(dos_header_block != NULL);
325 E : DCHECK(image_layout != NULL);
326 :
327 E : LOG(INFO) << "Building image layout.";
328 :
329 E : ImageLayoutBuilder builder(image_layout);
330 E : builder.set_padding(padding);
331 E : if (!builder.LayoutImageHeaders(dos_header_block)) {
332 i : LOG(ERROR) << "ImageLayoutBuilder::LayoutImageHeaders failed.";
333 i : return false;
334 : }
335 :
336 E : if (!builder.LayoutOrderedBlockGraph(ordered_block_graph)) {
337 i : LOG(ERROR) << "ImageLayoutBuilder::LayoutOrderedBlockGraph failed.";
338 i : return false;
339 : }
340 :
341 E : LOG(INFO) << "Finalizing image layout.";
342 E : if (!builder.Finalize()) {
343 i : LOG(ERROR) << "ImageLayoutBuilder::Finalize failed.";
344 i : return false;
345 : }
346 :
347 E : return true;
348 E : }
349 :
350 : // Writes the image.
351 E : bool WriteImage(const ImageLayout& image_layout, const FilePath& output_path) {
352 E : PEFileWriter writer(image_layout);
353 :
354 E : LOG(INFO) << "Writing image: " << output_path.value();
355 E : if (!writer.WriteImage(output_path)) {
356 i : LOG(ERROR) << "Failed to write image \"" << output_path.value() << "\".";
357 i : return false;
358 : }
359 :
360 E : return true;
361 E : }
362 :
363 : void BuildOmapVectors(const RelativeAddressRange& input_range,
364 : const ImageLayout& output_image_layout,
365 : std::vector<OMAP>* omap_to,
366 E : std::vector<OMAP>* omap_from) {
367 E : DCHECK(omap_to != NULL);
368 E : DCHECK(omap_from != NULL);
369 :
370 E : LOG(INFO) << "Building OMAP vectors.";
371 :
372 : // Get the range of the output image, sans headers. This is required for
373 : // generating OMAP information.
374 E : RelativeAddressRange output_range;
375 E : GetOmapRange(output_image_layout.sections, &output_range);
376 :
377 E : ImageSourceMap reverse_map;
378 E : BuildImageSourceMap(output_image_layout, &reverse_map);
379 :
380 E : ImageSourceMap forward_map;
381 E : if (reverse_map.ComputeInverse(&forward_map) != 0) {
382 E : LOG(WARNING) << "OMAPFROM not unique (there exist repeated source ranges).";
383 : }
384 :
385 : // Build the two OMAP vectors.
386 E : BuildOmapVectorFromImageSourceMap(output_range, reverse_map, omap_to);
387 E : BuildOmapVectorFromImageSourceMap(input_range, forward_map, omap_from);
388 E : }
389 :
390 : // Updates the OMAP and GUID info in the given PDB file.
391 : bool SetOmapAndGuid(const RelativeAddressRange input_range,
392 : const ImageLayout& image_layout,
393 : const GUID& guid,
394 E : PdbFile* pdb_file) {
395 E : DCHECK(pdb_file != NULL);
396 :
397 E : LOG(INFO) << "Updating OMAP and GUID information.";
398 :
399 E : std::vector<OMAP> omap_to, omap_from;
400 E : BuildOmapVectors(input_range, image_layout, &omap_to, &omap_from);
401 :
402 E : if (!pdb::SetGuid(guid, pdb_file)) {
403 i : LOG(ERROR) << "Unable to set PDB GUID.";
404 i : return false;
405 : }
406 :
407 E : if (!pdb::SetOmapToStream(omap_to, pdb_file)) {
408 i : LOG(ERROR) << "Unable to set OMAP_TO.";
409 i : return false;
410 : }
411 :
412 E : if (!pdb::SetOmapFromStream(omap_from, pdb_file)) {
413 i : LOG(ERROR) << "Unable to set OMAP_FROM.";
414 i : return false;
415 : }
416 :
417 E : return true;
418 E : }
419 :
420 E : bool WritePdbFile(const FilePath& output_pdb_path, const PdbFile& pdb_file) {
421 E : LOG(INFO) << "Writing PDB file: " << output_pdb_path.value();
422 :
423 E : FilePath temp_pdb;
424 : if (!file_util::CreateTemporaryFileInDir(output_pdb_path.DirName(),
425 E : &temp_pdb)) {
426 i : LOG(ERROR) << "Unable to create temporary PDB file.";
427 i : return false;
428 : }
429 :
430 E : pdb::PdbWriter pdb_writer;
431 E : if (!pdb_writer.Write(temp_pdb, pdb_file)) {
432 i : LOG(ERROR) << "Failed to write temporary PDB file to \""
433 : << temp_pdb.value() << "\".";
434 : }
435 :
436 E : if (!file_util::ReplaceFile(temp_pdb, output_pdb_path)) {
437 i : LOG(ERROR) << "Unable to move temporary PDB file to \""
438 : << output_pdb_path.value() << "\".";
439 i : file_util::Delete(temp_pdb, false);
440 i : return false;
441 : }
442 :
443 E : return true;
444 E : }
445 :
446 : // Get a specific named stream if it already exists, otherwise create one.
447 : // @param stream_name The name of the stream.
448 : // @param name_stream_map The map containing the names of the streams in the
449 : // PDB. If the stream doesn't already exist the map will be augmented with
450 : // another entry.
451 : // @param pdb_file The PDB file to which the stream will be added.
452 : // @param replace_stream If true, will cause a new stream to be created even if
453 : // another one already existed.
454 : // @return a pointer to the PDB stream on success, NULL on failure.
455 : PdbStream* GetOrCreatePdbStreamByName(const char* stream_name,
456 : bool replace_stream,
457 : NameStreamMap* name_stream_map,
458 E : PdbFile* pdb_file) {
459 E : DCHECK(name_stream_map != NULL);
460 E : DCHECK(pdb_file != NULL);
461 E : scoped_refptr<PdbStream> stream;
462 :
463 E : NameStreamMap::const_iterator name_it = name_stream_map->find(stream_name);
464 E : if (name_it != name_stream_map->end()) {
465 : // Replace the existing stream by a brand-new one if it's required.
466 i : if (replace_stream) {
467 i : stream = new PdbByteStream();
468 i : pdb_file->ReplaceStream(name_it->second, stream.get());
469 i : } else {
470 i : if (!pdb::EnsureStreamWritable(name_it->second, pdb_file)) {
471 i : LOG(ERROR) << "Failed to make " << stream_name << " stream writable.";
472 i : return NULL;
473 : }
474 i : stream = pdb_file->GetStream(name_it->second);
475 : }
476 i : } else {
477 E : stream = new PdbByteStream();
478 E : uint32 index = pdb_file->AppendStream(stream.get());
479 E : (*name_stream_map)[stream_name] = index;
480 : }
481 :
482 E : return stream.get();
483 E : }
484 :
485 : // This updates or creates the Syzygy history stream, appending the metadata
486 : // describing this module and transform. The history stream consists of
487 : // a named PDB stream with the name /Syzygy/History. It consists of:
488 : //
489 : // uint32 version
490 : // uint32 history_length
491 : // serialized pe::Metadata 0
492 : // ...
493 : // serialized pe::Metadata history_length - 1
494 : //
495 : // If the format is changed, be sure to update this documentation and
496 : // pdb::kSyzygyHistoryStreamVersion (in pdb_constants.h).
497 : bool WriteSyzygyHistoryStream(const FilePath& input_path,
498 : NameStreamMap* name_stream_map,
499 E : PdbFile* pdb_file) {
500 : // Get the history stream.
501 : scoped_refptr<PdbStream> history_reader =
502 : GetOrCreatePdbStreamByName(pdb::kSyzygyHistoryStreamName,
503 : false,
504 : name_stream_map,
505 E : pdb_file);
506 :
507 E : if (history_reader == NULL) {
508 i : LOG(ERROR) << "Failed to get the history stream.";
509 i : return false;
510 : }
511 :
512 : scoped_refptr<WritablePdbStream> history_writer =
513 E : history_reader->GetWritablePdbStream();
514 E : DCHECK(history_writer.get() != NULL);
515 :
516 : // Get the metadata.
517 E : Metadata metadata;
518 E : PEFile pe_file;
519 E : if (!pe_file.Init(input_path)) {
520 i : LOG(ERROR) << "Failed to initialize PE file for \"" << input_path.value()
521 : << "\".";
522 i : return false;
523 : }
524 :
525 E : PEFile::Signature pe_sig;
526 E : pe_file.GetSignature(&pe_sig);
527 E : if (!metadata.Init(pe_sig)) {
528 i : LOG(ERROR) << "Failed to initialize metadata for \"" << input_path.value()
529 : << "\".";
530 i : return false;
531 : }
532 :
533 : // Validate the history stream if it is non-empty.
534 E : if (history_reader->length() > 0) {
535 : // Read the header.
536 i : uint32 version = 0;
537 i : uint32 history_length = 0;
538 : if (!history_reader->Seek(0) ||
539 : !history_reader->Read(&version, 1) ||
540 i : !history_reader->Read(&history_length, 1)) {
541 i : LOG(ERROR) << "Failed to read existing Syzygy history stream header.";
542 i : return false;
543 : }
544 :
545 : // Check the version.
546 i : if (version != pdb::kSyzygyHistoryStreamVersion) {
547 i : LOG(ERROR) << "PDB contains unsupported Syzygy history stream version "
548 : << "(got " << version << ", expected "
549 : << pdb::kSyzygyHistoryStreamVersion << ").";
550 i : return false;
551 : }
552 :
553 : // Increment the history length and rewrite it.
554 i : history_length++;
555 i : history_writer->set_pos(sizeof(pdb::kSyzygyHistoryStreamVersion));
556 i : if (!history_writer->Write(history_length)) {
557 i : LOG(ERROR) << "Failed to write new Syzygy history stream length.";
558 i : return false;
559 : }
560 i : } else {
561 : // If there wasn't already a history stream, create one and write the
562 : // header.
563 E : DCHECK_EQ(0u, history_writer->pos());
564 E : const uint32 kHistoryLength = 1;
565 : if (!history_writer->Write(pdb::kSyzygyHistoryStreamVersion) ||
566 E : !history_writer->Write(kHistoryLength)) {
567 i : LOG(ERROR) << "Failed to write Syzygy history stream header.";
568 i : return false;
569 : }
570 : }
571 :
572 : // Append the metadata to the history.
573 E : history_writer->set_pos(history_writer->length());
574 E : PdbOutStream out_stream(history_writer.get());
575 E : core::OutArchive out_archive(&out_stream);
576 E : if (!out_archive.Save(metadata)) {
577 i : LOG(ERROR) << "Failed to write metadata to Syzygy history stream.";
578 i : return false;
579 : }
580 :
581 E : return true;
582 E : }
583 :
584 : // This writes the serialized block-graph and the image layout in a PDB stream
585 : // named /Syzygy/BlockGraph. If the format is changed, be sure to update this
586 : // documentation and pdb::kSyzygyBlockGraphStreamVersion (in pdb_constants.h).
587 : // The block graph stream will not include the data from the blocks of the
588 : // block-graph. If the strip-strings flag is set to true the strings contained
589 : // in the block-graph won't be saved.
590 : bool WriteSyzygyBlockGraphStream(const PEFile& pe_file,
591 : const ImageLayout& image_layout,
592 : bool strip_strings,
593 : bool compress,
594 : NameStreamMap* name_stream_map,
595 E : PdbFile* pdb_file) {
596 : // Get the redecomposition data stream.
597 : scoped_refptr<PdbStream> block_graph_reader =
598 : GetOrCreatePdbStreamByName(pdb::kSyzygyBlockGraphStreamName,
599 : true,
600 : name_stream_map,
601 E : pdb_file);
602 :
603 E : if (block_graph_reader == NULL) {
604 i : LOG(ERROR) << "Failed to get the block-graph stream.";
605 i : return false;
606 : }
607 E : DCHECK_EQ(0u, block_graph_reader->length());
608 :
609 : scoped_refptr<WritablePdbStream> block_graph_writer =
610 E : block_graph_reader->GetWritablePdbStream();
611 E : DCHECK(block_graph_writer.get() != NULL);
612 :
613 : // Write the version of the BlockGraph stream, and whether or not its
614 : // contents are compressed.
615 : if (!block_graph_writer->Write(pdb::kSyzygyBlockGraphStreamVersion) ||
616 E : !block_graph_writer->Write(static_cast<unsigned char>(compress))) {
617 i : LOG(ERROR) << "Failed to write Syzygy BlockGraph stream header.";
618 i : return false;
619 : }
620 :
621 : // Set up the output stream.
622 E : PdbOutStream pdb_out_stream(block_graph_writer.get());
623 E : core::OutStream* out_stream = &pdb_out_stream;
624 :
625 : // If requested, compress the output.
626 E : scoped_ptr<core::ZOutStream> zip_stream;
627 E : if (compress) {
628 E : zip_stream.reset(new core::ZOutStream(&pdb_out_stream));
629 E : out_stream = zip_stream.get();
630 E : if (!zip_stream->Init(core::ZOutStream::kZBestCompression)) {
631 i : LOG(ERROR) << "Failed to initialize zlib compressor.";
632 i : return false;
633 : }
634 : }
635 :
636 E : core::OutArchive out_archive(out_stream);
637 :
638 : // Set up the serialization properties.
639 E : block_graph::BlockGraphSerializer::Attributes attributes = 0;
640 E : if (strip_strings)
641 E : attributes |= block_graph::BlockGraphSerializer::OMIT_STRINGS;
642 :
643 : // And finally, perform the serialization.
644 : if (!SaveBlockGraphAndImageLayout(pe_file, attributes, image_layout,
645 E : &out_archive)) {
646 i : LOG(ERROR) << "SaveBlockGraphAndImageLayout failed.";
647 i : return false;
648 : }
649 :
650 : // We have to flush the stream in case it's a zstream.
651 E : out_stream->Flush();
652 :
653 E : return true;
654 E : }
655 :
656 : } // namespace
657 :
658 : PERelinker::PERelinker()
659 : : add_metadata_(true), allow_overwrite_(false), augment_pdb_(true),
660 : strip_strings_(false), padding_(0), inited_(false),
661 : input_image_layout_(&block_graph_), dos_header_block_(NULL),
662 E : output_guid_(GUID_NULL) {
663 E : }
664 :
665 E : void PERelinker::AppendTransform(Transform* transform) {
666 E : DCHECK(transform != NULL);
667 E : transforms_.push_back(transform);
668 E : }
669 :
670 E : void PERelinker::AppendTransforms(const std::vector<Transform*>& transforms) {
671 E : transforms_.insert(transforms_.end(), transforms.begin(), transforms.end());
672 E : }
673 :
674 E : void PERelinker::AppendOrderer(Orderer* orderer) {
675 E : DCHECK(orderer != NULL);
676 E : orderers_.push_back(orderer);
677 E : }
678 :
679 E : void PERelinker::AppendOrderers(const std::vector<Orderer*>& orderers) {
680 E : orderers_.insert(orderers_.end(), orderers.begin(), orderers.end());
681 E : }
682 :
683 E : void PERelinker::AppendPdbMutator(PdbMutatorInterface* pdb_mutator) {
684 E : DCHECK(pdb_mutator != NULL);
685 E : pdb_mutators_.push_back(pdb_mutator);
686 E : }
687 :
688 : void PERelinker::AppendPdbMutators(
689 E : const std::vector<PdbMutatorInterface*>& pdb_mutators) {
690 : pdb_mutators_.insert(pdb_mutators_.end(),
691 : pdb_mutators.begin(),
692 E : pdb_mutators.end());
693 E : }
694 :
695 E : bool PERelinker::Init() {
696 E : DCHECK(inited_ == false);
697 :
698 : // Initialize the paths.
699 : if (!InitializePaths(input_path_, output_path_, allow_overwrite_,
700 E : &input_pdb_path_, &output_pdb_path_)) {
701 E : return false;
702 : }
703 :
704 E : LOG(INFO) << "Input module : " << input_path_.value();
705 E : LOG(INFO) << "Input PDB : " << input_pdb_path_.value();
706 E : LOG(INFO) << "Output module: " << output_path_.value();
707 E : LOG(INFO) << "Output PDB : " << output_pdb_path_.value();
708 :
709 : // Open the input PE file.
710 E : if (!input_pe_file_.Init(input_path_)) {
711 i : LOG(ERROR) << "Unable to load \"" << input_path_.value() << "\".";
712 i : return false;
713 : }
714 :
715 : // Generate a GUID for the relinked image's PDB file.
716 E : if (FAILED(::CoCreateGuid(&output_guid_))) {
717 i : LOG(ERROR) << "Failed to create new PDB GUID.";
718 i : return false;
719 : }
720 :
721 : // Decompose the image.
722 : if (!Decompose(input_pe_file_, input_pdb_path_, &input_image_layout_,
723 E : &dos_header_block_)) {
724 i : return false;
725 : }
726 :
727 E : inited_ = true;
728 :
729 E : return true;
730 E : }
731 :
732 E : bool PERelinker::Relink() {
733 E : if (!inited_) {
734 i : LOG(ERROR) << "Init has not been successfully called.";
735 i : return false;
736 : }
737 :
738 : // Transform it.
739 : if (!ApplyTransforms(input_path_, output_pdb_path_, output_guid_,
740 : add_metadata_, &transforms_, &block_graph_,
741 E : dos_header_block_)) {
742 E : return false;
743 : }
744 :
745 : // Order it.
746 E : OrderedBlockGraph ordered_block_graph(&block_graph_);
747 E : if (!ApplyOrderers(&orderers_, &ordered_block_graph, dos_header_block_))
748 E : return false;
749 :
750 : // Lay it out.
751 E : ImageLayout output_image_layout(&block_graph_);
752 : if (!BuildImageLayout(padding_, ordered_block_graph, dos_header_block_,
753 E : &output_image_layout)) {
754 i : return false;
755 : }
756 :
757 : // Write the image.
758 E : if (!WriteImage(output_image_layout, output_path_))
759 i : return false;
760 :
761 : // From here on down we are processing the PDB file.
762 :
763 : // Read the PDB file.
764 E : LOG(INFO) << "Reading PDB file: " << input_pdb_path_.value();
765 E : pdb::PdbReader pdb_reader;
766 E : PdbFile pdb_file;
767 E : if (!pdb_reader.Read(input_pdb_path_, &pdb_file)) {
768 i : LOG(ERROR) << "Unable to read PDB file: " << input_pdb_path_.value();
769 i : return false;
770 : }
771 :
772 : // Apply the mutators to the PDB file.
773 E : if (!ApplyPdbMutators(pdb_mutators_, &pdb_file))
774 E : return false;
775 :
776 : // TODO(chrisha): Make the following 3 PDB updates PdbMutatorInterface
777 : // implementations.
778 :
779 : // Update the OMAP and GUID information.
780 E : RelativeAddressRange input_range;
781 E : GetOmapRange(input_image_layout_.sections, &input_range);
782 : if (!SetOmapAndGuid(input_range, output_image_layout, output_guid_,
783 E : &pdb_file)) {
784 i : return false;
785 : }
786 :
787 : // Parse the header and named streams.
788 E : pdb::PdbInfoHeader70 header = {};
789 E : pdb::NameStreamMap name_stream_map;
790 E : if (!pdb::ReadHeaderInfoStream(pdb_file, &header, &name_stream_map))
791 i : return false;
792 :
793 : // Update/create the Syzygy history stream.
794 E : if (!WriteSyzygyHistoryStream(input_path_, &name_stream_map, &pdb_file))
795 i : return false;
796 :
797 : // Add redecomposition data in another stream, only if augment_pdb_ is set.
798 E : if (augment_pdb_) {
799 E : LOG(INFO) << "The block-graph stream is being written to the PDB.";
800 :
801 E : PEFile new_pe_file;
802 E : if (!new_pe_file.Init(output_path_)) {
803 i : LOG(ERROR) << "Failed to read newly written PE file.";
804 i : return false;
805 : }
806 :
807 : if (!WriteSyzygyBlockGraphStream(new_pe_file,
808 : output_image_layout,
809 : strip_strings_,
810 : compress_pdb_,
811 : &name_stream_map,
812 E : &pdb_file)) {
813 i : return false;
814 : }
815 E : }
816 :
817 : // Write the updated name-stream map back to the header info stream.
818 E : if (!pdb::WriteHeaderInfoStream(header, name_stream_map, &pdb_file))
819 i : return false;
820 :
821 : // Stream 0 contains a copy of the previous PDB's directory. This, combined
822 : // with copy-on-write semantics of individual blocks makes the file contain
823 : // its whole edit history. Since we're writing a 'new' PDB file (we reset the
824 : // GUID and age), we have no history so can safely throw away this stream.
825 E : pdb_file.ReplaceStream(0, NULL);
826 :
827 : // Write the PDB file. We use a helper function that first writes it to a
828 : // temporary file and then moves it, enabling overwrites.
829 E : if (!WritePdbFile(output_pdb_path_, pdb_file))
830 i : return false;
831 :
832 E : return true;
833 E : }
834 :
835 : } // namespace pe
|