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