1    :  // Copyright 2011 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/wsdump/process_working_set.h"
  16    :  
  17    :  #include <psapi.h>
  18    :  #include <set>
  19    :  #include <vector>
  20    :  
  21    :  #include "base/at_exit.h"
  22    :  #include "base/memory/scoped_ptr.h"
  23    :  #include "base/strings/string_util.h"
  24    :  #include "base/win/scoped_handle.h"
  25    :  #include "gtest/gtest.h"
  26    :  #include "syzygy/common/com_utils.h"
  27    :  #include "syzygy/core/address_space.h"
  28    :  
  29    :  namespace wsdump {
  30    :  
  31    :  class TestingProcessWorkingSet: public ProcessWorkingSet {
  32    :   public:
  33    :    using ProcessWorkingSet::ScopedWsPtr;
  34    :    using ProcessWorkingSet::CaptureWorkingSet;
  35    :  
  36    :    using ProcessWorkingSet::ModuleAddressSpace;
  37    :    using ProcessWorkingSet::CaptureModules;
  38    :  };
  39    :  
  40  E :  TEST(ProcessWorkingSetTest, CaptureWorkingSet) {
  41  E :    TestingProcessWorkingSet::ScopedWsPtr working_set;
  42    :    EXPECT_TRUE(TestingProcessWorkingSet::CaptureWorkingSet(::GetCurrentProcess(),
  43  E :                                                            &working_set));
  44  E :    EXPECT_TRUE(working_set.get() != NULL);
  45  E :  }
  46    :  
  47    :  // This function gives us an address in our module.
  48  i :  static void dummy() {
  49  i :  }
  50    :  
  51  E :  TEST(ProcessWorkingSetTest, CaptureModules) {
  52    :    typedef TestingProcessWorkingSet::ModuleAddressSpace ModuleAddressSpace;
  53  E :    ModuleAddressSpace modules;
  54    :    EXPECT_TRUE(TestingProcessWorkingSet::CaptureModules(::GetCurrentProcessId(),
  55  E :                                                         &modules));
  56  E :    EXPECT_LT(0U, modules.ranges().size());
  57    :  
  58  E :    ModuleAddressSpace::Range range(reinterpret_cast<size_t>(&dummy), 1);
  59  E :    EXPECT_TRUE(modules.FindContaining(range) != modules.end());
  60  E :  }
  61    :  
  62  E :  TEST(ProcessWorkingSetTest, Initialize) {
  63  E :    ProcessWorkingSet ws;
  64  E :    ASSERT_TRUE(ws.Initialize(::GetCurrentProcessId()));
  65    :  
  66    :    // Double-check the accounting.
  67  E :    std::set<std::wstring> module_names;
  68  E :    ProcessWorkingSet::Stats total_modules;
  69  E :    ProcessWorkingSet::ModuleStatsVector::const_iterator it;
  70  E :    for (it = ws.module_stats().begin(); it != ws.module_stats().end(); ++it) {
  71  E :      const ProcessWorkingSet::ModuleStats& stats = *it;
  72    :  
  73    :      // Each module name must occur precisely once.
  74  E :      EXPECT_TRUE(module_names.insert(stats.module_name).second);
  75    :  
  76  E :      total_modules.pages += stats.pages;
  77  E :      total_modules.shareable_pages += stats.shareable_pages;
  78  E :      total_modules.shared_pages += stats.shared_pages;
  79  E :      total_modules.read_only_pages += stats.read_only_pages;
  80  E :      total_modules.writable_pages += stats.writable_pages;
  81  E :      total_modules.executable_pages += stats.executable_pages;
  82  E :    }
  83    :  
  84    :    // Our executable should be in the working set.
  85  E :    std::wstring exe_name;
  86    :    ASSERT_TRUE(::GetModuleFileName(NULL, base::WriteInto(&exe_name, MAX_PATH),
  87  E :                                    MAX_PATH));
  88  E :    exe_name.resize(wcslen(exe_name.c_str()));
  89    :  
  90  E :    EXPECT_TRUE(module_names.find(exe_name) != module_names.end());
  91    :  
  92    :    // And finally check the tally.
  93    :    EXPECT_EQ(ws.total_stats().pages,
  94  E :              total_modules.pages + ws.non_module_stats().pages);
  95    :    EXPECT_EQ(ws.total_stats().shareable_pages,
  96  E :        total_modules.shareable_pages + ws.non_module_stats().shareable_pages);
  97    :    EXPECT_EQ(ws.total_stats().shared_pages,
  98  E :        total_modules.shared_pages + ws.non_module_stats().shared_pages);
  99    :    EXPECT_EQ(ws.total_stats().read_only_pages,
 100  E :        total_modules.read_only_pages + ws.non_module_stats().read_only_pages);
 101    :    EXPECT_EQ(ws.total_stats().writable_pages,
 102  E :        total_modules.writable_pages + ws.non_module_stats().writable_pages);
 103    :    EXPECT_EQ(ws.total_stats().executable_pages,
 104  E :        total_modules.executable_pages + ws.non_module_stats().executable_pages);
 105  E :  }
 106    :  
 107    :  }  // namespace wsdump