// Copyright 2012 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Declares an image layout builder, a utility class for constructing valid
// PE ImageLayout objects.

#ifndef SYZYGY_PE_PE_IMAGE_LAYOUT_BUILDER_H_
#define SYZYGY_PE_PE_IMAGE_LAYOUT_BUILDER_H_

#include <windows.h>
#include <winnt.h>
#include <vector>

#include "syzygy/block_graph/block_graph.h"
#include "syzygy/block_graph/ordered_block_graph.h"
#include "syzygy/pe/image_layout.h"
#include "syzygy/pe/pe_coff_image_layout_builder.h"
#include "syzygy/pe/pe_file_parser.h"

namespace pe {

// A helper class that assists in assigning address space to PE image sections,
// building self-consistent PE image headers etc.
class PEImageLayoutBuilder : public PECoffImageLayoutBuilder {
 public:
  // Constructs a new image layout builder that populates the provided image
  // layout. The image layout must outlive the builder.
  explicit PEImageLayoutBuilder(ImageLayout* image_layout);

  // Accessors.
  BlockGraph::Block* dos_header_block() { return dos_header_block_; }
  BlockGraph::Block* nt_headers_block() { return nt_headers_block_; }
  const BlockGraph::Block* dos_header_block() const {
    return dos_header_block_;
  }
  const BlockGraph::Block* nt_headers_block() const {
    return nt_headers_block_;
  }

  // Lays out the image headers, and sets the file and section alignment using
  // the values from the header.
  // @param dos_header_block must be a block that's a valid DOS header
  //    and stub. This block must also refer to the NT headers block,
  //    which in turn must contain valid NT headers.
  // @returns true iff the dos_header_block is valid.
  // @pre OpenSection and LayoutBlock must not have been called.
  bool LayoutImageHeaders(BlockGraph::Block* dos_header_block);

  // Creates sections and lays out blocks using the provided ordered block
  // graph as a template. Lays out all sections except for the reloc section,
  // which must be the last section if it is present.
  // @param obg the ordered block graph to layout, which must be for the same
  //     block-graph as used in the constructor.
  // @returns true on success, false otherwise.
  // @pre LayoutImageHeaders has been called.
  bool LayoutOrderedBlockGraph(const OrderedBlockGraph& obg);

  // Finalizes the image layout. This builds the relocs, finalizes the headers,
  // and does any other PE touch-ups that are required to make the image
  // self-consistent. This may remove and/or modify blocks in the block-graph.
  // @returns true on success, false otherwise.
  bool Finalize();

 private:
  // Ensure that the Safe SEH Table is sorted.
  bool SortSafeSehTable();
  // Allocates and populates a new relocations section containing
  // relocations for all absolute references in address_space_.
  bool CreateRelocsSection();
  // Write the NT headers and section headers to the image.
  // After this is done, the image is "baked", and everything except for
  // the image checksum should be up to date.
  bool FinalizeHeaders();
  // Ensure that the image layout has the same number of blocks as the
  // block-graph. The relocs blocks that are in the block-graph but not in the
  // image layout will be removed. If there are extra blocks from other sections
  // in the block-graph an error will be returned.
  // @returns true if the block-graph and the image layout are consistent,
  //     false otherwise.
  bool ReconcileBlockGraphAndImageLayout();

  // The blocks that describe the DOS header and the NT headers.
  BlockGraph::Block* dos_header_block_;
  BlockGraph::Block* nt_headers_block_;

  DISALLOW_COPY_AND_ASSIGN(PEImageLayoutBuilder);
};

}  // namespace pe

#endif  // SYZYGY_PE_PE_IMAGE_LAYOUT_BUILDER_H_