Can now adjust the number of retries before the blacklist is disabled.

chrome_elf/blacklist/blacklist.cc

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 do, I would remove it (The code
that use this also seem to generally already handle closing the key on failure,
but it also isn't always checking for failure).

And if we remove that part, I'm not sure this code is worth still being a
function

[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 a helper function

[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 above block since only 1 can
happen at a time

[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 the value in this case

[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;
 }