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/refinery/validators/vftable_ptr_validator.h"
  16    :  
  17    :  #include <dia2.h>
  18    :  
  19    :  #include "base/containers/hash_tables.h"
  20    :  #include "base/win/scoped_comptr.h"
  21    :  #include "gmock/gmock.h"
  22    :  #include "gtest/gtest.h"
  23    :  #include "syzygy/pe/pe_file.h"
  24    :  #include "syzygy/refinery/process_state/process_state.h"
  25    :  #include "syzygy/refinery/process_state/process_state_util.h"
  26    :  
  27    :  namespace refinery {
  28    :  
  29    :  namespace {
  30    :  
  31    :  const Address kAddress = 1000ULL;  // Fits 32-bit.
  32    :  const Address kAddressOther = 2000ULL;  // Fits 32-bit.
  33    :  const Size kSize = 42U;
  34    :  const Size kSizeOther = 43U;
  35    :  const uint32 kChecksum = 11U;
  36    :  const uint32 kChecksumOther = 12U;
  37    :  const uint32 kTimestamp = 22U;
  38    :  const wchar_t kPath[] = L"c:\\path\\ModuleName";
  39    :  const wchar_t kPathOther[] = L"c:\\path\\ModuleNameOther";
  40    :  
  41    :  class MockDiaSymbolProvider : public DiaSymbolProvider {
  42    :   public:
  43    :    MOCK_METHOD2(FindOrCreateDiaSession,
  44    :                 bool(const pe::PEFile::Signature& signature,
  45  i :                      base::win::ScopedComPtr<IDiaSession>* session));
  46    :    MOCK_METHOD2(GetVFTableRVAs,
  47    :                 bool(const pe::PEFile::Signature& signature,
  48  E :                      base::hash_set<Address>* vftable_rvas));
  49    :  };
  50    :  
  51    :  class TestVftablePtrValidator : public VftablePtrValidator {
  52    :   public:
  53    :    using VftablePtrValidator::GetVFTableVAs;
  54    :  };
  55    :  
  56    :  }  // namespace
  57    :  
  58  E :  TEST(VftablePtrValidatorTest, BasicTest) {
  59    :    // TODO(manzagop): implement.
  60  E :  }
  61    :  
  62  E :  TEST(VftablePtrValidatorTest, GetVFTableVAs) {
  63    :    using testing::_;
  64    :    using testing::DoAll;
  65    :    using testing::Return;
  66    :    using testing::SetArgPointee;
  67    :  
  68    :    // Create a process state with 2 modules.
  69  E :    ProcessState state;
  70  E :    ModuleLayerAccessor accessor(&state);
  71    :    accessor.AddModuleRecord(AddressRange(kAddress, kSize), kChecksum, kTimestamp,
  72  E :                             kPath);
  73    :    accessor.AddModuleRecord(AddressRange(kAddressOther, kSizeOther),
  74  E :                             kChecksumOther, kTimestamp, kPathOther);
  75    :  
  76    :    // Set up the symbol provider.
  77  E :    scoped_refptr<MockDiaSymbolProvider> provider = new MockDiaSymbolProvider();
  78    :  
  79  E :    pe::PEFile::Signature signature;
  80  E :    ASSERT_TRUE(accessor.GetModuleSignature(kAddress, &signature));
  81  E :    signature.base_address = core::AbsoluteAddress(0U);
  82  E :    base::hash_set<Address> rvas;
  83  E :    rvas.insert(1ULL);
  84  E :    rvas.insert(2ULL);
  85    :    EXPECT_CALL(*provider, GetVFTableRVAs(signature, testing::_))
  86  E :        .WillOnce(DoAll(SetArgPointee<1>(rvas), Return(true)));
  87    :  
  88  E :    pe::PEFile::Signature signature_other;
  89  E :    ASSERT_TRUE(accessor.GetModuleSignature(kAddressOther, &signature_other));
  90  E :    signature_other.base_address = core::AbsoluteAddress(0U);
  91  E :    base::hash_set<Address> rvas_other;
  92  E :    rvas_other.insert(3ULL);
  93  E :    rvas_other.insert(4ULL);
  94    :    EXPECT_CALL(*provider, GetVFTableRVAs(signature_other, testing::_))
  95  E :        .WillOnce(DoAll(SetArgPointee<1>(rvas_other), Return(true)));
  96    :  
  97    :    // Retrieve VAs and validate.
  98  E :    base::hash_set<Address> vftable_vas;
  99    :    ASSERT_TRUE(
 100  E :        TestVftablePtrValidator::GetVFTableVAs(&state, provider, &vftable_vas));
 101    :  
 102  E :    base::hash_set<Address> expected_vftable_vas;
 103  E :    expected_vftable_vas.insert(kAddress + 1ULL);
 104  E :    expected_vftable_vas.insert(kAddress + 2ULL);
 105  E :    expected_vftable_vas.insert(kAddressOther + 3ULL);
 106  E :    expected_vftable_vas.insert(kAddressOther + 4ULL);
 107    :  
 108  E :    ASSERT_EQ(expected_vftable_vas, vftable_vas);
 109  E :  }
 110    :  
 111    :  }  // namespace refinery