1    :  // Copyright 2013 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/sampler/sampled_module_cache.h"
  16    :  
  17    :  #include "base/bind.h"
  18    :  #include "gmock/gmock.h"
  19    :  #include "gtest/gtest.h"
  20    :  
  21    :  namespace sampler {
  22    :  
  23    :  struct MockedCallbackStruct {
  24  E :    MOCK_METHOD1(OnDeadModule, void(const SampledModuleCache::Module*));
  25    :  };
  26    :  
  27    :  class SampledModuleCacheTest : public ::testing::Test {
  28    :   public:
  29  E :    virtual void SetUp() OVERRIDE {
  30  E :      testing::Test::SetUp();
  31    :  
  32    :      dead_module_callback = base::Bind(&MockedCallbackStruct::OnDeadModule,
  33  E :                                        base::Unretained(&mock));
  34  E :    }
  35    :  
  36  E :    bool IsAlive(const SampledModuleCache::Process* process) {
  37  E :      return process->alive();
  38  E :    }
  39    :  
  40  E :    bool IsAlive(const SampledModuleCache::Module* module) {
  41  E :      return module->alive();
  42  E :    }
  43    :  
  44    :    ::testing::StrictMock<MockedCallbackStruct> mock;
  45    :    SampledModuleCache::DeadModuleCallback dead_module_callback;
  46    :  };
  47    :  
  48  E :  TEST_F(SampledModuleCacheTest, ConstructorAndProperties) {
  49  E :    SampledModuleCache cache(2);
  50  E :    EXPECT_EQ(2u, cache.log2_bucket_size());
  51    :  
  52  E :    EXPECT_TRUE(cache.dead_module_callback().is_null());
  53  E :    cache.set_dead_module_callback(dead_module_callback);
  54  E :    EXPECT_FALSE(cache.dead_module_callback().is_null());
  55  E :  }
  56    :  
  57  E :  TEST_F(SampledModuleCacheTest, EndToEnd) {
  58  E :    SampledModuleCache cache(2);
  59  E :    cache.set_dead_module_callback(dead_module_callback);
  60  E :    EXPECT_EQ(0u, cache.processes().size());
  61    :  
  62    :    static const DWORD kAccess =
  63    :        PROCESS_DUP_HANDLE | PROCESS_QUERY_INFORMATION | PROCESS_VM_READ;
  64    :    base::win::ScopedHandle proc(
  65  E :        ::OpenProcess(kAccess, FALSE, ::GetCurrentProcessId()));
  66  E :    ASSERT_TRUE(proc.IsValid());
  67    :  
  68    :    // Try to add an invalid module. This should fail.
  69  E :    HMODULE module_handle = ::GetModuleHandle(NULL);
  70  E :    EXPECT_FALSE(cache.AddModule(proc.Get(), module_handle - 1));
  71  E :    EXPECT_EQ(0u, cache.processes().size());
  72    :  
  73    :    // Add this module. This should succeed.
  74  E :    EXPECT_TRUE(cache.AddModule(proc.Get(), module_handle));
  75    :  
  76  E :    EXPECT_EQ(1u, cache.processes().size());
  77    :    const SampledModuleCache::Process* process =
  78  E :        cache.processes().begin()->second;
  79  E :    ASSERT_TRUE(process != NULL);
  80  E :    EXPECT_TRUE(IsAlive(process));
  81    :  
  82  E :    EXPECT_EQ(1u, cache.processes().begin()->second->modules().size());
  83  E :    const SampledModuleCache::Module* module = process->modules().begin()->second;
  84  E :    ASSERT_TRUE(module != NULL);
  85  E :    EXPECT_TRUE(IsAlive(module));
  86    :  
  87    :    // Mark the modules as dead.
  88  E :    cache.MarkAllModulesDead();
  89  E :    EXPECT_FALSE(IsAlive(process));
  90  E :    EXPECT_FALSE(IsAlive(module));
  91    :  
  92    :    // Re-add the module. This should simply mark the existing module as alive.
  93  E :    cache.AddModule(proc.Get(), module_handle);
  94  E :    EXPECT_EQ(1u, cache.processes().size());
  95  E :    EXPECT_EQ(process, cache.processes().begin()->second);
  96  E :    EXPECT_TRUE(IsAlive(process));
  97  E :    EXPECT_EQ(1u, process->modules().size());
  98  E :    EXPECT_EQ(module, process->modules().begin()->second);
  99  E :    EXPECT_TRUE(IsAlive(module));
 100    :  
 101    :    // Clean up the modules. Nothing should be removed and the callback should
 102    :    // not be invoked.
 103  E :    cache.RemoveDeadModules();
 104  E :    EXPECT_EQ(1u, cache.processes().size());
 105  E :    EXPECT_EQ(process, cache.processes().begin()->second);
 106  E :    EXPECT_TRUE(IsAlive(process));
 107  E :    EXPECT_EQ(1u, process->modules().size());
 108  E :    EXPECT_EQ(module, process->modules().begin()->second);
 109  E :    EXPECT_TRUE(IsAlive(module));
 110    :  
 111  E :    EXPECT_CALL(mock, OnDeadModule(module)).Times(1);
 112    :  
 113    :    // Mark everything as dead and clean up the modules. The callback should be
 114    :    // invoked and the cache should now be empty.
 115  E :    cache.MarkAllModulesDead();
 116  E :    cache.RemoveDeadModules();
 117  E :    EXPECT_EQ(0u, cache.processes().size());
 118  E :  }
 119    :  
 120    :  }  // namespace sampler