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));
-    ::RegCloseKey(*key);
-    key = NULL;
-  }
+// Helper to set DWORD registry values.
+DWORD SetDWValue(HKEY* key, const wchar_t* property, DWORD value) {
+  return ::RegSetValueEx(*key,
+                         property,
+                         0,
+                         REG_DWORD,
+                         reinterpret_cast<LPBYTE>(&value),
+                         sizeof(value));
 }
 
 }  // 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 HandlePreviousBeacon(HKEY* key, DWORD blacklist_state) {

[csharp, 2014/06/11 13:34:43]

This function shouldn't be indented and can be moved into the unnamed namespace

[krstnmnlsn, 2014/06/11 16:35:27]

Done.

+    LONG result;
+    if (blacklist_state == BLACKLIST_SETUP_RUNNING) {
+      // Some part of the blacklist setup failed last time.  If this has occured
+      // 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) {

[csharp, 2014/06/11 13:34:43]

I'm not sure we need lines 88->95. It just sets the value of kBeaconAttempt, but
then the next bit of code increase the value, so we don't really need to write
it to the registry.

It might be cleaner if we just have "attempt_count=0" for when we hit this
error.

[krstnmnlsn, 2014/06/11 16:35:27]

Done.

+        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);

[csharp, 2014/06/11 13:34:43]

It probably is cleaner to move all the ::RegCloseKey calls to line 155, so you
only need to have it once (instead of needing to remember if for all the return
false statements here)

[krstnmnlsn, 2014/06/11 16:35:27]

Yea I realized after that might be better.  Done.

+        return false;
+      }
+
+      attempt_count = attempt_count + 1;

[csharp, 2014/06/11 13:34:43]

++attempt_count instead?

[krstnmnlsn, 2014/06/11 16:35:27]

Done.

+      result = SetDWValue(key, kBeaconAttemptCount, attempt_count);
+      if (result != ERROR_SUCCESS) {

[csharp, 2014/06/11 13:34:43]

I wonder if we shouldn't try to abort here. If we fail to change the key here,
then we don't try to disable the blacklist on line 106, which might be required,
so we might end up running too many times.

If we do make that change, the only downside is we might end up under reporting
how many time we try before failing, right? What are your thoughts?

[krstnmnlsn, 2014/06/11 16:35:27]

Yeah, I don't think taking it out even introduces any bad cases that weren't
already there.
On the other hand, right now I guess the worst case is that we return false, and
so skip setup but claim it is running .. and end up resetting the count to zero.
So I'll take it out! 

+        ::RegCloseKey(*key);
+        return false;
+      }
+
+      if (attempt_count >= blacklist::kBeaconMaxAttempts) {
+        blacklist_state = BLACKLIST_SETUP_FAILED;
+        SetDWValue(key, kBeaconState, blacklist_state);
+        ::RegCloseKey(*key);
+        return false;
+      }
+    } else if (blacklist_state == BLACKLIST_ENABLED) {
+      // If the blacklist succeeded on the previous run reset the failure
+      // counter.
+      result = SetDWValue(key, kBeaconAttemptCount, static_cast<DWORD>(0));
+      if (result != ERROR_SUCCESS) {
+        ::RegCloseKey(*key);
+        return false;
+      }
+    }
+    return true;
+  }
+
 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 (!HandlePreviousBeacon(&key, blacklist_state))

[csharp, 2014/06/11 13:34:43]

This name seems a bit vague, I don't really have a better idea right now, but
I'm not a bit fan of this.

[krstnmnlsn, 2014/06/11 16:35:27]

Yea I couldn't think of a reasonably short but also descriptive and not
misleading name.  I guess I could call it something like UpdateAttemptCount?
though it does slightly more than that.

+    return false;
+
+  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;
+  }
+
+  // Reaching this point with the setup running state means the setup
+  // succeeded and so we reset to enabled.  Any other state indicates that setup
+  // was skipped; in that case we leave the state alone for later recording.
+  if (blacklist_state == BLACKLIST_SETUP_RUNNING)
+    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.
   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