1 : // Copyright 2011 Google Inc.
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 <algorithm>
18 : #include <map>
19 :
20 : #include <psapi.h>
21 : #include <tlhelp32.h>
22 :
23 : #include "base/memory/scoped_ptr.h"
24 : #include "base/win/scoped_handle.h"
25 : #include "sawbuck/common/com_utils.h"
26 : #include "syzygy/core/address_space.h"
27 :
28 : namespace wsdump {
29 :
30 : namespace {
31 :
32 : const size_t kPageSize = 4096;
33 : // These are inferred from the MSDN page for QueryWorkingSet.
34 : const int kPageReadOnly = 0x001;
35 : const int kPageExecute = 0x002;
36 : const int kPageExecuteRead = 0x003;
37 : const int kPageReadWrite = 0x004;
38 : const int kPageWriteCopy = 0x005;
39 : const int kPageExecuteReadWrite = 0x006;
40 : const int kPageExecuteWriteCopy = 0x007;
41 :
42 : bool LessModuleName(const ProcessWorkingSet::ModuleStats& a,
43 E : const ProcessWorkingSet::ModuleStats& b) {
44 E : return a.module_name < b.module_name;
45 E : }
46 :
47 : } // namespace
48 :
49 :
50 E : bool ProcessWorkingSet::Initialize(DWORD process_id) {
51 E : ModuleAddressSpace modules;
52 E : if (!CaptureModules(process_id, &modules))
53 i : return false;
54 :
55 E : const DWORD kProcessPermissions = PROCESS_QUERY_INFORMATION | PROCESS_VM_READ;
56 : base::win::ScopedHandle process(
57 E : ::OpenProcess(kProcessPermissions, FALSE, process_id));
58 E : if (!process.IsValid()) {
59 i : DWORD err = ::GetLastError();
60 i : LOG(ERROR) << "OpenProcess failed: " << com::LogWe(err);
61 i : return false;
62 : }
63 :
64 E : ScopedWsPtr working_set;
65 E : if (!CaptureWorkingSet(process.Get(), &working_set))
66 i : return false;
67 :
68 : // The new stats we're building.
69 E : ModuleStatsVector new_stats;
70 :
71 : // This maps from module name to index in the above vector.
72 : typedef std::map<std::wstring, size_t> NameToIndexMap;
73 E : NameToIndexMap name_to_index;
74 E : for (size_t i = 0; i < working_set->NumberOfEntries; ++i) {
75 E : PSAPI_WORKING_SET_BLOCK entry = working_set->WorkingSetInfo[i];
76 :
77 E : size_t address = entry.VirtualPage * kPageSize;
78 E : ModuleAddressSpace::Range page_range(address, kPageSize);
79 : ModuleAddressSpace::RangeMap::const_iterator it =
80 E : modules.FindContaining(page_range);
81 :
82 E : Stats* stats = NULL;
83 E : if (it == modules.end()) {
84 E : stats = &non_module_stats_;
85 E : } else {
86 : // Find the module with this name, or add it if it's missing.
87 E : const std::wstring& module_name = it->second;
88 E : NameToIndexMap::const_iterator it = name_to_index.find(module_name);
89 E : if (it == name_to_index.end()) {
90 : // We haven't seen this module, add it to the end of the vector.
91 E : name_to_index[module_name] = new_stats.size();
92 E : new_stats.push_back(ModuleStats());
93 :
94 E : ModuleStats* module_stats = &new_stats.back();
95 E : module_stats->module_name = module_name;
96 :
97 E : stats = module_stats;
98 E : } else {
99 E : stats = &new_stats[it->second];
100 : }
101 E : }
102 :
103 E : DCHECK(stats != NULL);
104 :
105 E : total_stats_.pages++;
106 E : stats->pages++;
107 E : if (entry.Shared) {
108 E : total_stats_.shareable_pages++;
109 E : stats->shareable_pages++;
110 : }
111 :
112 E : if (entry.ShareCount > 1) {
113 E : total_stats_.shared_pages++;
114 E : stats->shared_pages++;
115 : }
116 :
117 E : if (entry.Protection & kPageReadWrite) {
118 E : total_stats_.writable_pages++;
119 E : stats->writable_pages++;
120 E : } else if (entry.Protection & kPageExecute) {
121 E : total_stats_.executable_pages++;
122 E : stats->executable_pages++;
123 E : } else if (entry.Protection & kPageReadOnly) {
124 E : total_stats_.read_only_pages++;
125 E : stats->read_only_pages++;
126 : }
127 E : }
128 :
129 E : std::sort(new_stats.begin(), new_stats.end(), LessModuleName);
130 E : new_stats.swap(module_stats_);
131 E : return true;
132 E : }
133 :
134 : bool ProcessWorkingSet::CaptureWorkingSet(HANDLE process,
135 E : ScopedWsPtr* working_set) {
136 E : DCHECK(working_set != NULL);
137 :
138 : // Estimate the starting buffer size by the current WS size.
139 E : PROCESS_MEMORY_COUNTERS counters = {};
140 E : if (!::GetProcessMemoryInfo(process, &counters, sizeof(counters))) {
141 i : DWORD err = ::GetLastError();
142 i : LOG(ERROR) << "Unable to get process memory info: " << com::LogWe(err);
143 i : return false;
144 : }
145 :
146 E : scoped_ptr<PSAPI_WORKING_SET_INFORMATION> buffer;
147 E : DWORD number_of_entries = counters.WorkingSetSize / kPageSize;
148 E : int retries = 5;
149 : for (;;) {
150 : DWORD buffer_size = sizeof(PSAPI_WORKING_SET_INFORMATION) +
151 E : (number_of_entries * sizeof(PSAPI_WORKING_SET_BLOCK));
152 :
153 : // If we can't expand the buffer, don't leak the previous
154 : // contents or pass a NULL pointer to QueryWorkingSet.
155 : buffer.reset(reinterpret_cast<PSAPI_WORKING_SET_INFORMATION*>(
156 E : new char[buffer_size]));
157 E : if (!buffer.get()) {
158 i : LOG(ERROR) << "Unable to allocate working set buffer.";
159 i : return false;
160 : }
161 : // Zero the buffer as Gary Nebbet warns that undefined bits may not be set
162 : // in the Windows NT/2000 Native API Reference.
163 E : memset(buffer.get(), 0, buffer_size);
164 :
165 : // Call the function once to get number of items.
166 E : if (::QueryWorkingSet(process, buffer.get(), buffer_size))
167 E : break;
168 :
169 E : if (::GetLastError() != ERROR_BAD_LENGTH) {
170 i : return false;
171 : }
172 :
173 E : number_of_entries = static_cast<DWORD>(buffer->NumberOfEntries);
174 :
175 : // Maybe some entries are being added right now. Increase the buffer to
176 : // take that into account.
177 E : number_of_entries = static_cast<DWORD>(number_of_entries * 1.25);
178 :
179 E : if (--retries == 0) {
180 i : LOG(ERROR) << "Out of retries to query working set.";
181 i : return false;
182 : }
183 E : }
184 :
185 E : working_set->swap(buffer);
186 E : return true;
187 E : }
188 :
189 : bool ProcessWorkingSet::CaptureModules(DWORD process_id,
190 E : ModuleAddressSpace* modules) {
191 E : DCHECK(modules != NULL);
192 :
193 : base::win::ScopedHandle snap(
194 E : ::CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, process_id));
195 E : if (!snap.IsValid()) {
196 i : DWORD err = ::GetLastError();
197 i : LOG(ERROR) << "CreateToolhelp32Snapshot failed: " << com::LogWe(err);
198 i : return false;
199 : }
200 :
201 E : MODULEENTRY32 module = { sizeof(module) };
202 E : if (!::Module32First(snap.Get(), &module)) {
203 i : DWORD err = ::GetLastError();
204 i : LOG(ERROR) << "Module32First failed: " << com::LogWe(err);
205 i : return false;
206 : }
207 :
208 : do {
209 : ModuleAddressSpace::Range range(
210 E : reinterpret_cast<size_t>(module.modBaseAddr), module.modBaseSize);
211 E : if (!modules->Insert(range, module.szExePath)) {
212 i : LOG(ERROR) << "Module insertion failed, overlapping modules?";
213 i : return false;
214 : }
215 E : } while (::Module32Next(snap.Get(), &module));
216 :
217 E : DWORD err = ::GetLastError();
218 E : if (err != ERROR_NO_MORE_FILES) {
219 i : LOG(ERROR) << "Module32Next failed: " << com::LogWe(err);
220 i : return false;
221 : }
222 :
223 E : return true;
224 E : }
225 :
226 : } // namespace wsdump
|