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(policy != NULL);
  34  E :    DCHECK(block_graph != NULL);
  35  E :    DCHECK(headers_block != NULL);
  36    :  
  37  E :    TypedBlock<IMAGE_FILE_HEADER> file_header;
  38  E :    if (!file_header.Init(0, headers_block)) {
  39  i :      LOG(ERROR) << "Unable to dereference COFF headers.";
  40  i :      return false;
  41    :    }
  42    :  
  43    :    // Wipe out references from headers to section blocks; these will be
  44    :    // rewritten during layout building.
  45  E :    if (!headers_block->RemoveAllReferences()) {
  46  i :      LOG(ERROR) << "Unable to remove references from COFF headers.";
  47  i :      return false;
  48    :    }
  49    :  
  50    :    // Resize the section table after the file header to reflect the number of
  51    :    // sections in the block graph. This ignores any optional header space, as
  52    :    // none should be included in the output COFF file.
  53    :    size_t new_headers_size = sizeof(IMAGE_FILE_HEADER) +
  54  E :        sizeof(IMAGE_SECTION_HEADER) * block_graph->sections().size();
  55  E :    size_t old_headers_size = file_header.block()->size();
  56    :    if (!headers_block->InsertOrRemoveData(
  57  E :            0, old_headers_size, new_headers_size, true)) {
  58  i :      LOG(ERROR) << "Unable to resize COFF headers.";
  59  i :      return false;
  60    :    }
  61    :  
  62  E :    file_header->NumberOfSections = block_graph->sections().size();
  63  E :    file_header->SizeOfOptionalHeader = 0;
  64    :  
  65  E :    return true;
  66  E :  }
  67    :  
  68    :  }  // namespace transforms
  69    :  }  // namespace pe