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    :  #include "syzygy/pe/transforms/add_imports_transform.h"
  16    :  
  17    :  #include "gtest/gtest.h"
  18    :  #include "syzygy/core/unittest_util.h"
  19    :  #include "syzygy/pe/decomposer.h"
  20    :  #include "syzygy/pe/pe_utils.h"
  21    :  #include "syzygy/pe/unittest_util.h"
  22    :  
  23    :  namespace pe {
  24    :  namespace transforms {
  25    :  
  26    :  using block_graph::BlockGraph;
  27    :  using core::RelativeAddress;
  28    :  typedef AddImportsTransform::ImportedModule ImportedModule;
  29    :  
  30    :  namespace {
  31    :  
  32    :  class AddImportsTransformTest : public testing::PELibUnitTest {
  33    :   public:
  34  E :    AddImportsTransformTest() : image_layout_(&block_graph_) {
  35  E :    }
  36    :  
  37  E :    virtual void SetUp() {
  38  E :      FilePath image_path(testing::GetExeRelativePath(kDllName));
  39    :  
  40  E :      ASSERT_TRUE(pe_file_.Init(image_path));
  41    :  
  42    :      // Decompose the test image and look at the result.
  43  E :      Decomposer decomposer(pe_file_);
  44  E :      ASSERT_TRUE(decomposer.Decompose(&image_layout_));
  45    :  
  46    :      // Retrieve and validate the DOS header.
  47    :      dos_header_block_ =
  48  E :          image_layout_.blocks.GetBlockByAddress(RelativeAddress(0));
  49  E :      ASSERT_TRUE(dos_header_block_ != NULL);
  50  E :      ASSERT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
  51  E :    }
  52    :  
  53    :    PEFile pe_file_;
  54    :    BlockGraph block_graph_;
  55    :    ImageLayout image_layout_;
  56    :    BlockGraph::Block* dos_header_block_;
  57    :  };
  58    :  
  59    :  // Given an ImportedModule tests that all of its symbols have been properly
  60    :  // processed.
  61  E :  void TestSymbols(const ImportedModule& module) {
  62  E :    for (size_t i = 0; i < module.size(); ++i) {
  63  E :      BlockGraph::Reference ref;
  64  E :      EXPECT_TRUE(module.GetSymbolReference(i, &ref));
  65  E :      EXPECT_TRUE(ref.referenced() != NULL);
  66  E :      EXPECT_GE(ref.offset(), 0);
  67    :      EXPECT_LT(ref.offset(),
  68  E :                static_cast<BlockGraph::Offset>(ref.referenced()->size()));
  69  E :    }
  70  E :  }
  71    :  
  72    :  }  // namespace
  73    :  
  74  E :  TEST_F(AddImportsTransformTest, AddImportsExisting) {
  75  E :    ImportedModule module("export_dll.dll");
  76  E :    size_t function1 = module.AddSymbol("function1");
  77  E :    size_t function3 = module.AddSymbol("function3");
  78  E :    EXPECT_EQ("function1", module.GetSymbolName(function1));
  79  E :    EXPECT_EQ("function3", module.GetSymbolName(function3));
  80    :  
  81  E :    AddImportsTransform transform;
  82  E :    transform.AddModule(&module);
  83    :    EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
  84  E :        &transform, &block_graph_, dos_header_block_));
  85  E :    EXPECT_EQ(0u, transform.modules_added());
  86  E :    EXPECT_EQ(0u, transform.symbols_added());
  87    :  
  88  E :    EXPECT_NO_FATAL_FAILURE(TestSymbols(module));
  89  E :  }
  90    :  
  91  E :  TEST_F(AddImportsTransformTest, AddImportsNewSymbol) {
  92  E :    ImportedModule module("export_dll.dll");
  93  E :    size_t function1 = module.AddSymbol("function1");
  94  E :    size_t function3 = module.AddSymbol("function3");
  95  E :    size_t function4 = module.AddSymbol("function4");
  96  E :    EXPECT_EQ("function1", module.GetSymbolName(function1));
  97  E :    EXPECT_EQ("function3", module.GetSymbolName(function3));
  98  E :    EXPECT_EQ("function4", module.GetSymbolName(function4));
  99    :  
 100  E :    AddImportsTransform transform;
 101  E :    transform.AddModule(&module);
 102    :    EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
 103  E :        &transform, &block_graph_, dos_header_block_));
 104  E :    EXPECT_EQ(0u, transform.modules_added());
 105  E :    EXPECT_EQ(1u, transform.symbols_added());
 106    :  
 107  E :    EXPECT_NO_FATAL_FAILURE(TestSymbols(module));
 108    :  
 109    :    // TODO(chrisha): Write the image and try to load it!
 110  E :  }
 111    :  
 112  E :  TEST_F(AddImportsTransformTest, AddImportsNewModule) {
 113  E :    ImportedModule module("call_trace_client_rpc.dll");
 114  E :    size_t indirect_penter = module.AddSymbol("_indirect_penter");
 115  E :    size_t indirect_penter_dllmain = module.AddSymbol("_indirect_penter_dllmain");
 116    :    EXPECT_EQ("_indirect_penter",
 117  E :              module.GetSymbolName(indirect_penter));
 118    :    EXPECT_EQ("_indirect_penter_dllmain",
 119  E :              module.GetSymbolName(indirect_penter_dllmain));
 120    :  
 121  E :    AddImportsTransform transform;
 122  E :    transform.AddModule(&module);
 123    :    EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
 124  E :        &transform, &block_graph_, dos_header_block_));
 125  E :    EXPECT_EQ(1u, transform.modules_added());
 126  E :    EXPECT_EQ(2u, transform.symbols_added());
 127    :  
 128  E :    EXPECT_NO_FATAL_FAILURE(TestSymbols(module));
 129    :  
 130    :    // TODO(chrisha): Write the image and try to load it!
 131  E :  }
 132    :  
 133    :  }  // namespace transforms
 134    :  }  // namespace pe