1    :  // Copyright 2011 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    :  #ifndef SYZYGY_PE_PE_FILE_WRITER_H_
  16    :  #define SYZYGY_PE_PE_FILE_WRITER_H_
  17    :  
  18    :  #include "base/files/file_path.h"
  19    :  #include "syzygy/block_graph/block_graph.h"
  20    :  #include "syzygy/core/address_space.h"
  21    :  #include "syzygy/pe/decomposer.h"
  22    :  #include "syzygy/pe/image_layout.h"
  23    :  #include "syzygy/pe/pe_file_parser.h"
  24    :  
  25  m :  namespace pe {
  26    :  
  27    :  // Given an address space and header information, writes a BlockGraph out
  28    :  // to a PE image file.
  29  m :  class PEFileWriter {
  30  m :   public:
  31  m :    typedef block_graph::BlockGraph BlockGraph;
  32  m :    typedef core::AbsoluteAddress AbsoluteAddress;
  33  m :    typedef core::FileOffsetAddress FileOffsetAddress;
  34  m :    typedef core::RelativeAddress RelativeAddress;
  35    :  
  36    :    // @param image_layout the image layout to write.
  37  m :    explicit PEFileWriter(const ImageLayout& image_layout);
  38    :  
  39    :    // Writes the image to path.
  40  m :    bool WriteImage(const base::FilePath& path);
  41    :  
  42    :    // Updates the checksum for the image @p path.
  43  m :    static bool UpdateFileChecksum(const base::FilePath& path);
  44    :  
  45  m :   protected:
  46    :    // Validates the DOS header and the NT headers in the image.
  47    :    // On success, sets the nt_headers_ pointer.
  48  m :    bool ValidateHeaders();
  49    :  
  50    :    // Validates that the section info is consistent and populates
  51    :    // section_file_range_map_ and section_index_space_.
  52  m :    bool CalculateSectionRanges();
  53    :  
  54    :    // Writes the entire image to the given file. Delegates to FlushSection and
  55    :    // WriteOneBlock.
  56  m :    bool WriteBlocks(FILE* file);
  57    :  
  58    :    // Closes off the writing of a section by adding any necessary padding to the
  59    :    // output buffer.
  60  m :    void FlushSection(size_t section_index,
  61  m :                      std::vector<uint8>* buffer);
  62    :  
  63    :    // Writes a single block to the buffer, first writing any necessary padding
  64    :    // (the content of which depends on the section type), followed by the
  65    :    // block data (containing finalized references).
  66  m :    bool WriteOneBlock(AbsoluteAddress image_base,
  67  m :                       size_t section_index,
  68  m :                       const BlockGraph::Block* block,
  69  m :                       std::vector<uint8>* buffer);
  70    :  
  71    :    // The file ranges of each section. This is populated by
  72    :    // CalculateSectionRanges and is a map from section index (as ordered in
  73    :    // the image layout) to section ranges on disk.
  74  m :    typedef core::AddressRange<core::FileOffsetAddress, size_t> FileRange;
  75  m :    typedef std::map<size_t, FileRange> SectionIndexFileRangeMap;
  76  m :    SectionIndexFileRangeMap section_file_range_map_;
  77    :  
  78    :    // This stores an address-space from RVAs to section indices and is populated
  79    :    // by CalculateSectionRanges. This can be used to map from a block's
  80    :    // address to the index of its section. This is needed for finalizing
  81    :    // references.
  82  m :    typedef core::AddressSpace<core::RelativeAddress, size_t, size_t>
  83  m :        SectionIndexSpace;
  84  m :    SectionIndexSpace section_index_space_;
  85    :  
  86    :    // Our image layout as provided to the constructor.
  87  m :    const ImageLayout& image_layout_;
  88    :  
  89    :    // Refers to the nt headers from the image during WriteImage.
  90  m :    const IMAGE_NT_HEADERS* nt_headers_;
  91    :  
  92  m :   private:
  93  m :    DISALLOW_COPY_AND_ASSIGN(PEFileWriter);
  94  m :  };
  95    :  
  96  m :  }  // namespace pe
  97    :  
  98    :  #endif  // SYZYGY_PE_PE_FILE_WRITER_H_