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    :  // Declares an image layout builder, a utility class for constructing valid
  16    :  // PE ImageLayout objects.
  17    :  
  18    :  #ifndef SYZYGY_PE_PE_IMAGE_LAYOUT_BUILDER_H_
  19    :  #define SYZYGY_PE_PE_IMAGE_LAYOUT_BUILDER_H_
  20    :  
  21    :  #include <windows.h>
  22    :  #include <winnt.h>
  23    :  #include <vector>
  24    :  
  25    :  #include "syzygy/block_graph/block_graph.h"
  26    :  #include "syzygy/block_graph/ordered_block_graph.h"
  27    :  #include "syzygy/pe/image_layout.h"
  28    :  #include "syzygy/pe/pe_coff_image_layout_builder.h"
  29    :  #include "syzygy/pe/pe_file_parser.h"
  30    :  
  31    :  namespace pe {
  32    :  
  33    :  // A helper class that assists in assigning address space to PE image sections,
  34    :  // building self-consistent PE image headers etc.
  35    :  class PEImageLayoutBuilder : public PECoffImageLayoutBuilder {
  36    :   public:
  37    :    // Constructs a new image layout builder that populates the provided image
  38    :    // layout. The image layout must outlive the builder.
  39    :    explicit PEImageLayoutBuilder(ImageLayout* image_layout);
  40    :  
  41    :    // Accessors.
  42  E :    BlockGraph::Block* dos_header_block() { return dos_header_block_; }
  43  E :    BlockGraph::Block* nt_headers_block() { return nt_headers_block_; }
  44    :    const BlockGraph::Block* dos_header_block() const {
  45    :      return dos_header_block_;
  46    :    }
  47    :    const BlockGraph::Block* nt_headers_block() const {
  48    :      return nt_headers_block_;
  49    :    }
  50    :  
  51    :    // Lays out the image headers, and sets the file and section alignment using
  52    :    // the values from the header.
  53    :    // @param dos_header_block must be a block that's a valid DOS header
  54    :    //    and stub. This block must also refer to the NT headers block,
  55    :    //    which in turn must contain valid NT headers.
  56    :    // @returns true iff the dos_header_block is valid.
  57    :    // @pre OpenSection and LayoutBlock must not have been called.
  58    :    bool LayoutImageHeaders(BlockGraph::Block* dos_header_block);
  59    :  
  60    :    // Creates sections and lays out blocks using the provided ordered block
  61    :    // graph as a template. Lays out all sections except for the reloc section,
  62    :    // which must be the last section if it is present.
  63    :    // @param obg the ordered block graph to layout, which must be for the same
  64    :    //     block-graph as used in the constructor.
  65    :    // @returns true on success, false otherwise.
  66    :    // @pre LayoutImageHeaders has been called.
  67    :    bool LayoutOrderedBlockGraph(const OrderedBlockGraph& obg);
  68    :  
  69    :    // Finalizes the image layout. This builds the relocs, finalizes the headers,
  70    :    // and does any other PE touch-ups that are required to make the image
  71    :    // self-consistent. This may remove and/or modify blocks in the block-graph.
  72    :    // @returns true on success, false otherwise.
  73    :    bool Finalize();
  74    :  
  75    :   private:
  76    :    // Ensure that the Safe SEH Table is sorted.
  77    :    bool SortSafeSehTable();
  78    :    // Allocates and populates a new relocations section containing
  79    :    // relocations for all absolute references in address_space_.
  80    :    bool CreateRelocsSection();
  81    :    // Write the NT headers and section headers to the image.
  82    :    // After this is done, the image is "baked", and everything except for
  83    :    // the image checksum should be up to date.
  84    :    bool FinalizeHeaders();
  85    :    // Ensure that the image layout has the same number of blocks as the
  86    :    // block-graph. The relocs blocks that are in the block-graph but not in the
  87    :    // image layout will be removed. If there are extra blocks from other sections
  88    :    // in the block-graph an error will be returned.
  89    :    // @returns true if the block-graph and the image layout are consistent,
  90    :    //     false otherwise.
  91    :    bool ReconcileBlockGraphAndImageLayout();
  92    :  
  93    :    // The blocks that describe the DOS header and the NT headers.
  94    :    BlockGraph::Block* dos_header_block_;
  95    :    BlockGraph::Block* nt_headers_block_;
  96    :  
  97    :    DISALLOW_COPY_AND_ASSIGN(PEImageLayoutBuilder);
  98    :  };
  99    :  
 100    :  }  // namespace pe
 101    :  
 102    :  #endif  // SYZYGY_PE_PE_IMAGE_LAYOUT_BUILDER_H_