1 // Copyright 2016 The Chromium Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #include "stdafx.h"
6
7 #include <algorithm>
8
9 #include "system_information_sampler.h"
10
11 // From ntdef.h
12 typedef struct _UNICODE_STRING {
13 USHORT Length;
14 USHORT MaximumLength;
15 PWCH Buffer;
16 } UNICODE_STRING;
17
18 // From <wdm.h>
19 typedef LONG KPRIORITY;
20 typedef LONG KWAIT_REASON; // Full definition is in wdm.h
21
22 // From ntddk.h
23 typedef struct _VM_COUNTERS {
24 SIZE_T PeakVirtualSize;
25 SIZE_T VirtualSize;
26 ULONG PageFaultCount;
27 // Padding here in 64-bit
28 SIZE_T PeakWorkingSetSize;
29 SIZE_T WorkingSetSize;
30 SIZE_T QuotaPeakPagedPoolUsage;
31 SIZE_T QuotaPagedPoolUsage;
32 SIZE_T QuotaPeakNonPagedPoolUsage;
33 SIZE_T QuotaNonPagedPoolUsage;
34 SIZE_T PagefileUsage;
35 SIZE_T PeakPagefileUsage;
36 } VM_COUNTERS;
37
38 // Two possibilities available from here:
39 // http://stackoverflow.com/questions/28858849/where-is-system-information-class-defined
40
41 typedef enum _SYSTEM_INFORMATION_CLASS {
42 SystemBasicInformation = 0,
43 SystemPerformanceInformation = 2,
44 SystemTimeOfDayInformation = 3,
45 SystemProcessInformation = 5, // This is the number that we need
46 SystemProcessorPerformanceInformation = 8,
47 SystemInterruptInformation = 23,
48 SystemExceptionInformation = 33,
49 SystemRegistryQuotaInformation = 37,
50 SystemLookasideInformation = 45
51 } SYSTEM_INFORMATION_CLASS;
52
53 // https://msdn.microsoft.com/en-us/library/gg750647.aspx?f=255&MSPPError=-2147217396
54 typedef struct {
55 HANDLE UniqueProcess; // Actually process ID
56 HANDLE UniqueThread; // Actually thread ID
57 } CLIENT_ID;
58
59 // From http://alax.info/blog/1182, with corrections and modifications
60 // Originally from
61 // http://undocumented.ntinternals.net/index.html?page=UserMode%2FUndocumented%20Functions%2FSystem%20Information%2FStructures%2FSYSTEM_THREAD.html
62 struct SYSTEM_THREAD_INFORMATION {
63 ULONGLONG KernelTime;
64 ULONGLONG UserTime;
65 ULONGLONG CreateTime;
66 ULONG WaitTime;
67 // Padding here in 64-bit
68 PVOID StartAddress;
69 CLIENT_ID ClientId;
70 KPRIORITY Priority;
71 LONG BasePriority;
72 ULONG ContextSwitchCount;
73 ULONG State;
74 KWAIT_REASON WaitReason;
75 };
76 #if _M_X64
77 static_assert(sizeof(SYSTEM_THREAD_INFORMATION) == 80,
78 "Structure size mismatch");
79 #else
80 static_assert(sizeof(SYSTEM_THREAD_INFORMATION) == 64,
81 "Structure size mismatch");
82 #endif
83
84 // From http://alax.info/blog/1182, with corrections and modifications
85 // Originally from
86 // http://undocumented.ntinternals.net/index.html?page=UserMode%2FUndocumented%20Functions%2FSystem%20Information%2FStructures%2FSYSTEM_THREAD.html
87 struct SYSTEM_PROCESS_INFORMATION {
88 ULONG NextEntryOffset;
89 ULONG NumberOfThreads;
90 // http://processhacker.sourceforge.net/doc/struct___s_y_s_t_e_m___p_r_o_c_e_s_s___i_n_f_o_r_m_a_t_i_o_n.html
91 ULONGLONG WorkingSetPrivateSize;
92 ULONG HardFaultCount;
93 ULONG Reserved1;
94 ULONGLONG CycleTime;
95 ULONGLONG CreateTime;
96 ULONGLONG UserTime;
97 ULONGLONG KernelTime;
98 UNICODE_STRING ImageName;
99 KPRIORITY BasePriority;
100 HANDLE ProcessId;
101 HANDLE ParentProcessId;
102 ULONG HandleCount;
103 ULONG Reserved2[2];
104 // Padding here in 64-bit
105 VM_COUNTERS VirtualMemoryCounters;
106 size_t Reserved3;
107 IO_COUNTERS IoCounters;
108 SYSTEM_THREAD_INFORMATION Threads[1];
109 };
110 #if _M_X64
111 static_assert(sizeof(SYSTEM_PROCESS_INFORMATION) == 336,
112 "Structure size mismatch");
113 #else
114 static_assert(sizeof(SYSTEM_PROCESS_INFORMATION) == 248,
115 "Structure size mismatch");
116 #endif
117
118 // ntstatus.h conflicts with windows.h so define this locally.
119 #define STATUS_SUCCESS ((NTSTATUS)0x00000000L)
120 #define STATUS_BUFFER_TOO_SMALL ((NTSTATUS)0xC0000023L)
121 #define STATUS_INFO_LENGTH_MISMATCH ((NTSTATUS)0xC0000004L)
122
123 typedef NTSTATUS(WINAPI* NTQUERYSYSTEMINFORMATION)(
124 SYSTEM_INFORMATION_CLASS SystemInformationClass,
125 PVOID SystemInformation,
126 ULONG SystemInformationLength,
127 PULONG ReturnLength);
128
oops(const char * pMessage)129 __declspec(noreturn) void oops(const char* pMessage) {
130 printf("%s\n", pMessage);
131 exit(0);
132 }
133
134 // Simple memory buffer wrapper for passing the data out of
135 // QuerySystemProcessInformation.
136 class ByteBuffer {
137 public:
ByteBuffer(size_t capacity)138 explicit ByteBuffer(size_t capacity) : size_(0), capacity_(0) {
139 if (capacity > 0)
140 grow(capacity);
141 }
142
~ByteBuffer()143 ~ByteBuffer() {}
144
data()145 BYTE* data() { return data_.get(); }
146
size()147 size_t size() { return size_; }
148
set_size(size_t new_size)149 void set_size(size_t new_size) { size_ = new_size; }
150
capacity()151 size_t capacity() { return capacity_; }
152
grow(size_t new_capacity)153 void grow(size_t new_capacity) {
154 capacity_ = new_capacity;
155 data_.reset(new BYTE[new_capacity]);
156 }
157
158 private:
159 std::unique_ptr<BYTE[]> data_;
160 size_t size_;
161 size_t capacity_;
162
163 ByteBuffer& operator=(const ByteBuffer&) = delete;
164 ByteBuffer(const ByteBuffer&) = delete;
165 };
166
167 // Wrapper for NtQuerySystemProcessInformation with buffer reallocation logic.
QuerySystemProcessInformation(ByteBuffer * buffer)168 bool QuerySystemProcessInformation(ByteBuffer* buffer) {
169 typedef NTSTATUS(WINAPI * NTQUERYSYSTEMINFORMATION)(
170 SYSTEM_INFORMATION_CLASS SystemInformationClass, PVOID SystemInformation,
171 ULONG SystemInformationLength, PULONG ReturnLength);
172
173 HMODULE ntdll = GetModuleHandle(L"ntdll.dll");
174 if (!ntdll) {
175 oops("Couldn't load ntdll.dll");
176 }
177
178 NTQUERYSYSTEMINFORMATION nt_query_system_information_ptr =
179 reinterpret_cast<NTQUERYSYSTEMINFORMATION>(
180 GetProcAddress(ntdll, "NtQuerySystemInformation"));
181 if (!nt_query_system_information_ptr)
182 oops("Couldn't find NtQuerySystemInformation");
183
184 NTSTATUS result;
185
186 // There is a potential race condition between growing the buffer and new
187 // processes being created. Try a few times before giving up.
188 for (int i = 0; i < 10; i++) {
189 ULONG data_size = 0;
190 ULONG buffer_size = static_cast<ULONG>(buffer->capacity());
191 result = nt_query_system_information_ptr(
192 SystemProcessInformation, buffer->data(), buffer_size, &data_size);
193
194 if (result == STATUS_SUCCESS) {
195 buffer->set_size(data_size);
196 break;
197 }
198
199 if (result == STATUS_INFO_LENGTH_MISMATCH ||
200 result == STATUS_BUFFER_TOO_SMALL) {
201 // Insufficient buffer. Grow to the returned |data_size| plus 10% extra
202 // to avoid frequent reallocations and try again.
203 buffer->grow(static_cast<ULONG>(data_size * 1.1));
204 } else {
205 // An error other than the two above.
206 break;
207 }
208 }
209
210 return result == STATUS_SUCCESS;
211 }
212
SystemInformationSampler(const wchar_t * process_name)213 SystemInformationSampler::SystemInformationSampler(
214 const wchar_t* process_name) {
215 lstrcpyn(target_process_name_, process_name,
216 sizeof(target_process_name_) / sizeof(wchar_t));
217
218 QueryPerformanceFrequency(&perf_frequency_);
219 QueryPerformanceCounter(&initial_counter_);
220 }
221
~SystemInformationSampler()222 SystemInformationSampler::~SystemInformationSampler() {}
223
TakeSnapshot()224 std::unique_ptr<ProcessDataSnapshot> SystemInformationSampler::TakeSnapshot() {
225 // Preallocate the buffer with the size determined on the previous call to
226 // QuerySystemProcessInformation. This should be sufficient most of the time.
227 // QuerySystemProcessInformation will grow the buffer if necessary.
228 ByteBuffer data_buffer(previous_buffer_size_);
229
230 if (!QuerySystemProcessInformation(&data_buffer))
231 return std::unique_ptr<ProcessDataSnapshot>();
232
233 previous_buffer_size_ = data_buffer.capacity();
234
235 std::unique_ptr<ProcessDataSnapshot> snapshot(new ProcessDataSnapshot);
236
237 LARGE_INTEGER perf_counter_value;
238 QueryPerformanceCounter(&perf_counter_value);
239 snapshot->timestamp = static_cast<double>(
240 (perf_counter_value.QuadPart - initial_counter_.QuadPart) /
241 perf_frequency_.QuadPart);
242
243 for (size_t offset = 0; offset < data_buffer.size();) {
244 auto pi = reinterpret_cast<const SYSTEM_PROCESS_INFORMATION*>(
245 data_buffer.data() + offset);
246
247 // Validate that the offset is valid and all needed data is within
248 // the buffer boundary.
249 if (offset + sizeof(SYSTEM_PROCESS_INFORMATION) > data_buffer.size())
250 break;
251 if (offset + sizeof(SYSTEM_PROCESS_INFORMATION) +
252 (pi->NumberOfThreads - 1) * sizeof(SYSTEM_THREAD_INFORMATION) >
253 data_buffer.size())
254 break;
255
256 if (pi->ImageName.Buffer) {
257 // Validate that the image name is within the buffer boundary.
258 // ImageName.Length seems to be in bytes rather than characters.
259 size_t image_name_offset =
260 reinterpret_cast<BYTE*>(pi->ImageName.Buffer) - data_buffer.data();
261 if (image_name_offset + pi->ImageName.Length > data_buffer.size())
262 break;
263
264 // Check if this is a chrome process. Ignore all other processes.
265 if (wcsncmp(target_process_name_filter(), pi->ImageName.Buffer,
266 lstrlen(target_process_name_filter())) == 0) {
267 // Collect enough data to be able to do a diff between two snapshots.
268 // Some threads might stop or new threads might be created between two
269 // snapshots. If a thread with a large number of context switches gets
270 // terminated the total number of context switches for the process might
271 // go down and the delta would be negative.
272 // To avoid that we need to compare thread IDs between two snapshots and
273 // not count context switches for threads that are missing in the most
274 // recent snapshot.
275 ProcessData process_data;
276
277 process_data.cpu_time = pi->KernelTime + pi->UserTime;
278 process_data.working_set = pi->WorkingSetPrivateSize;
279
280 // Iterate over threads and store each thread's ID and number of context
281 // switches.
282 for (ULONG thread_index = 0; thread_index < pi->NumberOfThreads;
283 ++thread_index) {
284 const SYSTEM_THREAD_INFORMATION* ti = &pi->Threads[thread_index];
285 if (ti->ClientId.UniqueProcess != pi->ProcessId)
286 continue;
287
288 ThreadData thread_data;
289 thread_data.thread_id = ti->ClientId.UniqueThread;
290 thread_data.context_switches = ti->ContextSwitchCount;
291 process_data.threads.push_back(thread_data);
292 }
293
294 // Order thread data by thread ID to help diff two snapshots.
295 std::sort(process_data.threads.begin(), process_data.threads.end(),
296 [](const ThreadData& l, const ThreadData r) {
297 return l.thread_id < r.thread_id;
298 });
299
300 snapshot->processes.insert(
301 std::make_pair(pi->ProcessId, std::move(process_data)));
302 }
303 }
304
305 // Check for end of the list.
306 if (!pi->NextEntryOffset)
307 break;
308
309 // Jump to the next entry.
310 offset += pi->NextEntryOffset;
311 }
312
313 return snapshot;
314 }
315