Prevent redirector from returning the current module in cases where it fails to look up the first...

chrome_frame/dll_redirector.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.
 
 #include "chrome_frame/dll_redirector.h"
 
 #include <aclapi.h>
 #include <atlbase.h>
 #include <atlsecurity.h>
 #include <sddl.h>
@@ skipping 24 matching lines @@
 
 DllRedirector::DllRedirector(const char* shared_memory_name)
     : shared_memory_name_(shared_memory_name), first_module_handle_(NULL) {
   shared_memory_.reset(new base::SharedMemory(ASCIIToWide(shared_memory_name)));
 }
 
 DllRedirector::~DllRedirector() {
   if (first_module_handle_) {
     if (first_module_handle_ != reinterpret_cast<HMODULE>(&__ImageBase)) {
       FreeLibrary(first_module_handle_);
+    } else {
+      NOTREACHED() << "Error, DllRedirector attempting to free self.";
     }
+
     first_module_handle_ = NULL;
   }
   UnregisterAsFirstCFModule();
 }
 
 // static
 DllRedirector* DllRedirector::GetInstance() {
   return Singleton<DllRedirector>::get();
 }
 
@@ skipping 79 matching lines @@
     lock_acquired = shared_memory_->Lock(kSharedMemoryLockTimeoutMs, &sec_attr);
   } else {
     lock_acquired = shared_memory_->Lock(kSharedMemoryLockTimeoutMs, NULL);
   }
 
   if (!lock_acquired) {
     // We couldn't get the lock in a reasonable amount of time, so fall
     // back to loading our current version. We return true to indicate that the
     // caller should not attempt to delegate to an already loaded version.
     dll_version_.swap(our_version);
-    first_module_handle_ = reinterpret_cast<HMODULE>(&__ImageBase);
     return true;
   }
 
   bool created_beacon = true;
   bool result = shared_memory_->CreateNamed(shared_memory_name_.c_str(),
                                             false,  // open_existing
                                             kSharedMemorySize);
 
   if (result) {
     // We created the beacon, now we need to mutate the security attributes
@@ skipping 19 matching lines @@
     if (result) {
       // Either write our own version number or read it in if it was already
       // present in the shared memory section.
       if (created_beacon) {
         dll_version_.swap(our_version);
 
         lstrcpynA(reinterpret_cast<char*>(shared_memory_->memory()),
                   dll_version_->GetString().c_str(),
                   std::min(kSharedMemorySize,
                            dll_version_->GetString().length() + 1));
-
-        // Mark ourself as the first module in.
-        first_module_handle_ = reinterpret_cast<HMODULE>(&__ImageBase);
       } else {
         char buffer[kSharedMemorySize] = {0};
         memcpy(buffer, shared_memory_->memory(), kSharedMemorySize - 1);
         dll_version_.reset(Version::GetVersionFromString(buffer));
 
         if (!dll_version_.get() || dll_version_->Equals(*our_version.get())) {
           // If we either couldn't parse a valid version out of the shared
           // memory or we did parse a version and it is the same as our own,
           // then pretend we're first in to avoid trying to load any other DLLs.
           dll_version_.reset(our_version.release());
-          first_module_handle_ = reinterpret_cast<HMODULE>(&__ImageBase);
           created_beacon = true;
         }
       }
     } else {
       NOTREACHED() << "Failed to map in version beacon.";
     }
   } else {
     NOTREACHED() << "Could not create file mapping for version beacon, gle: "
                  << ::GetLastError();
   }
@@ skipping 12 matching lines @@
       // deleted.
       shared_memory_->Close();
       shared_memory_->Unlock();
     }
   }
 }
 
 LPFNGETCLASSOBJECT DllRedirector::GetDllGetClassObjectPtr() {
   HMODULE first_module_handle = GetFirstModule();
 
-  LPFNGETCLASSOBJECT proc_ptr = reinterpret_cast<LPFNGETCLASSOBJECT>(
-      GetProcAddress(first_module_handle, "DllGetClassObject"));
-  if (!proc_ptr) {
-    DPLOG(ERROR) << "DllRedirector: Could not get address of DllGetClassObject "
-                    "from first loaded module.";
-    // Oh boink, the first module we loaded was somehow bogus, make ourselves
-    // the first module again.
-    first_module_handle = reinterpret_cast<HMODULE>(&__ImageBase);
+  LPFNGETCLASSOBJECT proc_ptr = NULL;
+  if (first_module_handle) {
+    proc_ptr = reinterpret_cast<LPFNGETCLASSOBJECT>(
+        GetProcAddress(first_module_handle, "DllGetClassObject"));
+    DPLOG_IF(ERROR, !proc_ptr) << "DllRedirector: Could not get address of "
+                                  "DllGetClassObject from first loaded module.";
   }
+
   return proc_ptr;
 }
 
 Version* DllRedirector::GetCurrentModuleVersion() {
   scoped_ptr<FileVersionInfo> file_version_info(
       FileVersionInfo::CreateFileVersionInfoForCurrentModule());
   DCHECK(file_version_info.get());
 
   Version* current_version = NULL;
   if (file_version_info.get()) {
      current_version = Version::GetVersionFromString(
          WideToASCII(file_version_info->file_version()));
     DCHECK(current_version);
   }
 
   return current_version;
 }
 
 HMODULE DllRedirector::GetFirstModule() {
   DCHECK(dll_version_.get())
       << "Error: Did you call RegisterAsFirstCFModule() first?";
 
   if (first_module_handle_ == NULL) {
     first_module_handle_ = LoadVersionedModule(dll_version_.get());
-    if (!first_module_handle_) {
-      first_module_handle_ = reinterpret_cast<HMODULE>(&__ImageBase);
-    }
+  }
+
+  if (first_module_handle_ == reinterpret_cast<HMODULE>(&__ImageBase)) {
+    NOTREACHED() << "Should not be loading own version.";
+    first_module_handle_ = NULL;
   }
 
   return first_module_handle_;
 }
 
 HMODULE DllRedirector::LoadVersionedModule(Version* version) {
   DCHECK(version);
 
   FilePath module_path;
   PathService::Get(base::FILE_MODULE, &module_path);
   DCHECK(!module_path.empty());
 
   // For a module located in
   // Foo\XXXXXXXXX\<module>.dll, load
   // Foo\<version>\<module>.dll:
   FilePath module_name = module_path.BaseName();
   module_path = module_path.DirName()
                            .DirName()
                            .Append(ASCIIToWide(version->GetString()))
                            .Append(module_name);
 
   HMODULE hmodule = LoadLibrary(module_path.value().c_str());
   if (hmodule == NULL) {
     DPLOG(ERROR) << "Could not load reported module version "
                  << version->GetString();
   }
 
   return hmodule;
 }