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/explode_basic_blocks_transform.h"
  16    :  
  17    :  #include "gtest/gtest.h"
  18    :  #include "syzygy/block_graph/orderers/random_orderer.h"
  19    :  #include "syzygy/core/unittest_util.h"
  20    :  #include "syzygy/pe/decomposer.h"
  21    :  #include "syzygy/pe/pe_relinker.h"
  22    :  #include "syzygy/pe/unittest_util.h"
  23    :  
  24    :  namespace pe {
  25    :  namespace transforms {
  26    :  
  27    :  using block_graph::BlockGraph;
  28    :  
  29    :  namespace {
  30    :  
  31    :  class ExplodeBasicBlocksTransformTest : public testing::PELibUnitTest {
  32    :   public:
  33    :    ExplodeBasicBlocksTransformTest()
  34    :        : image_layout_(&block_graph_),
  35    :          dos_header_block_(NULL),
  36  E :          input_path_(testing::GetExeRelativePath(testing::kTestDllName)) {
  37  E :    }
  38    :  
  39  E :    virtual void SetUp() override {
  40  E :      this->CreateTemporaryDir(&temp_dir_);
  41  E :      output_path_ = temp_dir_.Append(testing::kTestDllName);
  42  E :    }
  43    :  
  44  E :    void PerformRandomizationTest(ExplodeBasicBlocksTransform* transform) {
  45  E :      pe::PETransformPolicy policy;
  46  E :      pe::PERelinker relinker(&policy);
  47  E :      relinker.set_input_path(input_path_);
  48  E :      relinker.set_output_path(output_path_);
  49  E :      relinker.set_padding(8);
  50  E :      relinker.set_add_metadata(true);
  51  E :      relinker.set_allow_overwrite(true);
  52  E :      relinker.set_augment_pdb(true);
  53  E :      ASSERT_TRUE(relinker.Init());
  54    :  
  55  E :      relinker.AppendTransform(transform);
  56    :  
  57  E :      block_graph::orderers::RandomOrderer orderer(true, 123456);
  58  E :      relinker.AppendOrderer(&orderer);
  59    :  
  60    :      // Perform the actual relink.
  61  E :      ASSERT_TRUE(relinker.Relink());
  62    :  
  63    :      // Validate that the binary still loads.
  64  E :      ASSERT_NO_FATAL_FAILURE(CheckTestDll(output_path_));
  65  E :    }
  66    :  
  67    :    BlockGraph block_graph_;
  68    :    ImageLayout image_layout_;
  69    :    BlockGraph::Block* dos_header_block_;
  70    :    base::FilePath input_path_;
  71    :    base::FilePath temp_dir_;
  72    :    base::FilePath output_path_;
  73    :  };
  74    :  
  75    :  class DllMainRandomizer : public ExplodeBasicBlocksTransform {
  76    :   protected:
  77  E :    bool SkipThisBlock(const BlockGraph::Block* candidate) override {
  78  E :      return candidate->name() != "DllMain";
  79  E :    }
  80    :  };
  81    :  
  82    :  }  // namespace
  83    :  
  84  E :  TEST_F(ExplodeBasicBlocksTransformTest, RandomizeDllMain) {
  85  E :    DllMainRandomizer transform;
  86  E :    PerformRandomizationTest(&transform);
  87  E :  }
  88    :  
  89  E :  TEST_F(ExplodeBasicBlocksTransformTest, RandomizeAllBasicBlocks) {
  90  E :    ExplodeBasicBlocksTransform transform;
  91  E :    PerformRandomizationTest(&transform);
  92  E :  }
  93    :  
  94  E :  TEST_F(ExplodeBasicBlocksTransformTest, RandomizeAllBasicBlocksNoPadding) {
  95  E :    ExplodeBasicBlocksTransform transform;
  96  E :    transform.set_exclude_padding(true);
  97  E :    PerformRandomizationTest(&transform);
  98  E :  }
  99    :  
 100    :  }  // namespace transforms
 101    :  }  // namespace pe