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