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

chrome_elf/blacklist/blacklist.cc

 // Copyright 2013 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "chrome_elf/blacklist/blacklist.h"
 
 #include <assert.h>
 #include <string.h>
 
 #include "base/basictypes.h"
@@ skipping 28 matching lines @@
 // an extra allocation at run time.
 #pragma section(".crthunk",read,execute)
 __declspec(allocate(".crthunk")) sandbox::ThunkData g_thunk_storage;
 
 namespace {
 
 // Record if the blacklist was successfully initialized so processes can easily
 // 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
 
 namespace blacklist {
 
 #if defined(_WIN64)
   // Allocate storage for the pointer to the old NtMapViewOfSectionFunction.
 #pragma section(".oldntmap",write,read)
   __declspec(allocate(".oldntmap"))
     NtMapViewOfSectionFunction g_nt_map_view_of_section_func = NULL;
 #endif
 
 bool LeaveSetupBeacon() {
   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)
     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,
                              kBeaconState,
                              0,
                              &type,
                              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);
 }
 
 bool ResetBeacon() {
   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)
     return false;
 
-  DWORD blacklist_state = BLACKLIST_ENABLED;
-  result = ::RegSetValueEx(key,
-                           kBeaconState,
-                           0,
-                           REG_DWORD,
-                           reinterpret_cast<LPBYTE>(&blacklist_state),
-                           sizeof(blacklist_state));
+  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);
 }
 
 int BlacklistSize() {
   int size = -1;
   while (blacklist::g_troublesome_dlls[++size] != NULL) {}
 
   return size;
 }
 
@@ skipping 77 matching lines @@
 
 bool Initialize(bool force) {
   // Check to see that we found the functions we need in ntdll.
   if (!InitializeInterceptImports())
     return false;
 
   // Check to see if this is a non-browser process, abort if so.
   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;
 
   // It is possible for other dlls to have already patched code by now and
   // attempting to patch their code might result in crashes.
   const bool kRelaxed = false;
 
   // Create a thunk via the appropriate ServiceResolver instance.
   sandbox::ServiceResolverThunk* thunk = GetThunk(kRelaxed);
 
   // Don't try blacklisting on unsupported OS versions.
   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
   // ready to write to it.
   DWORD old_protect = 0;
   if (!VirtualProtect(&g_thunk_storage,
                       sizeof(g_thunk_storage),
                       PAGE_EXECUTE_READWRITE,
                       &old_protect)) {
-    RecordSuccessfulThunkSetup(&key);
     return false;
   }
 
   thunk->AllowLocalPatches();
 
   // We declare this early so it can be used in the 64-bit block below and
   // still work on 32-bit build when referenced at the end of the function.
   BOOL page_executable = false;
 
   // Replace the default NtMapViewOfSection with our patched version.
@@ skipping 31 matching lines @@
 
   // Record if we have initialized the blacklist.
   g_blacklist_initialized = NT_SUCCESS(ret);
 
   // Mark the thunk storage as executable and prevent any future writes to it.
   page_executable = page_executable && VirtualProtect(&g_thunk_storage,
                                                       sizeof(g_thunk_storage),
                                                       PAGE_EXECUTE_READ,
                                                       &old_protect);
 
-  RecordSuccessfulThunkSetup(&key);
-
   return NT_SUCCESS(ret) && page_executable;
 }
 
 }  // namespace blacklist