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    :  #include "syzygy/pe/transforms/pe_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::TypedBlock;
  24    :  
  25    :  typedef TypedBlock<IMAGE_DOS_HEADER> DosHeader;
  26    :  typedef TypedBlock<IMAGE_NT_HEADERS> NtHeaders;
  27    :  
  28    :  const char PEPrepareHeadersTransform::kTransformName[] =
  29    :      "PEPrepareHeadersTransform";
  30    :  
  31    :  bool PEPrepareHeadersTransform::TransformBlockGraph(
  32    :      const TransformPolicyInterface* policy,
  33    :      BlockGraph* block_graph,
  34  E :      BlockGraph::Block* dos_header_block) {
  35  E :    DCHECK_NE(reinterpret_cast<TransformPolicyInterface*>(NULL), policy);
  36  E :    DCHECK_NE(reinterpret_cast<BlockGraph*>(NULL), block_graph);
  37  E :    DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), dos_header_block);
  38  E :    DCHECK_EQ(BlockGraph::PE_IMAGE, block_graph->image_format());
  39    :  
  40  E :    DosHeader dos_header;
  41  E :    NtHeaders nt_headers;
  42    :    if (!dos_header.Init(0, dos_header_block) ||
  43  E :        !dos_header.Dereference(dos_header->e_lfanew, &nt_headers)) {
  44  i :      LOG(ERROR) << "Unable to dereference headers.";
  45  i :      return false;
  46    :    }
  47    :  
  48  E :    if (!UpdateDosHeader(dos_header_block)) {
  49  i :      LOG(ERROR) << "Unable to update DOS header.";
  50  i :      return false;
  51    :    }
  52    :  
  53    :    // Resize the NT headers to reflect the number of sections in the block graph.
  54    :    size_t new_nt_headers_size = sizeof(IMAGE_NT_HEADERS) +
  55  E :        sizeof(IMAGE_SECTION_HEADER) * block_graph->sections().size();
  56  E :    size_t old_nt_headers_size = nt_headers.block()->size();
  57    :    if (!nt_headers.block()->InsertOrRemoveData(
  58  E :            0, old_nt_headers_size, new_nt_headers_size, true)) {
  59  i :      LOG(ERROR) << "Unable to resize NT headers.";
  60  i :      return false;
  61    :    }
  62    :  
  63  E :    nt_headers->FileHeader.NumberOfSections = block_graph->sections().size();
  64  E :    nt_headers->OptionalHeader.CheckSum = 0;
  65    :    nt_headers->OptionalHeader.SizeOfHeaders =
  66    :        common::AlignUp(dos_header_block->size() + nt_headers.block()->size(),
  67  E :                        nt_headers->OptionalHeader.FileAlignment);
  68    :  
  69  E :    return true;
  70  E :  }
  71    :  
  72    :  }  // namespace transforms
  73    :  }  // namespace pe