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    :  #ifndef SYZYGY_PE_HOT_PATCHING_UNITTEST_UTIL_H_
  16    :  #define SYZYGY_PE_HOT_PATCHING_UNITTEST_UTIL_H_
  17    :  
  18    :  #include "base/files/file_path.h"
  19    :  #include "gtest/gtest.h"
  20    :  #include "syzygy/block_graph/transforms/iterative_transform.h"
  21    :  #include "syzygy/pe/pe_relinker.h"
  22    :  #include "syzygy/pe/unittest_util.h"
  23    :  #include "syzygy/pe/transforms/add_hot_patching_metadata_transform.h"
  24    :  
  25    :  namespace testing {
  26    :  
  27    :  // Prepares every code block of a module for hot patching.
  28    :  class TestHotPatchingTransform
  29    :      : public block_graph::transforms::IterativeTransformImpl<
  30    :            TestHotPatchingTransform> {
  31    :   public:
  32    :    typedef block_graph::BlockGraph BlockGraph;
  33    :    typedef block_graph::TransformPolicyInterface TransformPolicyInterface;
  34    :    typedef pe::transforms::AddHotPatchingMetadataTransform::BlockVector
  35    :        BlockVector;
  36    :  
  37    :    // The transform name.
  38    :    static const char kTransformName[];
  39    :  
  40    :    TestHotPatchingTransform();
  41    :    ~TestHotPatchingTransform();
  42    :  
  43    :    // After the transform has run, this function returns the blocks that have
  44    :    // been prepared for hot patching.
  45  E :    const BlockVector& blocks_prepared() const {
  46  E :      return blocks_prepared_;
  47  E :    }
  48    :  
  49    :    // @name IterativeTransformImpl implementation.
  50    :    // @{
  51    :    // Prepare every safe-to-decompose block for hot patching.
  52    :    // @param policy The policy object restricting how the transform is applied.
  53    :    // @param block_graph The block graph being modified.
  54    :    // @param block The block to explode, this must be in @p block_graph.
  55    :    // @returns true.
  56    :    bool OnBlock(const TransformPolicyInterface* policy,
  57    :                 BlockGraph* block_graph,
  58    :                 BlockGraph::Block* block);
  59    :  
  60    :    // Add the metadata stream to the BlockGraph.
  61    :    // @param policy The policy object restricting how the transform is applied.
  62    :    // @param block_graph the block graph being transformed.
  63    :    // @param header_block the header block.
  64    :    // @returns true on success, false otherwise.
  65    :    bool PostBlockGraphIteration(const TransformPolicyInterface* policy,
  66    :                                 BlockGraph* block_graph,
  67    :                                 BlockGraph::Block* header_block);
  68    :    // @}
  69    :  
  70    :   private:
  71    :    // Store the blocks that have been prepared for hot patching. This is used
  72    :    // to generate the metadata.
  73    :    BlockVector blocks_prepared_;
  74    :  
  75    :    DISALLOW_COPY_AND_ASSIGN(TestHotPatchingTransform);
  76    :  };
  77    :  
  78    :  // An unit test fixture that relinks the test_dll.dll with hot patching
  79    :  // information.
  80    :  class HotPatchingTestDllTest : public testing::PELibUnitTest {
  81    :   public:
  82    :    HotPatchingTestDllTest();
  83    :    ~HotPatchingTestDllTest();
  84    :  
  85    :    // Relinks test_dll.dll using TestHotPatchingTransform that prepares the
  86    :    // blocks for hot patching and adds hot patching metadata.
  87    :    void HotPatchInstrumentTestDll();
  88    :  
  89    :    // Creates a temporary directory for the transformed DLL.
  90    :    virtual void SetUp() override;
  91    :  
  92    :   protected:
  93    :    pe::PETransformPolicy policy_;
  94    :    pe::PERelinker relinker_;
  95    :  
  96    :    // Path of the original test_dll.dll.
  97    :    base::FilePath test_dll_path_;
  98    :    // Path of the temporary directory where the hot patchable DLL will be saved.
  99    :    base::FilePath temp_dir_;
 100    :    // Path of the hot patchable test_dll.dll.
 101    :    base::FilePath hp_test_dll_path_;
 102    :  
 103    :    // The transform used to make test_dll.dll hot patchable.
 104    :    TestHotPatchingTransform hp_transform_;
 105    :  
 106    :    DISALLOW_COPY_AND_ASSIGN(HotPatchingTestDllTest);
 107    :  };
 108    :  
 109    :  }  // namespace testing
 110    :  
 111    :  #endif  // SYZYGY_PE_HOT_PATCHING_UNITTEST_UTIL_H_