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/kasko/loader_lock.h"
  16    :  
  17    :  #include "Psapi.h"
  18    :  
  19    :  #include "base/files/file_path.h"
  20    :  #include "base/strings/utf_string_conversions.h"
  21    :  #include "gtest/gtest.h"
  22    :  #include "syzygy/common/unittest_util.h"
  23    :  #include "syzygy/pe/dia_util.h"
  24    :  #include "syzygy/pe/find.h"
  25    :  
  26    :  namespace kasko {
  27    :  
  28    :  namespace {
  29    :  
  30  E :  base::FilePath GetModulePath(HMODULE module) {
  31    :    wchar_t module_path[MAX_PATH];
  32  E :    DWORD size = ::GetModuleFileName(module, module_path, MAX_PATH);
  33  E :    DCHECK(size != 0);
  34    :  
  35  E :    return base::FilePath(module_path);
  36  E :  }
  37    :  
  38  E :  ULONGLONG GetModuleBaseAddress(HMODULE module) {
  39  E :    return reinterpret_cast<ULONGLONG>(module);
  40  E :  }
  41    :  
  42    :  class LoaderLockTest : public ::testing::Test {
  43    :   public:
  44  E :    LoaderLockTest() {}
  45  E :    ~LoaderLockTest() override {}
  46    :  
  47    :   private:
  48  E :    void SetUp() override { ASSERT_TRUE(scoped_symbol_path_.Setup()); }
  49    :  
  50    :    testing::ScopedSymbolPath scoped_symbol_path_;
  51    :  
  52    :    DISALLOW_COPY_AND_ASSIGN(LoaderLockTest);
  53    :  };
  54    :  
  55    :  }  // namespace
  56    :  
  57    :  // Tests that the offset used for finding the loader lock address in the Process
  58    :  // Environment Block is correct. This is done by looking into the PDB file for
  59    :  // ntdll.
  60  E :  TEST_F(LoaderLockTest, SymbolOffset) {
  61  E :    HMODULE ntdll_handle = ::GetModuleHandle(L"ntdll.dll");
  62    :  
  63  E :    base::FilePath ntdll_path = GetModulePath(ntdll_handle);
  64  E :    ASSERT_FALSE(ntdll_path.empty());
  65    :  
  66  E :    base::FilePath ntdll_pdb_path;
  67  E :    ASSERT_TRUE(pe::FindPdbForModule(ntdll_path, &ntdll_pdb_path));
  68  E :    ASSERT_FALSE(ntdll_pdb_path.empty());
  69    :  
  70    :    // Open the pdb file.
  71  E :    base::win::ScopedComPtr<IDiaDataSource> source;
  72  E :    ASSERT_TRUE(pe::CreateDiaSource(source.Receive()));
  73  E :    base::win::ScopedComPtr<IDiaSession> session;
  74    :    ASSERT_TRUE(
  75  E :        pe::CreateDiaSession(ntdll_pdb_path, source.get(), session.Receive()));
  76    :  
  77    :    // Set the load address of the dia session to get the computed virtual address
  78    :    // of the loader lock.
  79    :    ASSERT_HRESULT_SUCCEEDED(
  80  E :        session->put_loadAddress(GetModuleBaseAddress(ntdll_handle)));
  81    :  
  82  E :    base::win::ScopedComPtr<IDiaSymbol> global_scope;
  83  E :    ASSERT_HRESULT_SUCCEEDED(session->get_globalScope(global_scope.Receive()));
  84    :  
  85    :    // Find the loader lock using its symbol name.
  86  E :    base::win::ScopedComPtr<IDiaEnumSymbols> symbols_enum;
  87    :    ASSERT_HRESULT_SUCCEEDED(
  88    :        global_scope->findChildren(SymTagPublicSymbol, L"_LdrpLoaderLock",
  89  E :                                   nsfCaseSensitive, symbols_enum.Receive()));
  90    :  
  91    :    // Sanity check. Only one symbol should have been found.
  92  E :    LONG count = 0;
  93  E :    ASSERT_HRESULT_SUCCEEDED(symbols_enum->get_Count(&count));
  94  E :    ASSERT_EQ(1, count);
  95    :  
  96  E :    base::win::ScopedComPtr<IDiaSymbol> loader_lock_symbol;
  97  E :    ASSERT_HRESULT_SUCCEEDED(symbols_enum->Item(0, loader_lock_symbol.Receive()));
  98  E :    ULONGLONG loader_lock_va = 0;
  99    :    ASSERT_HRESULT_SUCCEEDED(
 100  E :        loader_lock_symbol->get_virtualAddress(&loader_lock_va));
 101    :  
 102  E :    ASSERT_EQ(loader_lock_va, reinterpret_cast<uintptr_t>(GetLoaderLock()));
 103  E :  }
 104    :  
 105    :  }  // namespace kasko