1    :  // Copyright 2012 Google Inc.
   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    :  // Unittests for ThunkImportReferencesTransform.
  16    :  
  17    :  #include "syzygy/instrument/transforms/thunk_import_references_transform.h"
  18    :  
  19    :  #include <vector>
  20    :  
  21    :  #include "base/scoped_temp_dir.h"
  22    :  #include "gtest/gtest.h"
  23    :  #include "syzygy/block_graph/typed_block.h"
  24    :  #include "syzygy/core/unittest_util.h"
  25    :  #include "syzygy/pe/decomposer.h"
  26    :  #include "syzygy/pe/transforms/add_imports_transform.h"
  27    :  
  28    :  namespace instrument {
  29    :  namespace transforms {
  30    :  
  31    :  namespace {
  32    :  
  33    :  using block_graph::BlockGraph;
  34    :  using block_graph::TypedBlock;
  35    :  using core::AbsoluteAddress;
  36    :  using pe::transforms::AddImportsTransform;
  37    :  
  38    :  typedef BlockGraph::Block::ReferrerSet ReferrerSet;
  39    :  
  40    :  // Expose protected members for testing.
  41    :  class TestThunkImportReferencesTransform
  42    :      : public ThunkImportReferencesTransform {
  43    :   public:
  44    :    using ThunkImportReferencesTransform::Thunk;
  45    :  
  46  E :    AddImportsTransform& add_imports_transform() {
  47  E :      return ThunkImportReferencesTransform::add_imports_transform();
  48  E :    }
  49    :  };
  50    :  
  51    :  class ThunkImportReferencesTransformTest : public testing::Test {
  52    :   public:
  53  E :    ThunkImportReferencesTransformTest()
  54    :        : input_image_layout_(&block_graph_), dos_header_block_(NULL) {
  55  E :    }
  56    :  
  57  E :    virtual void SetUp() {
  58  E :      tmp_dir_.CreateUniqueTempDir();
  59  E :      ASSERT_TRUE(input_pe_file_.Init(
  60    :          testing::GetOutputRelativePath(L"test_dll.dll")));
  61  E :      ASSERT_NO_FATAL_FAILURE(Decompose());
  62  E :    }
  63    :  
  64  E :    void Decompose() {
  65    :      // Decompose the input image.
  66  E :      pe::Decomposer decomposer(input_pe_file_);
  67    :      decomposer.set_pdb_path(
  68  E :          FilePath(testing::GetOutputRelativePath(L"test_dll.pdb")));
  69  E :      ASSERT_TRUE(decomposer.Decompose(&input_image_layout_))
  70    :          << "Unable to decompose module: "
  71  E :          << input_pe_file_.path().value();
  72    :  
  73    :      // Get the DOS header block.
  74    :      dos_header_block_ =
  75    :          input_image_layout_.blocks.GetBlockByAddress(
  76  E :              BlockGraph::RelativeAddress(0));
  77  E :      ASSERT_TRUE(dos_header_block_ != NULL)
  78  E :          << "Unable to find the DOS header block.";
  79  E :    }
  80    :  
  81    :   protected:
  82    :    pe::PEFile input_pe_file_;
  83    :    pe::ImageLayout input_image_layout_;
  84    :    BlockGraph block_graph_;
  85    :    BlockGraph::Block* dos_header_block_;
  86    :    ScopedTempDir tmp_dir_;
  87    :  };
  88    :  
  89    :  }  // namespace
  90    :  
  91  E :  TEST_F(ThunkImportReferencesTransformTest, TestImports) {
  92    :    // Check that we don't have a thunks section yet.
  93    :    BlockGraph::Section* null_thunks_section =
  94  E :        block_graph_.FindSection(".thunks");
  95  E :    EXPECT_EQ(NULL, null_thunks_section);
  96    :  
  97  E :    TestThunkImportReferencesTransform transform;
  98    :    ASSERT_TRUE(ApplyBlockGraphTransform(
  99  E :        &transform, &block_graph_, dos_header_block_));
 100    :  
 101    :    // Check that we now have a thunks section.
 102  E :    BlockGraph::Section* thunks_section = block_graph_.FindSection(".thunks");
 103  E :    ASSERT_TRUE(thunks_section != NULL);
 104    :  
 105  E :    BlockGraph::SectionId thunks_section_id = thunks_section->id();
 106    :  
 107    :    // The only CODE_BLOCK referrers to the IAT should be in the thunk section.
 108  E :    AddImportsTransform& ait = transform.add_imports_transform();
 109    :    const ReferrerSet& referrers  =
 110  E :        ait.import_address_table_block()->referrers();
 111  E :    ReferrerSet::const_iterator ref_iter(referrers.begin()),
 112  E :                                ref_end(referrers.end());
 113  E :    for (; ref_iter != ref_end; ++ref_iter) {
 114    :      // Use this check to exclude the NT headers and the IDT blocks.
 115  E :      if (ref_iter->first->type() == BlockGraph::CODE_BLOCK)
 116  E :        EXPECT_EQ(thunks_section_id, ref_iter->first->section());
 117  E :    }
 118  E :  }
 119    :  
 120    :  }  // namespace transforms
 121    :  }  // namespace instrument