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 <psapi.h>
18 :
19 : #include "base/stringprintf.h"
20 : #include "sawbuck/common/com_utils.h"
21 : #include "syzygy/common/align.h"
22 : #include "syzygy/pe/pe_file.h"
23 :
24 : namespace sampler {
25 :
26 : namespace {
27 :
28 : typedef SampledModuleCache::Process::ModuleMap ModuleMap;
29 :
30 : } // namespace
31 :
32 : SampledModuleCache::SampledModuleCache(size_t log2_bucket_size)
33 E : : log2_bucket_size_(log2_bucket_size) {
34 E : DCHECK_LE(2u, log2_bucket_size);
35 E : DCHECK_GE(31u, log2_bucket_size);
36 E : }
37 :
38 E : SampledModuleCache::~SampledModuleCache() {
39 : // Force a clean up of all modules (and consequently all processes).
40 E : MarkAllModulesDead();
41 E : RemoveDeadModules();
42 E : }
43 :
44 E : bool SampledModuleCache::AddModule(HANDLE process, HMODULE module) {
45 E : DCHECK(process != INVALID_HANDLE_VALUE);
46 :
47 : // Create or find the process object. We don't actually insert it into the
48 : // map until everything has succeeded, saving us the cleanup on failure.
49 E : DWORD pid = ::GetProcessId(process);
50 E : scoped_ptr<Process> scoped_proc;
51 E : Process* proc = NULL;
52 E : ProcessMap::iterator proc_it = processes_.find(pid);
53 E : if (proc_it == processes_.end()) {
54 E : HANDLE temp_handle = INVALID_HANDLE_VALUE;
55 : if (!::DuplicateHandle(::GetCurrentProcess(), process,
56 : ::GetCurrentProcess(), &temp_handle,
57 : 0, FALSE, DUPLICATE_SAME_ACCESS) ||
58 E : temp_handle == INVALID_HANDLE_VALUE) {
59 i : DWORD error = ::GetLastError();
60 i : LOG(ERROR) << "Failed to duplicate handle to process " << pid << ": "
61 : << com::LogWe(error);
62 i : return false;
63 : }
64 :
65 E : scoped_proc.reset(new Process(temp_handle, pid));
66 E : proc = scoped_proc.get();
67 E : } else {
68 E : proc = proc_it->second;
69 : }
70 E : DCHECK(proc != NULL);
71 :
72 E : if (!proc->AddModule(module, log2_bucket_size_))
73 E : return false;
74 :
75 E : if (scoped_proc.get() != NULL) {
76 : // Initialization was successful so we can safely insert the newly created
77 : // process into the map.
78 E : processes_.insert(std::make_pair(pid, proc));
79 E : scoped_proc.release();
80 : }
81 :
82 E : return true;
83 E : }
84 :
85 E : void SampledModuleCache::MarkAllModulesDead() {
86 E : for (ProcessMap::iterator proc_it = processes_.begin();
87 E : proc_it != processes_.end(); ++proc_it) {
88 E : proc_it->second->MarkDead();
89 E : }
90 E : }
91 :
92 E : void SampledModuleCache::RemoveDeadModules() {
93 E : ProcessMap::iterator proc_it = processes_.begin();
94 E : ProcessMap::iterator proc_it_next = proc_it;
95 E : ++proc_it_next;
96 :
97 E : while (proc_it != processes_.end()) {
98 : // Remove any dead modules from the process.
99 E : proc_it->second->RemoveDeadModules(dead_module_callback_);
100 :
101 : // If the process itself is dead (contains no more profiling modules) then
102 : // remove it.
103 E : if (!proc_it->second->alive()) {
104 E : Process* proc = proc_it->second;
105 E : delete proc;
106 E : processes_.erase(proc_it);
107 : }
108 :
109 E : proc_it = proc_it_next;
110 E : ++proc_it_next;
111 E : }
112 E : }
113 :
114 : SampledModuleCache::Process::Process(HANDLE process, DWORD pid)
115 E : : process_(process), pid_(pid), alive_(true) {
116 E : DCHECK(process != INVALID_HANDLE_VALUE);
117 E : }
118 :
119 E : SampledModuleCache::Process::~Process() {
120 E : MarkDead();
121 E : RemoveDeadModules(DeadModuleCallback());
122 E : }
123 :
124 : bool SampledModuleCache::Process::AddModule(HMODULE module,
125 E : size_t log2_bucket_size) {
126 E : DCHECK(module != INVALID_HANDLE_VALUE);
127 E : DCHECK_LE(2u, log2_bucket_size);
128 E : DCHECK_GE(31u, log2_bucket_size);
129 :
130 E : ModuleMap::iterator mod_it = modules_.find(module);
131 E : if (mod_it != modules_.end()) {
132 : // The module is already being profiled. Simply mark it as being alive.
133 E : mod_it->second->MarkAlive();
134 :
135 : // And mark ourselves as being alive while we're at it.
136 E : MarkAlive();
137 :
138 E : return true;
139 : }
140 :
141 : // Create a new module object. We don't actually insert it into the map until
142 : // everything has succeeded, saving us the cleanup on failure.
143 E : scoped_ptr<Module> mod(new Module(this, module, log2_bucket_size));
144 :
145 E : if (!mod->Init())
146 E : return false;
147 :
148 E : if (!mod->Start())
149 i : return false;
150 :
151 : // Initialization was successful so we can safely insert the initialized
152 : // (and currently profiling) module into the map.
153 E : modules_.insert(std::make_pair(module, mod.release()));
154 :
155 E : return true;
156 E : }
157 :
158 E : void SampledModuleCache::Process::MarkDead() {
159 : // Mark all of our children as dead, and ourselves.
160 E : alive_ = false;
161 E : for (ModuleMap::iterator it = modules_.begin(); it != modules_.end(); ++it)
162 E : it->second->MarkDead();
163 E : }
164 :
165 : void SampledModuleCache::Process::RemoveDeadModules(
166 E : DeadModuleCallback callback) {
167 E : ModuleMap::iterator mod_it = modules_.begin();
168 E : ModuleMap::iterator mod_it_next = mod_it;
169 E : ++mod_it_next;
170 :
171 E : while (mod_it != modules_.end()) {
172 E : DCHECK(mod_it->second != NULL);
173 :
174 E : if (!mod_it->second->alive()) {
175 : // Stop profiling.
176 E : mod_it->second->Stop();
177 :
178 : // Return the results to the callback if one has been provided.
179 E : if (!callback.is_null())
180 E : callback.Run(mod_it->second);
181 :
182 : // And clean things up.
183 E : Module* mod = mod_it->second;
184 E : delete mod;
185 E : modules_.erase(mod_it);
186 : }
187 :
188 E : mod_it = mod_it_next;
189 E : ++mod_it_next;
190 E : }
191 E : }
192 :
193 : SampledModuleCache::Module::Module(Process* process,
194 : HMODULE module,
195 : size_t log2_bucket_size)
196 : : process_(process),
197 : module_(module),
198 : module_size_(0),
199 : module_checksum_(0),
200 : module_time_date_stamp_(0),
201 : buckets_begin_(NULL),
202 : buckets_end_(NULL),
203 : log2_bucket_size_(log2_bucket_size),
204 : profiling_start_time_(0),
205 E : alive_(true) {
206 E : DCHECK(process != NULL);
207 E : DCHECK(module_ != INVALID_HANDLE_VALUE);
208 E : DCHECK_LE(2u, log2_bucket_size);
209 E : DCHECK_GE(31u, log2_bucket_size);
210 E : }
211 :
212 E : bool SampledModuleCache::Module::Init() {
213 : // Read the headers.
214 E : char headers[4096] = {};
215 E : size_t net_bytes_read = 0;
216 E : size_t empty_reads = 0;
217 E : while (net_bytes_read < sizeof(headers)) {
218 E : SIZE_T bytes_read = 0;
219 : if (::ReadProcessMemory(process_->process(),
220 : module_,
221 : headers + net_bytes_read,
222 : sizeof(headers) - net_bytes_read,
223 E : &bytes_read) == FALSE) {
224 E : DWORD error = ::GetLastError();
225 E : LOG(ERROR) << "ReadProcessMemory failed for module at address "
226 : << base::StringPrintf("0x%08X", module_)
227 : << " of process " << process_->pid() << ": "
228 : << com::LogWe(error);
229 E : return false;
230 : }
231 E : if (bytes_read == 0) {
232 i : if (++empty_reads == 3) {
233 i : LOG(ERROR) << "ReadProcessMemory unable to read headers for module at "
234 : << "address " << base::StringPrintf("0x%08X", module_)
235 : << " of process " << process_->pid() << ".";
236 i : return false;
237 : }
238 i : } else {
239 E : net_bytes_read += bytes_read;
240 E : empty_reads = 0;
241 : }
242 E : }
243 :
244 : const IMAGE_DOS_HEADER* dos_header =
245 E : reinterpret_cast<const IMAGE_DOS_HEADER*>(headers);
246 : COMPILE_ASSERT(sizeof(IMAGE_DOS_HEADER) <= sizeof(headers),
247 : headers_must_be_big_enough_for_DOS_headers);
248 :
249 : // Get the NT headers and make sure they're fully contained in the block we
250 : // read.
251 E : if (dos_header->e_lfanew > sizeof(headers))
252 i : return false;
253 : const IMAGE_NT_HEADERS* nt_headers =
254 E : reinterpret_cast<const IMAGE_NT_HEADERS*>(headers + dos_header->e_lfanew);
255 E : if (reinterpret_cast<const char*>(nt_headers + 1) - headers > sizeof(headers))
256 i : return false;
257 :
258 : // Get the section headers and make sure they're fully contained in the
259 : // block we read.
260 E : size_t section_count = nt_headers->FileHeader.NumberOfSections;
261 : const IMAGE_SECTION_HEADER* section_headers =
262 E : reinterpret_cast<const IMAGE_SECTION_HEADER*>(nt_headers + 1);
263 : if (reinterpret_cast<const char*>(section_headers + section_count) - headers >
264 E : sizeof(headers)) {
265 i : return false;
266 : }
267 :
268 E : module_size_ = nt_headers->OptionalHeader.SizeOfImage;
269 E : module_checksum_ = nt_headers->OptionalHeader.CheckSum;
270 E : module_time_date_stamp_ = nt_headers->FileHeader.TimeDateStamp;
271 :
272 : // Find the RVA range associated with any text segments in the module.
273 E : DWORD text_begin = ~0;
274 E : DWORD text_end = 0;
275 E : for (size_t i = 0; i < section_count; ++i) {
276 E : const IMAGE_SECTION_HEADER& sh = section_headers[i];
277 : static const DWORD kExecFlags = IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_CNT_CODE;
278 E : if ((sh.Characteristics & kExecFlags) == 0)
279 E : continue;
280 :
281 E : DWORD sec_begin = sh.VirtualAddress;
282 E : DWORD sec_end = sec_begin + sh.Misc.VirtualSize;
283 E : if (sec_begin < text_begin)
284 E : text_begin = sec_begin;
285 E : if (sec_end > text_end)
286 E : text_end = sec_end;
287 E : }
288 :
289 : // Adjust the address range for the bucket size.
290 E : DWORD bucket_size = 1 << log2_bucket_size_;
291 E : text_begin = (text_begin / bucket_size) * bucket_size;
292 E : text_end = ((text_end + bucket_size - 1) / bucket_size) * bucket_size;
293 :
294 : // Calculate the number of buckets.
295 E : DCHECK_EQ(0u, (text_end - text_begin) % bucket_size);
296 E : DWORD bucket_count = (text_end - text_begin) / bucket_size;
297 :
298 : // Calculate the bucket range in the remote address space.
299 : buckets_begin_ = reinterpret_cast<const void*>(
300 E : reinterpret_cast<const char*>(module_) + text_begin);
301 : buckets_end_ = reinterpret_cast<const void*>(
302 E : reinterpret_cast<const char*>(module_) + text_end);
303 :
304 : // Initialize the profiler.
305 : if (!profiler_.Initialize(process_->process(),
306 : const_cast<void*>(buckets_begin_),
307 : text_end - text_begin,
308 E : log2_bucket_size_)) {
309 i : LOG(ERROR) << "Failed to initialize profiler for address range "
310 : << base::StringPrintf("0x%08X - 0x%08X",
311 : buckets_begin_,
312 : buckets_end_)
313 : << " of process " << process_->pid() << ".";
314 i : return false;
315 : }
316 E : DCHECK_EQ(bucket_count, profiler_.buckets().size());
317 :
318 E : return true;
319 E : }
320 :
321 : } // namespace sampler
|