Index: chrome_elf/nt_registry/nt_registry.cc |
diff --git a/chrome_elf/nt_registry/nt_registry.cc b/chrome_elf/nt_registry/nt_registry.cc |
index 4c65af8daff7c76ea5a9e4a0255da0225aeb58a4..6e1039c707cbc53144ce0fd6d67042659e54300e 100644 |
--- a/chrome_elf/nt_registry/nt_registry.cc |
+++ b/chrome_elf/nt_registry/nt_registry.cc |
@@ -542,6 +542,47 @@ bool ParseFullRegPath(const std::wstring& converted_root, |
return true; |
} |
+// String safety. |
+// - NOTE: only working with wchar_t here. |
+// - Also ensures the content of |value_bytes| is at least a terminator. |
+// - Pass "true" for |multi| for MULTISZ. |
+void EnsureTerminatedSZ(std::vector<BYTE>* value_bytes, bool multi) { |
+ DWORD terminator_size = sizeof(wchar_t); |
+ |
+ if (multi) |
+ terminator_size = 2 * sizeof(wchar_t); |
+ |
+ // Ensure content is at least the size of a terminator. |
+ for (size_t i = value_bytes->size(); i < terminator_size; i++) |
grt (UTC plus 2)
2017/07/20 05:22:53
if (value_bytes->size() < terminator_size) {
v
penny
2017/07/20 18:38:37
Done. <sigh> Thank you for your 1337 c++ skillz
|
+ value_bytes->push_back(0); |
+ |
+ // Sanity check content size based on character size. |
+ DWORD modulo = value_bytes->size() % sizeof(wchar_t); |
+ while (modulo) { |
grt (UTC plus 2)
2017/07/20 05:22:53
value_bytes->insert(value_bytes->end(), modulo, 0)
penny
2017/07/20 18:38:37
Done.
|
+ value_bytes->push_back(0); |
+ modulo--; |
+ } |
+ |
+ // Now finally check for trailing terminator. |
+ bool terminated = true; |
+ size_t last_element = value_bytes->size() - 1; |
+ for (size_t i = 0; i < terminator_size; i++) { |
+ if ((*value_bytes)[last_element - i] != 0) { |
+ terminated = false; |
+ break; |
+ } |
+ } |
+ |
+ if (terminated) |
+ return; |
+ |
+ // Append a full terminator to be safe. |
+ for (size_t i = 0; i < terminator_size; i++) |
grt (UTC plus 2)
2017/07/20 05:22:53
value_bytes->insert(value_bytes->end(), terminator
penny
2017/07/20 18:38:37
Done.
|
+ value_bytes->push_back(0); |
+ |
+ return; |
+} |
+ |
//------------------------------------------------------------------------------ |
// Misc wrapper functions - LOCAL |
//------------------------------------------------------------------------------ |
@@ -749,8 +790,7 @@ void CloseRegKey(HANDLE key) { |
bool QueryRegKeyValue(HANDLE key, |
const wchar_t* value_name, |
ULONG* out_type, |
- BYTE** out_buffer, |
- DWORD* out_size) { |
+ std::vector<BYTE>* out_buffer) { |
if (!g_initialized) |
InitNativeRegApi(); |
@@ -774,11 +814,14 @@ bool QueryRegKeyValue(HANDLE key, |
value_info, size_needed, &size_needed); |
if (NT_SUCCESS(ntstatus)) { |
*out_type = value_info->Type; |
- *out_size = value_info->DataLength; |
- *out_buffer = new BYTE[*out_size]; |
- ::memcpy(*out_buffer, |
- (reinterpret_cast<BYTE*>(value_info) + value_info->DataOffset), |
- *out_size); |
+ DWORD data_size = value_info->DataLength; |
+ out_buffer->clear(); |
grt (UTC plus 2)
2017/07/20 05:22:53
nit: by using assign() below, clear() is only need
penny
2017/07/20 18:38:36
Done.
|
+ if (data_size) { |
+ // Move the data into |out_buffer| vector. |
+ BYTE* data = reinterpret_cast<BYTE*>(value_info) + value_info->DataOffset; |
+ for (size_t i = 0; i < data_size; i++) |
grt (UTC plus 2)
2017/07/20 05:22:53
out_buffer->assign(data, data + data_size);
penny
2017/07/20 18:38:37
assign. thank you.
|
+ out_buffer->push_back(data[i]); |
+ } |
success = true; |
} |
@@ -791,16 +834,18 @@ bool QueryRegValueDWORD(HANDLE key, |
const wchar_t* value_name, |
DWORD* out_dword) { |
ULONG type = REG_NONE; |
- BYTE* value_bytes = nullptr; |
- DWORD ret_size = 0; |
+ std::vector<BYTE> value_bytes; |
+ |
+ if (!QueryRegKeyValue(key, value_name, &type, &value_bytes) || |
+ type != REG_DWORD) { |
+ return false; |
+ } |
- if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) || |
- type != REG_DWORD) |
+ if (value_bytes.size() < sizeof(*out_dword)) |
return false; |
- *out_dword = *(reinterpret_cast<DWORD*>(value_bytes)); |
+ *out_dword = *(reinterpret_cast<DWORD*>(value_bytes.data())); |
- delete[] value_bytes; |
return true; |
} |
@@ -828,17 +873,18 @@ bool QueryRegValueDWORD(ROOT_KEY root, |
bool QueryRegValueSZ(HANDLE key, |
const wchar_t* value_name, |
std::wstring* out_sz) { |
- BYTE* value_bytes = nullptr; |
- DWORD ret_size = 0; |
+ std::vector<BYTE> value_bytes; |
ULONG type = REG_NONE; |
- if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) || |
- type != REG_SZ) |
+ if (!QueryRegKeyValue(key, value_name, &type, &value_bytes) || |
+ (type != REG_SZ && type != REG_EXPAND_SZ)) { |
return false; |
+ } |
+ |
+ EnsureTerminatedSZ(&value_bytes, false); |
- *out_sz = reinterpret_cast<wchar_t*>(value_bytes); |
+ *out_sz = reinterpret_cast<wchar_t*>(value_bytes.data()); |
grt (UTC plus 2)
2017/07/20 05:22:53
should this preserve embedded nulls? if yes, use o
penny
2017/07/20 18:38:37
I never expected the result of the QueryRegValueSZ
grt (UTC plus 2)
2017/07/21 08:09:00
Sounds reasonable to me. Please document this beha
penny
2017/07/25 18:30:47
Done.
|
- delete[] value_bytes; |
return true; |
} |
@@ -866,18 +912,20 @@ bool QueryRegValueSZ(ROOT_KEY root, |
bool QueryRegValueMULTISZ(HANDLE key, |
const wchar_t* value_name, |
std::vector<std::wstring>* out_multi_sz) { |
- BYTE* value_bytes = nullptr; |
- DWORD ret_size = 0; |
+ std::vector<BYTE> value_bytes; |
ULONG type = REG_NONE; |
- if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) || |
- type != REG_MULTI_SZ) |
+ if (!QueryRegKeyValue(key, value_name, &type, &value_bytes) || |
+ type != REG_MULTI_SZ) { |
return false; |
+ } |
- // Make sure the vector is empty to start. |
- (*out_multi_sz).resize(0); |
+ EnsureTerminatedSZ(&value_bytes, true); |
- wchar_t* pointer = reinterpret_cast<wchar_t*>(value_bytes); |
+ // Make sure the out vector is empty to start. |
+ (*out_multi_sz).clear(); |
+ |
+ wchar_t* pointer = reinterpret_cast<wchar_t*>(value_bytes.data()); |
std::wstring temp = pointer; |
robertshield
2017/07/20 17:14:04
Won't this truncate any characters after the first
penny
2017/07/20 18:38:37
No. It merely creates a new wstring object, and c
|
// Loop. Each string is separated by '\0'. Another '\0' at very end (so 2 in |
// a row). |
@@ -888,11 +936,6 @@ bool QueryRegValueMULTISZ(HANDLE key, |
temp = pointer; |
} |
- // Handle the case of "empty multi_sz". |
- if (out_multi_sz->size() == 0) |
- out_multi_sz->push_back(L""); |
- |
- delete[] value_bytes; |
return true; |
} |