OLD | NEW |
---|---|
(Empty) | |
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 "chrome_elf/chrome_elf_reg.h" | |
6 | |
7 #include "chrome_elf/chrome_elf_util.h" // IsSystemInstall() | |
8 | |
9 namespace { | |
10 | |
11 // function pointers used for registry access. | |
12 RtlInitUnicodeStringFunction rtl_init_unicode_string = nullptr; | |
13 NtCreateKeyFunction nt_create_key = nullptr; | |
14 NtDeleteKeyFunction nt_delete_key = nullptr; | |
15 NtOpenKeyExFunction nt_open_key_ex = nullptr; | |
16 NtCloseFunction nt_close = nullptr; | |
17 NtQueryValueKeyFunction nt_query_value_key = nullptr; | |
18 NtSetValueKeyFunction nt_set_value_key = nullptr; | |
19 | |
20 // lazy init. No concern about concurrency in chrome_elf. | |
21 bool initialized = false; | |
22 bool system_install = false; | |
23 const size_t kRegMaxPathLen = 255; | |
24 wchar_t kRegPathHKLM[] = L"\\Registry\\Machine\\"; | |
25 wchar_t kRegPathHKCU[kRegMaxPathLen] = L""; | |
26 base::string16 current_user_sid_string; | |
27 base::string16 override_path; | |
robertshield
2016/04/20 16:55:53
These are globals, please label them as g_override
penny
2016/05/28 01:34:23
Done.
| |
28 | |
29 bool InitNativeRegApi() { | |
30 HMODULE ntdll = ::GetModuleHandleW(L"ntdll.dll"); | |
31 | |
32 // Setup the global function pointers for registry access. | |
33 rtl_init_unicode_string = reinterpret_cast<RtlInitUnicodeStringFunction>( | |
34 ::GetProcAddress(ntdll, "RtlInitUnicodeString")); | |
35 | |
36 nt_create_key = reinterpret_cast<NtCreateKeyFunction>( | |
37 ::GetProcAddress(ntdll, "NtCreateKey")); | |
38 | |
39 nt_delete_key = reinterpret_cast<NtDeleteKeyFunction>( | |
40 ::GetProcAddress(ntdll, "NtDeleteKey")); | |
41 | |
42 nt_open_key_ex = reinterpret_cast<NtOpenKeyExFunction>( | |
43 ::GetProcAddress(ntdll, "NtOpenKeyEx")); | |
44 | |
45 nt_close = | |
46 reinterpret_cast<NtCloseFunction>(::GetProcAddress(ntdll, "NtClose")); | |
47 | |
48 nt_query_value_key = reinterpret_cast<NtQueryValueKeyFunction>( | |
49 ::GetProcAddress(ntdll, "NtQueryValueKey")); | |
50 | |
51 nt_set_value_key = reinterpret_cast<NtSetValueKeyFunction>( | |
52 ::GetProcAddress(ntdll, "NtSetValueKey")); | |
53 | |
54 if (!rtl_init_unicode_string || !nt_create_key || !nt_open_key_ex || | |
55 !nt_delete_key || !nt_close || !nt_query_value_key || !nt_set_value_key) | |
56 return false; | |
57 | |
58 // We need to set HKCU based on the sid of the current user account. | |
59 RtlFormatCurrentUserKeyPathFunction rtl_current_user_string = | |
60 reinterpret_cast<RtlFormatCurrentUserKeyPathFunction>( | |
61 ::GetProcAddress(ntdll, "RtlFormatCurrentUserKeyPath")); | |
62 | |
63 RtlFreeUnicodeStringFunction rtl_free_unicode_str = | |
64 reinterpret_cast<RtlFreeUnicodeStringFunction>( | |
65 ::GetProcAddress(ntdll, "RtlFreeUnicodeString")); | |
66 | |
67 if (!rtl_current_user_string || !rtl_free_unicode_str) | |
68 return false; | |
69 | |
70 UNICODE_STRING current_user_reg_path; | |
71 if (!NT_SUCCESS(rtl_current_user_string(¤t_user_reg_path))) | |
72 return false; | |
73 | |
74 // Finish setting up global HKCU path. | |
75 ::wcsncat(kRegPathHKCU, current_user_reg_path.Buffer, kRegMaxPathLen - 1); | |
76 ::wcsncat(kRegPathHKCU, L"\\", (kRegMaxPathLen - ::wcslen(kRegPathHKCU) - 1)); | |
77 // Keep the sid string as well. | |
78 wchar_t* ptr = ::wcsrchr(current_user_reg_path.Buffer, L'\\'); | |
79 ptr++; | |
80 current_user_sid_string.assign(ptr); | |
81 rtl_free_unicode_str(¤t_user_reg_path); | |
82 | |
83 // Figure out if we're a system or user install. | |
84 system_install = IsSystemInstall(GetCommandLineW()); | |
85 | |
86 initialized = true; | |
87 return true; | |
88 } | |
89 | |
90 const wchar_t* ConvertRootKey(nt::ROOT_KEY root) { | |
91 nt::ROOT_KEY key = root; | |
92 | |
93 if (!root) { | |
94 // AUTO | |
95 key = system_install ? nt::HKLM : nt::HKCU; | |
96 } | |
97 | |
98 if ((key == nt::HKCU) && (!nt::HKCU_override.empty())) { | |
99 override_path.assign(kRegPathHKCU); | |
100 override_path.append(nt::HKCU_override.c_str()); | |
101 override_path.append(L"\\"); | |
102 return override_path.c_str(); | |
103 } else if ((key == nt::HKLM) && (!nt::HKLM_override.empty())) { | |
104 override_path.assign(kRegPathHKCU); | |
105 override_path.append(nt::HKLM_override.c_str()); | |
106 override_path.append(L"\\"); | |
107 return override_path.c_str(); | |
108 } | |
109 | |
110 if (key == nt::HKCU) | |
111 return kRegPathHKCU; | |
112 else | |
113 return kRegPathHKLM; | |
114 } | |
115 | |
116 } // namespace | |
117 | |
118 namespace nt { | |
119 | |
120 base::string16 HKLM_override; | |
121 base::string16 HKCU_override; | |
122 | |
123 bool CreateRegKey(ROOT_KEY root, | |
124 const wchar_t* key_path, | |
125 ACCESS_MASK access, | |
126 HANDLE* out_handle OPTIONAL) { | |
127 if (!initialized) | |
128 InitNativeRegApi(); | |
129 | |
130 NTSTATUS status = STATUS_UNSUCCESSFUL; | |
131 UNICODE_STRING key_path_uni = {}; | |
132 OBJECT_ATTRIBUTES obj = {}; | |
133 HANDLE key_handle = INVALID_HANDLE_VALUE; | |
134 | |
135 base::string16 full_path(ConvertRootKey(root)); | |
136 full_path.append(key_path); | |
137 | |
138 rtl_init_unicode_string(&key_path_uni, full_path.c_str()); | |
139 InitializeObjectAttributes(&obj, &key_path_uni, OBJ_CASE_INSENSITIVE, NULL, | |
140 NULL); | |
141 // Assuming CreateOptions to preserve across reboot, and currently | |
142 // don't care whether already existed or not. | |
143 status = nt_create_key(&key_handle, access, &obj, 0, nullptr, | |
144 REG_OPTION_NON_VOLATILE, nullptr); | |
145 | |
146 if (NT_SUCCESS(status)) { | |
147 if (out_handle) | |
148 *out_handle = key_handle; | |
149 else | |
150 CloseRegKey(key_handle); | |
151 return true; | |
152 } | |
153 | |
154 return false; | |
155 } | |
156 | |
157 bool DeleteRegKey(HANDLE key) { | |
158 if (!initialized) | |
159 InitNativeRegApi(); | |
160 | |
161 NTSTATUS status = STATUS_UNSUCCESSFUL; | |
162 | |
163 status = nt_delete_key(key); | |
164 | |
165 if (NT_SUCCESS(status)) | |
166 return true; | |
167 | |
168 return false; | |
169 } | |
170 | |
171 // wrapper function | |
172 bool DeleteRegKey(ROOT_KEY root, const wchar_t* key_path) { | |
173 HANDLE key = INVALID_HANDLE_VALUE; | |
174 | |
175 if (!OpenRegKey(root, key_path, DELETE, &key, nullptr)) | |
176 return false; | |
177 | |
178 if (!DeleteRegKey(key)) { | |
179 CloseRegKey(key); | |
180 return false; | |
181 } | |
182 | |
183 CloseRegKey(key); | |
184 return true; | |
185 } | |
186 | |
187 bool OpenRegKey(ROOT_KEY root, | |
188 const wchar_t* key_path, | |
189 ACCESS_MASK access, | |
190 HANDLE* out_handle, | |
191 NTSTATUS* error_code OPTIONAL) { | |
192 if (!initialized) | |
193 InitNativeRegApi(); | |
194 | |
195 NTSTATUS status = STATUS_UNSUCCESSFUL; | |
196 UNICODE_STRING key_path_uni = {}; | |
197 OBJECT_ATTRIBUTES obj = {}; | |
198 *out_handle = INVALID_HANDLE_VALUE; | |
199 | |
200 base::string16 full_path(ConvertRootKey(root)); | |
201 full_path.append(key_path); | |
202 | |
203 rtl_init_unicode_string(&key_path_uni, full_path.c_str()); | |
204 InitializeObjectAttributes(&obj, &key_path_uni, OBJ_CASE_INSENSITIVE, NULL, | |
205 NULL); | |
206 | |
207 status = nt_open_key_ex(out_handle, access, &obj, 0); | |
208 // See if caller wants the NTSTATUS. | |
209 if (error_code) | |
210 *error_code = status; | |
211 | |
212 if (NT_SUCCESS(status)) | |
213 return true; | |
214 | |
215 return false; | |
216 } | |
217 | |
218 void CloseRegKey(HANDLE key) { | |
219 if (!initialized) | |
220 InitNativeRegApi(); | |
221 nt_close(key); | |
222 } | |
223 | |
224 //------------------------------------------------------------------------------ | |
225 // Getter functions | |
226 //------------------------------------------------------------------------------ | |
227 | |
228 bool QueryRegKeyValue(HANDLE key, | |
229 const wchar_t* value_name, | |
230 ULONG* out_type, | |
231 BYTE** out_buffer, | |
232 DWORD* out_size) { | |
233 if (!initialized) | |
234 InitNativeRegApi(); | |
235 | |
236 NTSTATUS ntstatus = STATUS_UNSUCCESSFUL; | |
237 UNICODE_STRING value_uni = {}; | |
238 rtl_init_unicode_string(&value_uni, value_name); | |
239 DWORD size_needed = 0; | |
240 bool success = false; | |
241 | |
242 // First call to find out how much room we need for the value! | |
243 ntstatus = nt_query_value_key(key, &value_uni, KeyValueFullInformation, | |
244 nullptr, 0, &size_needed); | |
245 if (ntstatus != STATUS_BUFFER_TOO_SMALL) | |
246 return false; | |
247 | |
248 KEY_VALUE_FULL_INFORMATION* value_info = | |
249 reinterpret_cast<KEY_VALUE_FULL_INFORMATION*>(new BYTE[size_needed]); | |
250 | |
251 // Second call to get the value. | |
252 ntstatus = nt_query_value_key(key, &value_uni, KeyValueFullInformation, | |
253 value_info, size_needed, &size_needed); | |
254 if (NT_SUCCESS(ntstatus)) { | |
255 *out_type = value_info->Type; | |
256 *out_size = value_info->DataLength; | |
257 *out_buffer = new BYTE[*out_size]; | |
258 ::memcpy(*out_buffer, | |
259 (reinterpret_cast<BYTE*>(value_info) + value_info->DataOffset), | |
260 *out_size); | |
261 success = true; | |
262 } | |
263 | |
264 delete[] value_info; | |
265 return success; | |
266 } | |
267 | |
268 // wrapper function | |
269 bool GetRegValue_DWORD(HANDLE key, | |
270 const wchar_t* value_name, | |
271 DWORD* out_dword) { | |
272 ULONG type = REG_NONE; | |
273 BYTE* value_bytes = nullptr; | |
274 DWORD ret_size = 0; | |
275 | |
276 if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) || | |
277 type != REG_DWORD) | |
278 return false; | |
279 | |
280 *out_dword = *(reinterpret_cast<DWORD*>(value_bytes)); | |
281 | |
282 delete[] value_bytes; | |
283 return true; | |
284 } | |
285 | |
286 // wrapper function | |
287 bool GetRegValue_DWORD(ROOT_KEY root, | |
288 const wchar_t* key_path, | |
289 const wchar_t* value_name, | |
290 DWORD* out_dword) { | |
291 HANDLE key = INVALID_HANDLE_VALUE; | |
292 | |
293 if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | KEY_WOW64_32KEY, &key, | |
294 NULL)) | |
295 return false; | |
296 | |
297 if (!GetRegValue_DWORD(key, value_name, out_dword)) { | |
298 CloseRegKey(key); | |
299 return false; | |
300 } | |
301 | |
302 CloseRegKey(key); | |
303 return true; | |
304 } | |
305 | |
306 // wrapper function | |
307 bool GetRegValue_SZ(HANDLE key, | |
308 const wchar_t* value_name, | |
309 base::string16* out_sz) { | |
310 BYTE* value_bytes = nullptr; | |
311 DWORD ret_size = 0; | |
312 ULONG type = REG_NONE; | |
313 | |
314 if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) || | |
315 type != REG_SZ) | |
316 return false; | |
317 | |
318 *out_sz = reinterpret_cast<wchar_t*>(value_bytes); | |
319 | |
320 delete[] value_bytes; | |
321 return true; | |
322 } | |
323 | |
324 // wrapper function | |
325 bool GetRegValue_SZ(ROOT_KEY root, | |
326 const wchar_t* key_path, | |
327 const wchar_t* value_name, | |
328 base::string16* out_sz) { | |
329 HANDLE key = INVALID_HANDLE_VALUE; | |
330 | |
331 if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | KEY_WOW64_32KEY, &key, | |
332 NULL)) | |
333 return false; | |
334 | |
335 if (!GetRegValue_SZ(key, value_name, out_sz)) { | |
336 CloseRegKey(key); | |
337 return false; | |
338 } | |
339 | |
340 CloseRegKey(key); | |
341 return true; | |
342 } | |
343 | |
344 // wrapper function | |
345 bool GetRegValue_MULTI_SZ(HANDLE key, | |
346 const wchar_t* value_name, | |
347 std::vector<base::string16>* out_multi_sz) { | |
348 BYTE* value_bytes = nullptr; | |
349 DWORD ret_size = 0; | |
350 ULONG type = REG_NONE; | |
351 | |
352 if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) || | |
353 type != REG_MULTI_SZ) | |
354 return false; | |
355 | |
356 // Make sure the vector is empty to start. | |
357 (*out_multi_sz).resize(0); | |
358 | |
359 wchar_t* pointer = reinterpret_cast<wchar_t*>(value_bytes); | |
360 base::string16 temp = pointer; | |
361 // Loop. Each string is separated by '\0'. Another '\0' at very end (so 2 in | |
362 // a row). | |
363 while (temp.length() != 0) { | |
364 (*out_multi_sz).push_back(temp); | |
365 | |
366 pointer += temp.length() + 1; | |
367 temp = pointer; | |
368 } | |
369 | |
370 // Handle the case of "empty multi_sz". | |
371 if (out_multi_sz->size() == 0) | |
372 out_multi_sz->push_back(L""); | |
373 | |
374 delete[] value_bytes; | |
375 return true; | |
376 } | |
377 | |
378 // wrapper function | |
379 bool GetRegValue_MULTI_SZ(ROOT_KEY root, | |
380 const wchar_t* key_path, | |
381 const wchar_t* value_name, | |
382 std::vector<base::string16>* out_multi_sz) { | |
383 HANDLE key = INVALID_HANDLE_VALUE; | |
384 | |
385 if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | KEY_WOW64_32KEY, &key, | |
386 NULL)) | |
387 return false; | |
388 | |
389 if (!GetRegValue_MULTI_SZ(key, value_name, out_multi_sz)) { | |
390 CloseRegKey(key); | |
391 return false; | |
392 } | |
393 | |
394 CloseRegKey(key); | |
395 return true; | |
396 } | |
397 | |
398 //------------------------------------------------------------------------------ | |
399 // Setter functions | |
400 //------------------------------------------------------------------------------ | |
401 | |
402 bool SetRegKeyValue(HANDLE key, | |
403 const wchar_t* value_name, | |
404 ULONG type, | |
405 BYTE* data, | |
406 DWORD data_size) { | |
407 if (!initialized) | |
408 InitNativeRegApi(); | |
409 | |
410 NTSTATUS ntstatus = STATUS_UNSUCCESSFUL; | |
411 UNICODE_STRING value_uni = {}; | |
412 rtl_init_unicode_string(&value_uni, value_name); | |
413 | |
414 ntstatus = nt_set_value_key(key, &value_uni, 0, type, data, data_size); | |
415 | |
416 if (NT_SUCCESS(ntstatus)) | |
417 return true; | |
418 | |
419 return false; | |
420 } | |
421 | |
422 // wrapper function | |
423 bool SetRegValue_DWORD(HANDLE key, const wchar_t* value_name, DWORD value) { | |
424 return SetRegKeyValue(key, value_name, REG_DWORD, | |
425 reinterpret_cast<BYTE*>(&value), sizeof(value)); | |
426 } | |
427 | |
428 // wrapper function | |
429 bool SetRegValue_DWORD(ROOT_KEY root, | |
430 const wchar_t* key_path, | |
431 const wchar_t* value_name, | |
432 DWORD value) { | |
433 HANDLE key = INVALID_HANDLE_VALUE; | |
434 | |
435 if (!OpenRegKey(root, key_path, KEY_SET_VALUE | KEY_WOW64_32KEY, &key, NULL)) | |
436 return false; | |
437 | |
438 if (!SetRegValue_DWORD(key, value_name, value)) { | |
439 CloseRegKey(key); | |
440 return false; | |
441 } | |
442 | |
443 return true; | |
444 } | |
445 | |
446 // wrapper function | |
447 bool SetRegValue_SZ(HANDLE key, | |
448 const wchar_t* value_name, | |
449 base::string16* value) { | |
450 wchar_t* string = const_cast<wchar_t*>(value->c_str()); | |
451 if (std::numeric_limits<DWORD>::max() < ::wcslen(string)) | |
452 return false; | |
453 | |
454 DWORD length = static_cast<DWORD>(::wcslen(string)); | |
455 | |
456 return SetRegKeyValue(key, value_name, REG_SZ, | |
457 reinterpret_cast<BYTE*>(string), | |
458 (length + 1) * sizeof(wchar_t)); | |
459 } | |
460 | |
461 // wrapper function | |
462 bool SetRegValue_SZ(ROOT_KEY root, | |
463 const wchar_t* key_path, | |
464 const wchar_t* value_name, | |
465 base::string16* value) { | |
466 HANDLE key = INVALID_HANDLE_VALUE; | |
467 | |
468 if (!OpenRegKey(root, key_path, KEY_SET_VALUE | KEY_WOW64_32KEY, &key, NULL)) | |
469 return false; | |
470 | |
471 if (!SetRegValue_SZ(key, value_name, value)) { | |
472 CloseRegKey(key); | |
473 return false; | |
474 } | |
475 | |
476 return true; | |
477 } | |
478 | |
479 // wrapper function | |
480 bool SetRegValue_MULTI_SZ(HANDLE key, | |
481 const wchar_t* value_name, | |
482 std::vector<base::string16>* values) { | |
483 std::vector<wchar_t> builder; | |
484 | |
485 for (auto& string : *values) { | |
486 // Just in case someone is passing in an illegal empty string | |
487 // (not allowed in REG_MULTI_SZ), ignore it. | |
488 if (!string.empty()) { | |
489 for (wchar_t& w : string) { | |
490 builder.push_back(w); | |
491 } | |
492 builder.push_back(L'\0'); | |
493 } | |
494 } | |
495 // Add second null terminator to end REG_MULTI_SZ. | |
496 builder.push_back(L'\0'); | |
497 // Handle rare case where the vector passed in was empty, | |
498 // or only had an empty string. | |
499 if (builder.size() == 1) | |
500 builder.push_back(L'\0'); | |
501 | |
502 if (std::numeric_limits<DWORD>::max() < builder.size()) | |
503 return false; | |
504 | |
505 return SetRegKeyValue( | |
506 key, value_name, REG_MULTI_SZ, reinterpret_cast<BYTE*>(builder.data()), | |
507 (static_cast<DWORD>(builder.size()) + 1) * sizeof(wchar_t)); | |
508 } | |
509 | |
510 // wrapper function | |
511 bool SetRegValue_MULTI_SZ(ROOT_KEY root, | |
512 const wchar_t* key_path, | |
513 const wchar_t* value_name, | |
514 std::vector<base::string16>* values) { | |
515 HANDLE key = INVALID_HANDLE_VALUE; | |
516 | |
517 if (!OpenRegKey(root, key_path, KEY_SET_VALUE | KEY_WOW64_32KEY, &key, NULL)) | |
518 return false; | |
519 | |
520 if (!SetRegValue_MULTI_SZ(key, value_name, values)) { | |
521 CloseRegKey(key); | |
522 return false; | |
523 } | |
524 | |
525 return true; | |
526 } | |
527 | |
528 //------------------------------------------------------------------------------ | |
529 // Utils | |
530 //------------------------------------------------------------------------------ | |
531 | |
532 base::string16 GetCurrentUserSidString() { | |
533 if (!initialized) | |
534 InitNativeRegApi(); | |
535 | |
536 return current_user_sid_string; | |
537 } | |
538 | |
539 }; // namespace nt | |
OLD | NEW |