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/detectors/unittest_util.h"
  16    :  
  17    :  #include "base/environment.h"
  18    :  #include "base/strings/utf_string_conversions.h"
  19    :  #include "syzygy/pe/find.h"
  20    :  #include "syzygy/pe/pe_file.h"
  21    :  #include "syzygy/refinery/types/dia_crawler.h"
  22    :  
  23  m :  namespace testing {
  24    :  
  25  m :  namespace {
  26    :  
  27    :  // TODO(siggi): Remove dupes of this function.
  28  m :  bool GetNtdllTypes(refinery::TypeRepository* repo) {
  29    :    // As of 28/10/2015 the symbol file for ntdll.dll on Win7 is missing the
  30    :    // crucial symbols for heap enumeration. This code deserves to either die
  31    :    // in a fire, or else be updated to find symbols that are close to the
  32    :    // system in version and bitness.
  33  m :    pe::PEFile::Signature ntdll_sig(L"ntdll.dll", core::AbsoluteAddress(0),
  34  m :                                    0x141000, 0, 0x560D708C);
  35    :  
  36  m :    std::unique_ptr<base::Environment> env(base::Environment::Create());
  37  m :    std::string search_path;
  38  m :    if (!env->GetVar("_NT_SYMBOL_PATH", &search_path)) {
  39    :      // TODO(siggi): Set a default when it's missing.
  40  m :      LOG(ERROR) << "Missing symbol path.";
  41  m :      return false;
  42  m :    }
  43    :  
  44  m :    base::FilePath ntdll_path;
  45  m :    if (!pe::FindModuleBySignature(ntdll_sig, base::UTF8ToUTF16(search_path),
  46  m :                                   &ntdll_path)) {
  47  m :      LOG(ERROR) << "Failed to locate NTDLL.";
  48  m :      return false;
  49  m :    }
  50    :  
  51  m :    refinery::DiaCrawler crawler;
  52  m :    if (!crawler.InitializeForFile(base::FilePath(ntdll_path)) ||
  53  m :        !crawler.GetTypes(repo)) {
  54  m :      LOG(ERROR) << "Failed to get ntdll types.";
  55  m :      return false;
  56  m :    }
  57    :  
  58  m :    return true;
  59  m :  }
  60    :  
  61  m :  }  // namespace
  62    :  
  63  m :  LFHDetectorTest::LFHDetectorTest() {
  64  m :  }
  65    :  
  66  m :  void LFHDetectorTest::SetUp() {
  67  m :    ASSERT_TRUE(scoped_symbol_path_.Setup());
  68    :  
  69  m :    repo_ = new refinery::TypeRepository;
  70  m :    ASSERT_TRUE(scoped_heap_.Create());
  71  m :    ASSERT_TRUE(testing::GetNtdllTypes(repo_.get()));
  72  m :  }
  73    :  
  74  m :  void LFHDetectorTest::TearDown() {
  75  m :  }
  76    :  
  77  m :  refinery::Address LFHDetectorTest::AllocateLFHBucket(size_t block_size) {
  78  m :    for (size_t i = 0; i < 10000; ++i) {
  79  m :      void* ptr = scoped_heap_.Allocate(block_size);
  80    :  
  81  m :      if (scoped_heap_.IsLFHBlock(ptr))
  82  m :        return reinterpret_cast<refinery::Address>(ptr);
  83  m :    }
  84    :  
  85  m :    return 0;
  86  m :  }
  87    :  
  88  m :  }  // namespace testing