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  m :  namespace kasko {
  27    :  
  28  m :  namespace {
  29    :  
  30  m :  base::FilePath GetModulePath(HMODULE module) {
  31  m :    wchar_t module_path[MAX_PATH];
  32  m :    DWORD size = ::GetModuleFileName(module, module_path, MAX_PATH);
  33  m :    DCHECK(size != 0);
  34    :  
  35  m :    return base::FilePath(module_path);
  36  m :  }
  37    :  
  38  m :  ULONGLONG GetModuleBaseAddress(HMODULE module) {
  39  m :    return reinterpret_cast<ULONGLONG>(module);
  40  m :  }
  41    :  
  42  m :  class LoaderLockTest : public ::testing::Test {
  43  m :   public:
  44  m :    LoaderLockTest() {}
  45  m :    ~LoaderLockTest() override {}
  46    :  
  47  m :   private:
  48  m :    void SetUp() override { ASSERT_TRUE(scoped_symbol_path_.Setup()); }
  49    :  
  50  m :    testing::ScopedSymbolPath scoped_symbol_path_;
  51    :  
  52  m :    DISALLOW_COPY_AND_ASSIGN(LoaderLockTest);
  53  m :  };
  54    :  
  55  m :  }  // 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    :  //
  61    :  // NOTE: Currently disabled as it relies on being able to retrieve the symbols
  62    :  // for ntdll.dll, which doesn't always work on the bots.
  63  m :  TEST_F(LoaderLockTest, DISABLED_SymbolOffset) {
  64  m :    HMODULE ntdll_handle = ::GetModuleHandle(L"ntdll.dll");
  65    :  
  66  m :    base::FilePath ntdll_path = GetModulePath(ntdll_handle);
  67  m :    ASSERT_FALSE(ntdll_path.empty());
  68    :  
  69  m :    base::FilePath ntdll_pdb_path;
  70  m :    ASSERT_TRUE(pe::FindPdbForModule(ntdll_path, &ntdll_pdb_path));
  71  m :    ASSERT_FALSE(ntdll_pdb_path.empty());
  72    :  
  73    :    // Open the pdb file.
  74  m :    base::win::ScopedComPtr<IDiaDataSource> source;
  75  m :    ASSERT_TRUE(pe::CreateDiaSource(source.Receive()));
  76  m :    base::win::ScopedComPtr<IDiaSession> session;
  77  m :    ASSERT_TRUE(
  78  m :        pe::CreateDiaSession(ntdll_pdb_path, source.get(), session.Receive()));
  79    :  
  80    :    // Set the load address of the dia session to get the computed virtual address
  81    :    // of the loader lock.
  82  m :    ASSERT_HRESULT_SUCCEEDED(
  83  m :        session->put_loadAddress(GetModuleBaseAddress(ntdll_handle)));
  84    :  
  85  m :    base::win::ScopedComPtr<IDiaSymbol> global_scope;
  86  m :    ASSERT_HRESULT_SUCCEEDED(session->get_globalScope(global_scope.Receive()));
  87    :  
  88    :    // Find the loader lock using its symbol name.
  89  m :    base::win::ScopedComPtr<IDiaEnumSymbols> symbols_enum;
  90  m :    ASSERT_HRESULT_SUCCEEDED(
  91  m :        global_scope->findChildren(SymTagPublicSymbol, L"_LdrpLoaderLock",
  92  m :                                   nsfCaseSensitive, symbols_enum.Receive()));
  93    :  
  94    :    // Sanity check. Only one symbol should have been found.
  95  m :    LONG count = 0;
  96  m :    ASSERT_HRESULT_SUCCEEDED(symbols_enum->get_Count(&count));
  97  m :    ASSERT_EQ(1, count);
  98    :  
  99  m :    base::win::ScopedComPtr<IDiaSymbol> loader_lock_symbol;
 100  m :    ASSERT_HRESULT_SUCCEEDED(symbols_enum->Item(0, loader_lock_symbol.Receive()));
 101  m :    ULONGLONG loader_lock_va = 0;
 102  m :    ASSERT_HRESULT_SUCCEEDED(
 103  m :        loader_lock_symbol->get_virtualAddress(&loader_lock_va));
 104    :  
 105  m :    ASSERT_EQ(loader_lock_va, reinterpret_cast<uintptr_t>(GetLoaderLock()));
 106  m :  }
 107    :  
 108  m :  }  // namespace kasko