Index: chrome_elf/blacklist/blacklist.cc |
diff --git a/chrome_elf/blacklist/blacklist.cc b/chrome_elf/blacklist/blacklist.cc |
index a168c9a044c4ff0a5a54efa57eeeb6556ddcab3b..649b1c98988dffff25ebf6e9815ebf6ce4982e4c 100644 |
--- a/chrome_elf/blacklist/blacklist.cc |
+++ b/chrome_elf/blacklist/blacklist.cc |
@@ -46,20 +46,17 @@ namespace { |
// determine if the blacklist is enabled for them. |
bool g_blacklist_initialized = false; |
-// Record that the thunk setup completed succesfully and close the registry |
-// key handle since it is no longer needed. |
-void RecordSuccessfulThunkSetup(HKEY* key) { |
- if (key != NULL) { |
- DWORD blacklist_state = blacklist::BLACKLIST_SETUP_RUNNING; |
- ::RegSetValueEx(*key, |
- blacklist::kBeaconState, |
- 0, |
- REG_DWORD, |
- reinterpret_cast<LPBYTE>(&blacklist_state), |
- sizeof(blacklist_state)); |
+// Helper to set DWORD registry values. Closes the key on failure. |
csharp
2014/06/10 14:07:02
Closing the key only value seems an odd thing to d
krstnmnlsn
2014/06/10 22:03:27
Half done :)
|
+DWORD SetDWValue(HKEY* key, const wchar_t* property, DWORD value) { |
+ DWORD result = ::RegSetValueEx(*key, |
+ property, |
+ 0, |
+ REG_DWORD, |
+ reinterpret_cast<LPBYTE>(&value), |
+ sizeof(value)); |
+ if (result != ERROR_SUCCESS) |
::RegCloseKey(*key); |
- key = NULL; |
- } |
+ return result; |
} |
} // namespace |
@@ -89,7 +86,7 @@ bool LeaveSetupBeacon() { |
return false; |
// Retrieve the current blacklist state. |
- DWORD blacklist_state = BLACKLIST_DISABLED; |
+ DWORD blacklist_state = BLACKLIST_STATE_MAX; |
DWORD blacklist_state_size = sizeof(blacklist_state); |
DWORD type = 0; |
result = ::RegQueryValueEx(key, |
@@ -99,21 +96,55 @@ bool LeaveSetupBeacon() { |
reinterpret_cast<LPBYTE>(&blacklist_state), |
&blacklist_state_size); |
- if (blacklist_state != BLACKLIST_ENABLED || |
- result != ERROR_SUCCESS || type != REG_DWORD) { |
+ if (blacklist_state == BLACKLIST_DISABLED || result != ERROR_SUCCESS || |
+ type != REG_DWORD) { |
::RegCloseKey(key); |
return false; |
} |
- // Mark the blacklist setup code as running so if it crashes the blacklist |
- // won't be enabled for the next run. |
- blacklist_state = BLACKLIST_SETUP_RUNNING; |
- result = ::RegSetValueEx(key, |
- kBeaconState, |
- 0, |
- REG_DWORD, |
- reinterpret_cast<LPBYTE>(&blacklist_state), |
- sizeof(blacklist_state)); |
+ if (blacklist_state == BLACKLIST_SETUP_RUNNING) { |
+ // Some part of the blacklist setup failed last time. If this has occured |
csharp
2014/06/10 14:07:02
Could you move the code in this if statement into
krstnmnlsn
2014/06/10 22:03:27
Done.
|
+ // some minimum number of times in a row we switch the state to failed and |
+ // skip setting up the blacklist. |
+ DWORD attempt_count = 0; |
+ DWORD attempt_count_size = sizeof(attempt_count); |
+ DWORD attempt_count_type = 0; |
+ result = ::RegQueryValueEx(key, |
+ kBeaconAttemptCount, |
+ 0, |
+ &attempt_count_type, |
+ reinterpret_cast<LPBYTE>(&attempt_count), |
+ &attempt_count_size); |
+ |
+ if (result == ERROR_FILE_NOT_FOUND) { |
+ attempt_count = 0; |
+ attempt_count_type = REG_DWORD; |
+ result = SetDWValue(&key, kBeaconAttemptCount, attempt_count); |
+ } |
+ |
+ if (result != ERROR_SUCCESS || attempt_count_type != REG_DWORD) { |
+ ::RegCloseKey(key); |
+ return false; |
+ } |
+ |
+ attempt_count = attempt_count + 1; |
+ result = SetDWValue(&key, kBeaconAttemptCount, attempt_count); |
+ if (result != ERROR_SUCCESS) |
+ return false; |
+ |
+ if (attempt_count >= blacklist::kBeaconMaxAttempts) { |
+ blacklist_state = BLACKLIST_SETUP_FAILED; |
+ SetDWValue(&key, kBeaconState, blacklist_state); |
+ ::RegCloseKey(key); |
+ return false; |
+ } |
+ } |
+ |
+ // If the blacklist succeeded on the previous run, reset the failure counter. |
+ if (blacklist_state == BLACKLIST_ENABLED) |
csharp
2014/06/10 14:07:02
This should probably be join via else with the abo
krstnmnlsn
2014/06/10 22:03:27
Done.
|
+ SetDWValue(&key, kBeaconAttemptCount, static_cast<DWORD>(0)); |
+ |
+ result = SetDWValue(&key, kBeaconState, BLACKLIST_SETUP_RUNNING); |
::RegCloseKey(key); |
return (result == ERROR_SUCCESS); |
@@ -134,15 +165,25 @@ bool ResetBeacon() { |
if (result != ERROR_SUCCESS) |
return false; |
- DWORD blacklist_state = BLACKLIST_ENABLED; |
- result = ::RegSetValueEx(key, |
- kBeaconState, |
- 0, |
- REG_DWORD, |
- reinterpret_cast<LPBYTE>(&blacklist_state), |
- sizeof(blacklist_state)); |
- ::RegCloseKey(key); |
+ DWORD blacklist_state = BLACKLIST_STATE_MAX; |
+ DWORD blacklist_state_size = sizeof(blacklist_state); |
+ DWORD type = 0; |
+ result = ::RegQueryValueEx(key, |
+ kBeaconState, |
+ 0, |
+ &type, |
+ reinterpret_cast<LPBYTE>(&blacklist_state), |
+ &blacklist_state_size); |
+ if (result != ERROR_SUCCESS || type != REG_DWORD) { |
+ ::RegCloseKey(key); |
+ return false; |
+ } |
+ |
+ if (blacklist_state == BLACKLIST_SETUP_RUNNING) |
csharp
2014/06/10 14:07:02
Please add a comment explaining why we only reset
krstnmnlsn
2014/06/10 22:03:27
Done.
|
+ result = SetDWValue(&key, kBeaconState, BLACKLIST_ENABLED); |
+ |
+ ::RegCloseKey(key); |
return (result == ERROR_SUCCESS); |
} |
@@ -240,7 +281,8 @@ bool Initialize(bool force) { |
if (IsNonBrowserProcess()) |
return false; |
- // Check to see if a beacon is present, abort if so. |
+ // Check to see if the blacklist beacon is still set to running (indicating a |
+ // failure) or disabled and abort if so. |
csharp
2014/06/10 14:07:02
nit: disabled -> disabled,
krstnmnlsn
2014/06/10 22:03:27
Done.
|
if (!force && !LeaveSetupBeacon()) |
return false; |
@@ -255,30 +297,6 @@ bool Initialize(bool force) { |
if (!thunk) |
return false; |
- // Record that we are starting the thunk setup code. |
- HKEY key = NULL; |
- DWORD disposition = 0; |
- LONG result = ::RegCreateKeyEx(HKEY_CURRENT_USER, |
- kRegistryBeaconPath, |
- 0, |
- NULL, |
- REG_OPTION_NON_VOLATILE, |
- KEY_QUERY_VALUE | KEY_SET_VALUE, |
- NULL, |
- &key, |
- &disposition); |
- if (result == ERROR_SUCCESS) { |
- DWORD blacklist_state = BLACKLIST_THUNK_SETUP; |
- ::RegSetValueEx(key, |
- kBeaconState, |
- 0, |
- REG_DWORD, |
- reinterpret_cast<LPBYTE>(&blacklist_state), |
- sizeof(blacklist_state)); |
- } else { |
- key = NULL; |
- } |
- |
BYTE* thunk_storage = reinterpret_cast<BYTE*>(&g_thunk_storage); |
// Mark the thunk storage as readable and writeable, since we |
@@ -288,7 +306,6 @@ bool Initialize(bool force) { |
sizeof(g_thunk_storage), |
PAGE_EXECUTE_READWRITE, |
&old_protect)) { |
- RecordSuccessfulThunkSetup(&key); |
return false; |
} |
@@ -340,8 +357,6 @@ bool Initialize(bool force) { |
PAGE_EXECUTE_READ, |
&old_protect); |
- RecordSuccessfulThunkSetup(&key); |
- |
return NT_SUCCESS(ret) && page_executable; |
} |