1    :  // Copyright 2013 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    :  #include "syzygy/pe/transforms/coff_prepare_headers_transform.h"
  16    :  
  17    :  #include "syzygy/block_graph/typed_block.h"
  18    :  #include "syzygy/pe/pe_utils.h"
  19    :  
  20    :  namespace pe {
  21    :  namespace transforms {
  22    :  
  23    :  using block_graph::BlockGraph;
  24    :  using block_graph::TypedBlock;
  25    :  
  26    :  const char CoffPrepareHeadersTransform::kTransformName[] =
  27    :      "CoffPrepareHeadersTransform";
  28    :  
  29    :  bool CoffPrepareHeadersTransform::TransformBlockGraph(
  30    :      const TransformPolicyInterface* policy,
  31    :      BlockGraph* block_graph,
  32  E :      BlockGraph::Block* headers_block) {
  33  E :    DCHECK_NE(reinterpret_cast<TransformPolicyInterface*>(NULL), policy);
  34  E :    DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
  35  E :    DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), headers_block);
  36  E :    DCHECK_EQ(BlockGraph::COFF_IMAGE, block_graph->image_format());
  37    :  
  38  E :    TypedBlock<IMAGE_FILE_HEADER> file_header;
  39  E :    if (!file_header.Init(0, headers_block)) {
  40  i :      LOG(ERROR) << "Unable to dereference COFF headers.";
  41  i :      return false;
  42    :    }
  43    :  
  44    :    // Wipe out references from headers to section blocks; these will be
  45    :    // rewritten during layout building.
  46  E :    if (!headers_block->RemoveAllReferences()) {
  47  i :      LOG(ERROR) << "Unable to remove references from COFF headers.";
  48  i :      return false;
  49    :    }
  50    :  
  51    :    // Resize the section table after the file header to reflect the number of
  52    :    // sections in the block graph. This ignores any optional header space, as
  53    :    // none should be included in the output COFF file.
  54    :    size_t new_headers_size = sizeof(IMAGE_FILE_HEADER) +
  55  E :        sizeof(IMAGE_SECTION_HEADER) * block_graph->sections().size();
  56  E :    size_t old_headers_size = file_header.block()->size();
  57    :    if (!headers_block->InsertOrRemoveData(
  58  E :            0, old_headers_size, new_headers_size, true)) {
  59  i :      LOG(ERROR) << "Unable to resize COFF headers.";
  60  i :      return false;
  61    :    }
  62    :  
  63  E :    file_header->NumberOfSections = block_graph->sections().size();
  64  E :    file_header->SizeOfOptionalHeader = 0;
  65    :  
  66  E :    return true;
  67  E :  }
  68    :  
  69    :  }  // namespace transforms
  70    :  }  // namespace pe