Make Chrome Frame registration and unregistration more robust.

chrome_frame/chrome_tab.cc

 // Copyright (c) 2011 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.
 
 // chrome_tab.cc : Implementation of DLL Exports.
 
 // Include without path to make GYP build see it.
 #include "chrome_tab.h"  // NOLINT
 
 #include <atlsecurity.h>
 
 #include "base/at_exit.h"
 #include "base/command_line.h"
 #include "base/file_util.h"
 #include "base/file_version_info.h"
 #include "base/logging.h"
 #include "base/logging_win.h"
+#include "base/memory/scoped_vector.h"
 #include "base/path_service.h"
 #include "base/string_number_conversions.h"
 #include "base/string_piece.h"
 #include "base/string_util.h"
 #include "base/sys_string_conversions.h"
 #include "base/win/registry.h"
 #include "base/win/windows_version.h"
 #include "chrome/common/chrome_constants.h"
 #include "chrome/common/chrome_switches.h"
 #include "chrome/installer/util/google_update_settings.h"
@@ skipping 156 matching lines @@
     return hr;
   }
 
   // See comments in AddCommonRGSReplacements
   bool do_system_registration_;
 };
 
 ChromeTabModule _AtlModule;
 
 base::AtExitManager* g_exit_manager = NULL;
-bool RegisterSecuredMimeHandler(bool enable, bool is_system);  // forward
+HRESULT RegisterSecuredMimeHandler(bool enable, bool is_system);  // forward
 
 // DLL Entry Point
 extern "C" BOOL WINAPI DllMain(HINSTANCE instance,
                                DWORD reason,
                                LPVOID reserved) {
   UNREFERENCED_PARAMETER(instance);
   if (reason == DLL_PROCESS_ATTACH) {
 #ifndef NDEBUG
     // Silence traces from the ATL registrar to reduce the log noise.
     ATL::CTrace::s_trace.ChangeCategory(atlTraceRegistrar, 0,
@@ skipping 136 matching lines @@
     // needs doing.
     result = S_FALSE;
   }
 
   return result;
 }
 
 // Helper method called for user-level installs where we don't have admin
 // permissions. Starts up the long running process and registers it to get it
 // started at next boot.
-void SetupUserLevelHelper() {
+HRESULT SetupUserLevelHelper() {
+  HRESULT hr = S_OK;
+
   // Remove existing run-at-startup entry.
   base::win::RemoveCommandFromAutoRun(HKEY_CURRENT_USER, kRunKeyName);
 
   // Build the chrome_frame_helper command line.
   FilePath module_path;
   FilePath helper_path;
   if (PathService::Get(base::FILE_MODULE, &module_path)) {
     module_path = module_path.DirName();
     helper_path = module_path.Append(kChromeFrameHelperExe);
     if (!file_util::PathExists(helper_path)) {
       // If we can't find the helper in the current directory, try looking
       // one up (this is the layout in the build output folder).
       module_path = module_path.DirName();
       helper_path = module_path.Append(kChromeFrameHelperExe);
       DCHECK(file_util::PathExists(helper_path)) <<
           "Could not find chrome_frame_helper.exe.";
     }
+
+    // Find window handle of existing instance.
+    HWND old_window = FindWindow(kChromeFrameHelperWindowClassName,
+                                 kChromeFrameHelperWindowName);
+
+    if (file_util::PathExists(helper_path)) {
+      std::wstring helper_path_cmd(L"\"");
+      helper_path_cmd += helper_path.value();
+      helper_path_cmd += L"\" ";
+      helper_path_cmd += kChromeFrameHelperStartupArg;
+
+      // Add new run-at-startup entry.
+      base::win::AddCommandToAutoRun(HKEY_CURRENT_USER, kRunKeyName,

[grt (UTC plus 2), 2011/09/06 14:56:55]

if this fails, installation will continue but GCF won't work after the user logs
out and back on.  is this desirable?

[robertshield, 2011/09/06 20:23:42]

Nope, we now fail on this as well.

+                                     helper_path_cmd);
+
+      // Start new instance.
+      base::LaunchOptions options;
+      options.start_hidden = true;
+      bool launched = base::LaunchProcess(helper_path.value(), options, NULL);
+      if (!launched) {
+        hr = E_FAIL;

[grt (UTC plus 2), 2011/09/06 14:56:55]

this will now cause installation to fail and rollback (it didn't previously). 
this is desired?

[robertshield, 2011/09/06 20:23:42]

Yes, this is desired. New IE instances will not render in Chrome Frame if this
does not succeed.

+        NOTREACHED();
+        LOG(ERROR) << "Could not launch helper process.";

[grt (UTC plus 2), 2011/09/06 14:56:55]

Do you think the last-error code might be useful here?  Change to PLOG if so.

[robertshield, 2011/09/06 20:23:42]

So changed, although sadly this won't show up in the install log. Me wonders if
it is worth it to have the installer listen for logging messages (would be a
separate CL).

[grt (UTC plus 2), 2011/09/07 15:00:48]

frown.  so the logging stuff in DllMain is a noop during installation?  does ETW
logging (as seen in sawbuck) not work either?

[robertshield, 2011/09/07 17:58:11]

As far as I can tell, this is the case. It may be possible to get selfreg dlls
to hand off logging information to the installer but that is outside of the
scope of the current cl.

+      }
+
+      // Kill old instance using window handle.
+      if (IsWindow(old_window)) {
+        BOOL result = PostMessage(old_window, WM_CLOSE, 0, 0);
+        if (!result) {
+          LOG(ERROR) << "Failed to post close message to old helper process: "

[grt (UTC plus 2), 2011/09/06 14:56:55]

PLOG and remove GetLastError

[robertshield, 2011/09/06 20:23:42]

Done.

+                     << GetLastError();
+        }
+      }
+    } else {
+      hr = ERROR_FILE_NOT_FOUND;

[grt (UTC plus 2), 2011/09/06 14:56:55]

hr = HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND);

[robertshield, 2011/09/06 20:23:42]

Done.

+    }
   } else {
+    hr = E_UNEXPECTED;
     NOTREACHED();
   }
 
-  // Find window handle of existing instance.
-  HWND old_window = FindWindow(kChromeFrameHelperWindowClassName,
-                               kChromeFrameHelperWindowName);
-
-  if (file_util::PathExists(helper_path)) {
-    std::wstring helper_path_cmd(L"\"");
-    helper_path_cmd += helper_path.value();
-    helper_path_cmd += L"\" ";
-    helper_path_cmd += kChromeFrameHelperStartupArg;
-
-    // Add new run-at-startup entry.
-    base::win::AddCommandToAutoRun(HKEY_CURRENT_USER, kRunKeyName,
-                                   helper_path_cmd);
-
-    // Start new instance.
-    base::LaunchOptions options;
-    options.start_hidden = true;
-    bool launched = base::LaunchProcess(helper_path.value(), options, NULL);
-    if (!launched) {
-      NOTREACHED();
-      LOG(ERROR) << "Could not launch helper process.";
-    }
-
-    // Kill old instance using window handle.
-    if (IsWindow(old_window)) {
-      BOOL result = PostMessage(old_window, WM_CLOSE, 0, 0);
-      if (!result) {
-        LOG(ERROR) << "Failed to post close message to old helper process: "
-                   << GetLastError();
-      }
-    }
-  }
+  return hr;
 }
 
 
 // Used to determine whether the DLL can be unloaded by OLE
 STDAPI DllCanUnloadNow() {
   return _AtlModule.DllCanUnloadNow();
 }
 
 // Returns a class factory to create an object of the requested type
 STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID* ppv) {
@@ skipping 18 matching lines @@
 const wchar_t kChromeFramePrefix[] = L"chromeframe/";
 
 // To delete the user agent, set value to NULL.
 // The is_system parameter indicates whether this is a per machine or a per
 // user installation.
 HRESULT SetChromeFrameUA(bool is_system, const wchar_t* value) {
   HRESULT hr = E_FAIL;
   HKEY parent_hive = is_system ? HKEY_LOCAL_MACHINE : HKEY_CURRENT_USER;
 
   RegKey ua_key;
   if (ua_key.Create(parent_hive, kPostPlatformUAKey,

[grt (UTC plus 2), 2011/09/06 14:56:55]

capture the return value here

[robertshield, 2011/09/06 20:23:42]

Done.

                     KEY_READ | KEY_WRITE) == ERROR_SUCCESS) {
     // Make sure that we unregister ChromeFrame UA strings registered previously
     wchar_t value_name[MAX_PATH + 1] = {};
     wchar_t value_data[MAX_PATH + 1] = {};
 
     DWORD value_index = 0;
     while (value_index < ua_key.ValueCount()) {
       DWORD name_size = arraysize(value_name);
       DWORD value_size = arraysize(value_data);
       DWORD type = 0;
@@ skipping 12 matching lines @@
       }
     }
 
     std::wstring chrome_frame_ua_value_name = kChromeFramePrefix;
     chrome_frame_ua_value_name += GetCurrentModuleVersion();
     if (value) {
       ua_key.WriteValue(chrome_frame_ua_value_name.c_str(), value);
     }
     hr = S_OK;
   } else {
     DLOG(ERROR) << __FUNCTION__ << ": " << kPostPlatformUAKey;

[grt (UTC plus 2), 2011/09/06 14:56:55]

log the captured return value here

[robertshield, 2011/09/06 20:23:42]

Done.

     hr = E_UNEXPECTED;

[grt (UTC plus 2), 2011/09/06 14:56:55]

hr = HRESULT_FROM_WIN32(captured value)

[robertshield, 2011/09/06 20:23:42]

Done.

   }
   return hr;
 }
 

[grt (UTC plus 2), 2011/09/06 14:56:55]

consider adding a comment indicating that what follows are the ordered
registration step functions.

[robertshield, 2011/09/06 20:23:42]

Done.

+HRESULT RegisterActiveDoc(bool reg, bool is_system) {
+  // We have to call the static T::UpdateRegistry function instead of
+  // _AtlModule.UpdateRegistryFromResourceS(IDR_CHROMEFRAME_ACTIVEDOC, reg)
+  // because there is specific OLEMISC replacement.
+  return ChromeActiveDocument::UpdateRegistry(reg);
+}
+
+HRESULT RegisterActiveX(bool reg, bool is_system) {
+  // We have to call the static T::UpdateRegistry function instead of
+  // _AtlModule.UpdateRegistryFromResourceS(IDR_CHROMEFRAME_ACTIVEX, reg)
+  // because there is specific OLEMISC replacement.
+  return ChromeFrameActivex::UpdateRegistry(reg);
+}
+
+HRESULT RegisterElevationPolicy(bool reg, bool is_system) {
+  if (reg && base::win::GetVersion() >= base::win::VERSION_VISTA) {
+    // Register the elevation policy. We do this only for developer
+    // convenience as the installer is really responsible for doing this.

[grt (UTC plus 2), 2011/09/06 14:56:55]

This comment says that the installer is responsible for doing this.  Does the
installer, in fact, do it; or has responsibility been shifted to this function?

[robertshield, 2011/09/06 20:23:42]

It does, in fact, do it. See install_worker.cc, AddElevationPolicyWorkItems and
RefreshElevationPolicy.

+    // Because of that, we do not unregister this policy and just leave that
+    // up to the installer. Also we do not gate registration on this
+    // succeeding.
+    _AtlModule.UpdateRegistryFromResourceS(IDR_CHROMEFRAME_ELEVATION,
+                                           reg);
+    RefreshElevationPolicy();
+  }
+  return S_OK;
+}
+
+HRESULT RegisterProtocol(bool reg, bool is_system) {
+  return _AtlModule.UpdateRegistryFromResourceS(IDR_CHROMEPROTOCOL, reg);
+}
+
+HRESULT RegisterBhoClsid(bool reg, bool is_system) {
+  return Bho::UpdateRegistry(reg);
+}
+
+HRESULT RegisterBhoIE(bool reg, bool is_system) {
+  if (is_system) {
+    return _AtlModule.UpdateRegistryFromResourceS(IDR_REGISTER_BHO, reg);
+  } else {
+    if (reg) {
+      // Setup the long running process:
+      return SetupUserLevelHelper();
+    } else {
+      // Unschedule the user-level helper. Note that we don't kill it here
+      // so that during updates we don't have a time window with no running
+      // helper. Uninstalls and updates will explicitly kill the helper from
+      // within the installer. Unregister existing run-at-startup entry.
+      return base::win::RemoveCommandFromAutoRun(HKEY_CURRENT_USER,

[grt (UTC plus 2), 2011/09/06 14:56:55]

return base::win::Remove...(...) ? S_OK : E_FAIL;

[robertshield, 2011/09/06 20:23:42]

Done.

+                                                 kRunKeyName);
+    }
+  }
+}
+
+HRESULT RegisterTypeLib(bool reg, bool is_system) {
+  if (reg && !is_system) {
+    // Enables the RegisterTypeLib Function function to override default
+    // registry mappings under Windows Vista Service Pack 1 (SP1),
+    // Windows Server 2008, and later operating system versions
+    typedef void (WINAPI* OaEnablePerUserTypeLibReg)(void);
+    OaEnablePerUserTypeLibReg per_user_typelib_func =
+        reinterpret_cast<OaEnablePerUserTypeLibReg>(
+            GetProcAddress(GetModuleHandle(L"oleaut32.dll"),
+                           "OaEnablePerUserTLibRegistration"));
+    if (per_user_typelib_func) {
+      (*per_user_typelib_func)();
+    }
+  }
+  return reg ?
+      UtilRegisterTypeLib(_AtlComModule.m_hInstTypeLib,
+                          NULL, !is_system) :
+      UtilUnRegisterTypeLib(_AtlComModule.m_hInstTypeLib,
+                            NULL, !is_system);
+}
+
+HRESULT RegisterLegacyNPAPICleanup(bool reg, bool is_system) {
+  if (!reg) {
+    _AtlModule.UpdateRegistryFromResourceS(IDR_CHROMEFRAME_NPAPI, reg);
+    UtilRemovePersistentNPAPIMarker();
+  }
+  // Ignore failures.
+  return S_OK;
+}
+

[grt (UTC plus 2), 2011/09/06 14:56:55]

remove extra newline

[robertshield, 2011/09/06 20:23:42]

Done.

+
+HRESULT RegisterAppId(bool reg, bool is_system) {
+  return _AtlModule.UpdateRegistryAppId(reg);
+}
+

[grt (UTC plus 2), 2011/09/06 14:56:55]

remove extra newline

[robertshield, 2011/09/06 20:23:42]

Done.

+
+HRESULT RegisterUserAgent(bool reg, bool is_system) {
+  if (reg) {
+    return SetChromeFrameUA(is_system, L"1");
+  } else {
+    return SetChromeFrameUA(is_system, NULL);
+  }
+}
+
+enum RegistrationSteps {

[grt (UTC plus 2), 2011/09/06 14:56:55]

nit: make this singular since an instance of the type represents one step (much
like an instance of an "int" represents one integer). since RegistrationStep is
already taken, maybe RegistrationStepId.

[robertshield, 2011/09/06 20:23:42]

Done.

+  kStepSecuredMimeHandler = 0,
+  kStepActiveDoc          = 1,
+  kStepActiveX            = 2,
+  kStepElevationPolicy    = 3,
+  kStepProtocol           = 4,
+  kStepBhoClsid           = 5,
+  kStepBhoRegistration    = 6,
+  kStepRegisterTypeLib    = 7,
+  kStepNpapiCleanup       = 8,
+  kStepAppId              = 9,
+  kStepUserAgent          = 10,
+  kStepEnd                = 11
+};
+
 enum RegistrationFlags {
   ACTIVEX             = 0x0001,
   ACTIVEDOC           = 0x0002,
   GCF_PROTOCOL        = 0x0004,
   BHO_CLSID           = 0x0008,
   BHO_REGISTRATION    = 0x0010,
   TYPELIB             = 0x0020,
 
   ALL                 = 0xFFFF
 };
 
-STDAPI CustomRegistration(UINT reg_flags, BOOL reg, bool is_system) {
-  UINT flags = reg_flags;
+// Mux the failure step into the hresult. We take only the first four bits
+// and stick those into the top four bits of the facility code. We also set the
+// Customer bit to be polite. Graphically, we write our error code to the
+// bits marked with ^:
+// |                    1                   2                   3  |
+// |0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1|
+// |S|R|C|N|X|Facility             |Code                           |
+//      ^     ^ ^ ^ ^
+// See http://msdn.microsoft.com/en-us/library/cc231198(PROT.10).aspx for
+// more details on HRESULTS.
+//
+// The resulting error can be extracted by:
+// error_code = (fiddled_hr & 0x07800000) >> 23
+HRESULT MuxErrorIntoHRESULT(HRESULT hr, int error_code) {
+  DCHECK_GE(error_code, 0);
+  DCHECK_LT(error_code, kStepEnd);
+  COMPILE_ASSERT(kStepEnd <= 0xF, update_error_muxing_too_many_steps);
 
-  if (reg && (flags & (ACTIVEDOC | ACTIVEX)))
-    flags |= (TYPELIB | GCF_PROTOCOL);
+  // Check that our four desired bits are clear.
+  // 0xF87FFFFF == 11111000011111111111111111111111
+  DCHECK_EQ(static_cast<HRESULT>(hr & 0xF87FFFFF), hr);
 
-  HRESULT hr = S_OK;
+  HRESULT fiddled_hr = ((error_code & 0xF) << 23) | hr;
+  fiddled_hr |= 1 << 29;  // Set the customer bit.
+
+  return fiddled_hr;
+}
+
+HRESULT CustomRegistration(uint16 reg_flags, bool reg, bool is_system) {
+  if (reg && (reg_flags & (ACTIVEDOC | ACTIVEX)))
+    reg_flags |= (TYPELIB | GCF_PROTOCOL);
 
   // Set the flag that gets checked in AddCommonRGSReplacements before doing
   // registration work.
   _AtlModule.do_system_registration_ = is_system;
 
-  if ((hr == S_OK) && (flags & ACTIVEDOC)) {
-    // Don't fail to unregister if we can't undo the secure mime
-    // handler registration. This was observed getting hit during
-    // uninstallation.
-    if (!RegisterSecuredMimeHandler(reg ? true : false, is_system) && reg)
-      return E_FAIL;
-    hr = ChromeActiveDocument::UpdateRegistry(reg);
-  }
+  typedef HRESULT (*RegistrationFn)(bool reg, bool is_system);
+  struct RegistrationStep {
+    int id;

[grt (UTC plus 2), 2011/09/06 14:56:55]

how about RegistrationStepId rather than int?  this make looking at the steps in
a debugger friendlier.  and it's more typesafe.

[robertshield, 2011/09/06 20:23:42]

Done.

+    uint32 condition;

[grt (UTC plus 2), 2011/09/06 14:56:55]

since reg_flags is a uint16, should this be, too?

[robertshield, 2011/09/06 20:23:42]

Done.

+    RegistrationFn func;
+  };
+  static const RegistrationStep registration_steps[] = {
+    { kStepSecuredMimeHandler, ACTIVEDOC, &RegisterSecuredMimeHandler },
+    { kStepActiveDoc, ACTIVEDOC, &RegisterActiveDoc },
+    { kStepActiveX, ACTIVEX, &RegisterActiveX },
+    { kStepElevationPolicy, (ACTIVEDOC | ACTIVEX), &RegisterElevationPolicy },
+    { kStepProtocol, GCF_PROTOCOL, &RegisterProtocol },
+    { kStepBhoClsid, BHO_CLSID, &RegisterBhoClsid },
+    { kStepBhoRegistration, BHO_REGISTRATION, &RegisterBhoIE },
+    { kStepRegisterTypeLib, TYPELIB, &RegisterTypeLib },
+    { kStepNpapiCleanup, ALL, &RegisterLegacyNPAPICleanup },
+    { kStepAppId, ALL, &RegisterAppId },
+    { kStepUserAgent, ALL, &RegisterUserAgent }
+  };
 
-  if ((hr == S_OK) && (flags & ACTIVEX)) {
-    // We have to call the static T::UpdateRegistry function instead of
-    // _AtlModule.UpdateRegistryFromResourceS(IDR_CHROMEFRAME_ACTIVEX, reg)
-    // because there is specific OLEMISC replacement.
-    hr = ChromeFrameActivex::UpdateRegistry(reg);
-  }
-
-  // Register the elevation policy.  We do this only for developer convenience
-  // as the installer is really responsible for doing this.
-  // Because of that, we do not unregister this policy and just leave that up
-  // to the installer.
-  if (hr == S_OK && (flags & (ACTIVEDOC | ACTIVEX)) && reg) {
-    _AtlModule.UpdateRegistryFromResourceS(IDR_CHROMEFRAME_ELEVATION, reg);
-    RefreshElevationPolicy();
-  }
-
-  if ((hr == S_OK) && (flags & GCF_PROTOCOL)) {
-    hr = _AtlModule.UpdateRegistryFromResourceS(IDR_CHROMEPROTOCOL, reg);
-  }
-
-  if ((hr == S_OK) && (flags & BHO_CLSID)) {
-    hr = Bho::UpdateRegistry(reg);
-  }
-
-  if ((hr == S_OK) && (flags & BHO_REGISTRATION)) {
-    if (is_system) {
-      _AtlModule.UpdateRegistryFromResourceS(IDR_REGISTER_BHO, reg);
-    } else {
-      if (reg) {
-        // Setup the long running process:
-        SetupUserLevelHelper();
-      } else {
-        // Unschedule the user-level helper. Note that we don't kill it here so
-        // that during updates we don't have a time window with no running
-        // helper. Uninstalls and updates will explicitly kill the helper from
-        // within the installer. Unregister existing run-at-startup entry.
-        base::win::RemoveCommandFromAutoRun(HKEY_CURRENT_USER, kRunKeyName);
+  HRESULT hr = S_OK;
+  bool rollback = false;
+  int failure_step = 0;
+  for (int step = 0; step < arraysize(registration_steps); ++step) {
+    if ((reg_flags & registration_steps[step].condition) != 0) {
+      hr = registration_steps[step].func(reg, is_system);
+      if (FAILED(hr) && reg) {
+        rollback = true;
+        failure_step = step;
+        break;
       }
     }
   }
 
-  if ((hr == S_OK) && (flags & TYPELIB)) {
-    if (reg && !is_system) {
-      // Enables the RegisterTypeLib Function function to override default
-      // registry mappings under Windows Vista Service Pack 1 (SP1),
-      // Windows Server 2008, and later operating system versions
-      typedef void (WINAPI* OaEnablePerUserTypeLibReg)(void);
-      OaEnablePerUserTypeLibReg per_user_typelib_func =
-          reinterpret_cast<OaEnablePerUserTypeLibReg>(
-              GetProcAddress(GetModuleHandle(L"oleaut32.dll"),
-                             "OaEnablePerUserTLibRegistration"));
-      if (per_user_typelib_func) {
-        (*per_user_typelib_func)();
-      }
+  if (rollback) {
+    DCHECK(reg);
+    // Rollback the failing action and all preceding ones.
+    for (int rollback_step = failure_step;

[grt (UTC plus 2), 2011/09/06 14:56:55]

how about reusing "step" so this doesn't need so much wrapping?

[robertshield, 2011/09/06 20:23:42]

Done.

+         rollback_step >= 0;
+         --rollback_step) {
+      registration_steps[rollback_step].func(!reg, is_system);
     }
-    hr = (reg)?
-        UtilRegisterTypeLib(_AtlComModule.m_hInstTypeLib, NULL, !is_system) :
-        UtilUnRegisterTypeLib(_AtlComModule.m_hInstTypeLib, NULL, !is_system);
   }
 
-  // Unconditionally remove NPAPI registration when unregistering any component.
-  if ((hr == S_OK) && !reg) {
-    // Ignore failures.
-    _AtlModule.UpdateRegistryFromResourceS(IDR_CHROMEFRAME_NPAPI, reg);
-    UtilRemovePersistentNPAPIMarker();
-  }
-
-  if (hr == S_OK) {
-    hr = _AtlModule.UpdateRegistryAppId(reg);
-  }
-
-  if (hr == S_OK) {
-    if (reg) {
-      hr = SetChromeFrameUA(is_system, L"1");
-    } else {
-      hr = SetChromeFrameUA(is_system, NULL);
-    }
-  }
-  return hr;
+  return MuxErrorIntoHRESULT(hr, failure_step);
 }
 
-
-
 // DllRegisterServer - Adds entries to the system registry
 STDAPI DllRegisterServer() {
-  UINT flags =  ACTIVEX | ACTIVEDOC | TYPELIB | GCF_PROTOCOL |
-                BHO_CLSID | BHO_REGISTRATION;
+  uint16 flags =  ACTIVEX | ACTIVEDOC | TYPELIB | GCF_PROTOCOL |
+                  BHO_CLSID | BHO_REGISTRATION;
 
-  HRESULT hr = CustomRegistration(flags, TRUE, true);
+  HRESULT hr = CustomRegistration(flags, true, true);
   if (SUCCEEDED(hr)) {
     SetupRunOnce();
   }
 
   return hr;
 }
 
 // DllUnregisterServer - Removes entries from the system registry
 STDAPI DllUnregisterServer() {
-  HRESULT hr = CustomRegistration(ALL, FALSE, true);
+  HRESULT hr = CustomRegistration(ALL, false, true);
   return hr;
 }
 
-// DllRegisterServer - Adds entries to the HKCU hive in the registry
+// DllRegisterUserServer - Adds entries to the HKCU hive in the registry.
 STDAPI DllRegisterUserServer() {
   UINT flags =  ACTIVEX | ACTIVEDOC | TYPELIB | GCF_PROTOCOL |
                 BHO_CLSID | BHO_REGISTRATION;
 
   HRESULT hr = CustomRegistration(flags, TRUE, false);
   if (SUCCEEDED(hr)) {
     SetupRunOnce();
   }
 
   return hr;
 }
 
-// DllRegisterServer - Removes entries from the HKCU hive in the registry.
+// DllUnregisterUserServer - Removes entries from the HKCU hive in the registry.
 STDAPI DllUnregisterUserServer() {
   HRESULT hr = CustomRegistration(ALL, FALSE, false);
   return hr;
 }
 
 class SecurityDescBackup {
  public:
   explicit SecurityDescBackup(const std::wstring& backup_key)
       : backup_key_name_(backup_key) {}
   ~SecurityDescBackup() {}
@@ skipping 99 matching lines @@
     return user_;
   }
 
  private:
   CAccessToken token_;
   CTokenPrivileges take_ownership_;
   CTokenPrivileges restore_;
   CSid user_;
 };
 
-static bool SetOrDeleteMimeHandlerKey(bool set, HKEY root_key) {
+static HRESULT SetOrDeleteMimeHandlerKey(bool set, HKEY root_key) {

[grt (UTC plus 2), 2011/09/06 14:56:55]

nit: move this (and the above classes) into the unnamed namespace and remove
"static".

[robertshield, 2011/09/06 20:23:42]

Done.

   std::wstring key_name = kInternetSettings;
   key_name.append(L"\\Secure Mime Handlers");
   RegKey key(root_key, key_name.c_str(), KEY_READ | KEY_WRITE);
   if (!key.Valid())
     return false;
 
   LONG result1 = ERROR_SUCCESS;
   LONG result2 = ERROR_SUCCESS;
   if (set) {
     result1 = key.WriteValue(L"ChromeTab.ChromeActiveDocument", 1);
     result2 = key.WriteValue(L"ChromeTab.ChromeActiveDocument.1", 1);
   } else {
     result1 = key.DeleteValue(L"ChromeTab.ChromeActiveDocument");
     result2 = key.DeleteValue(L"ChromeTab.ChromeActiveDocument.1");
   }
 
-  return (result2 == ERROR_SUCCESS) && (result2 == ERROR_SUCCESS);
+  return result1 != ERROR_SUCCESS ? HRESULT_FROM_WIN32(result1) :
+                                    HRESULT_FROM_WIN32(result2);
 }
 
-bool RegisterSecuredMimeHandler(bool enable, bool is_system) {
+HRESULT RegisterSecuredMimeHandler(bool enable, bool is_system) {

[grt (UTC plus 2), 2011/09/06 14:56:55]

nit: move this so that it lives with the other registration funcs, and is in the
proper order.  mega-nit: move all of them into the unnamed namespace.

[robertshield, 2011/09/06 20:23:42]

Done.

   if (!is_system) {
     return SetOrDeleteMimeHandlerKey(enable, HKEY_CURRENT_USER);
   } else if (base::win::GetVersion() < base::win::VERSION_VISTA) {
     return SetOrDeleteMimeHandlerKey(enable, HKEY_LOCAL_MACHINE);
   }
 
   std::wstring mime_key = kInternetSettings;
   mime_key.append(L"\\Secure Mime Handlers");
   std::wstring backup_key = kInternetSettings;
   backup_key.append(L"\\__backup_SMH__");
   std::wstring object_name = L"MACHINE\\";
   object_name.append(mime_key);
 
   TokenWithPrivileges token_;
   if (!token_.EnablePrivileges())
-    return false;
+    return E_ACCESSDENIED;
 
   // If there is a backup key - something bad happened; try to restore
   // security on "Secure Mime Handlers" from the backup.
   SecurityDescBackup backup(backup_key);
   backup.RestoreSecurity(object_name.c_str());
 
   // Read old security descriptor of the Mime key first.
   CSecurityDesc sd;
   if (!AtlGetSecurityDescriptor(object_name.c_str(), SE_REGISTRY_KEY, &sd)) {
-    return false;
+    return E_FAIL;
   }
 
   backup.SaveSecurity(sd);
-  bool result = false;
+  HRESULT hr = E_FAIL;
   // set new owner
   if (AtlSetOwnerSid(object_name.c_str(), SE_REGISTRY_KEY, token_.GetUser())) {
     // set new dacl
     CDacl new_dacl;
     sd.GetDacl(&new_dacl);
     new_dacl.AddAllowedAce(token_.GetUser(), GENERIC_WRITE | GENERIC_READ);
     if (AtlSetDacl(object_name.c_str(), SE_REGISTRY_KEY, new_dacl)) {
-      result = SetOrDeleteMimeHandlerKey(enable, HKEY_LOCAL_MACHINE);
+      hr = SetOrDeleteMimeHandlerKey(enable, HKEY_LOCAL_MACHINE);
     }
   }
 
   backup.RestoreSecurity(object_name.c_str());
-  return result;
+  return hr;
 }