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    :  namespace testing {
  24    :  
  25    :  namespace {
  26    :  
  27    :  // TODO(siggi): Remove dupes of this function.
  28  E :  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    :    pe::PEFile::Signature ntdll_sig(L"ntdll.dll", core::AbsoluteAddress(0),
  34  E :                                    0x141000, 0, 0x560D708C);
  35    :  
  36  E :    scoped_ptr<base::Environment> env(base::Environment::Create());
  37  E :    std::string search_path;
  38  E :    if (!env->GetVar("_NT_SYMBOL_PATH", &search_path)) {
  39    :      // TODO(siggi): Set a default when it's missing.
  40  i :      LOG(ERROR) << "Missing symbol path.";
  41  i :      return false;
  42    :    }
  43    :  
  44  E :    base::FilePath ntdll_path;
  45    :    if (!pe::FindModuleBySignature(ntdll_sig, base::UTF8ToUTF16(search_path),
  46  E :                                   &ntdll_path)) {
  47  i :      LOG(ERROR) << "Failed to locate NTDLL.";
  48  i :      return false;
  49    :    }
  50    :  
  51  E :    refinery::DiaCrawler crawler;
  52    :    if (!crawler.InitializeForFile(base::FilePath(ntdll_path)) ||
  53  E :        !crawler.GetTypes(repo)) {
  54  i :      LOG(ERROR) << "Failed to get ntdll types.";
  55  i :      return false;
  56    :    }
  57    :  
  58  E :    return true;
  59  E :  }
  60    :  
  61    :  }  // namespace
  62    :  
  63  E :  LFHDetectorTest::LFHDetectorTest() {
  64  E :  }
  65    :  
  66  E :  void LFHDetectorTest::SetUp() {
  67  E :    ASSERT_TRUE(scoped_symbol_path_.Setup());
  68    :  
  69  E :    repo_ = new refinery::TypeRepository;
  70  E :    ASSERT_TRUE(scoped_heap_.Create());
  71  E :    ASSERT_TRUE(testing::GetNtdllTypes(repo_.get()));
  72  E :  }
  73    :  
  74  E :  void LFHDetectorTest::TearDown() {
  75  E :  }
  76    :  
  77  E :  refinery::Address LFHDetectorTest::AllocateLFHBucket(size_t block_size) {
  78  E :    for (size_t i = 0; i < 10000; ++i) {
  79  E :      void* ptr = scoped_heap_.Allocate(block_size);
  80    :  
  81  E :      if (scoped_heap_.IsLFHBlock(ptr))
  82  E :        return reinterpret_cast<refinery::Address>(ptr);
  83  E :    }
  84    :  
  85  i :    return 0;
  86  E :  }
  87    :  
  88    :  }  // namespace testing