1    :  // Copyright 2015 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/hot_patching_unittest_util.h"
  16    :  
  17    :  #include "syzygy/pe/transforms/pe_hot_patching_basic_block_transform.h"
  18    :  
  19    :  namespace testing {
  20    :  
  21    :  const char TestHotPatchingTransform::kTransformName[] =
  22    :      "TestHotPatchingTransform";
  23    :  
  24  E :  TestHotPatchingTransform::TestHotPatchingTransform() { }
  25    :  
  26  E :  TestHotPatchingTransform::~TestHotPatchingTransform() { }
  27    :  
  28    :  bool TestHotPatchingTransform::OnBlock(const TransformPolicyInterface* policy,
  29    :                                         BlockGraph* block_graph,
  30  E :                                         BlockGraph::Block* block) {
  31  E :    DCHECK_NE(static_cast<TransformPolicyInterface*>(nullptr), policy);
  32  E :    DCHECK_NE(static_cast<BlockGraph*>(nullptr), block_graph);
  33  E :    DCHECK_NE(static_cast<BlockGraph::Block*>(nullptr), block);
  34    :  
  35  E :    if (!policy->BlockIsSafeToBasicBlockDecompose(block))
  36  E :      return true;
  37    :  
  38    :    // Apply the basic block transform to make the block hot patchable.
  39  E :    pe::transforms::PEHotPatchingBasicBlockTransform transform;
  40  E :    block_graph::BlockVector new_blocks;
  41    :    if (!ApplyBasicBlockSubGraphTransform(
  42  E :            &transform, policy, block_graph, block, &new_blocks)) {
  43  i :      return false;
  44    :    }
  45    :  
  46    :    // One new code block should be created.
  47  E :    DCHECK_EQ(1U, new_blocks.size());
  48    :  
  49    :    // Collect transformed blocks.
  50  E :    blocks_prepared_.push_back(new_blocks[0]);
  51    :  
  52  E :    return true;
  53  E :  }
  54    :  
  55    :  bool TestHotPatchingTransform::PostBlockGraphIteration(
  56    :      const TransformPolicyInterface* policy,
  57    :      BlockGraph* block_graph,
  58  E :      BlockGraph::Block* header_block) {
  59  E :    DCHECK_NE(static_cast<TransformPolicyInterface*>(nullptr), policy);
  60  E :    DCHECK_NE(static_cast<BlockGraph*>(nullptr), block_graph);
  61  E :    DCHECK_NE(static_cast<BlockGraph::Block*>(nullptr), header_block);
  62    :  
  63    :    // Insert the hot patching metadata section.
  64  E :    pe::transforms::AddHotPatchingMetadataTransform hp_metadata_transform;
  65  E :    hp_metadata_transform.set_blocks_prepared(&blocks_prepared_);
  66    :    if (!ApplyBlockGraphTransform(&hp_metadata_transform, policy,
  67  E :                                  block_graph, header_block)) {
  68  i :      return false;
  69    :    }
  70    :  
  71  E :    return true;
  72  E :  }
  73    :  
  74    :  HotPatchingTestDllTest::HotPatchingTestDllTest()
  75    :      : relinker_(&policy_),
  76  E :        test_dll_path_(testing::GetExeRelativePath(testing::kTestDllName)) {
  77  E :  }
  78    :  
  79  E :  HotPatchingTestDllTest::~HotPatchingTestDllTest() { }
  80    :  
  81  E :  void HotPatchingTestDllTest::HotPatchInstrumentTestDll() {
  82    :    // Set up relinker.
  83  E :    relinker_.set_input_path(test_dll_path_);
  84  E :    relinker_.set_output_path(hp_test_dll_path_);
  85  E :    relinker_.set_allow_overwrite(true);
  86  E :    ASSERT_TRUE(relinker_.Init());
  87    :  
  88    :    // Make test.dll hot patchable.
  89  E :    relinker_.AppendTransform(&hp_transform_);
  90    :  
  91    :    // Perform the actual relink.
  92  E :    ASSERT_TRUE(relinker_.Relink());
  93    :  
  94    :    // Validate that the binary still loads.
  95  E :    ASSERT_NO_FATAL_FAILURE(CheckTestDll(hp_test_dll_path_));
  96  E :  }
  97    :  
  98  E :  void HotPatchingTestDllTest::SetUp() {
  99  E :    this->CreateTemporaryDir(&temp_dir_);
 100  E :    hp_test_dll_path_ = temp_dir_.Append(testing::kTestDllName);
 101  E :  }
 102    :  
 103    :  }  // namespace testing