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