This fixes a crash in IE8 with ChromeFrame when a new tab was created. ...

chrome_frame/protocol_sink_wrap.cc

 // Copyright (c) 2009 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 <htiframe.h>
 
 #include "chrome_frame/protocol_sink_wrap.h"
 
 #include "base/scoped_bstr_win.h"
 #include "base/logging.h"
 #include "base/registry.h"
 #include "base/scoped_bstr_win.h"
 #include "base/singleton.h"
 #include "base/string_util.h"
 
 #include "chrome_frame/utils.h"
-#include "chrome_frame/vtable_patch_manager.h"
 
-// BINDSTATUS_SERVER_MIMETYPEAVAILABLE == 54. Introduced in IE 8, so 
+// BINDSTATUS_SERVER_MIMETYPEAVAILABLE == 54. Introduced in IE 8, so
 // not in everyone's headers yet. See:
 // http://msdn.microsoft.com/en-us/library/ms775133(VS.85,loband).aspx
 #ifndef BINDSTATUS_SERVER_MIMETYPEAVAILABLE
 #define BINDSTATUS_SERVER_MIMETYPEAVAILABLE 54
 #endif
 
 static const wchar_t* kChromeMimeType = L"application/chromepage";
 static const char kTextHtmlMimeType[] = "text/html";
+const wchar_t kUrlMonDllName[] = L"urlmon.dll";
 
 static const int kInternetProtocolStartIndex = 3;
 static const int kInternetProtocolReadIndex = 9;
 static const int kInternetProtocolStartExIndex = 13;
 
+// TODO(ananta)
+// We should avoid duplicate VTable declarations.
 BEGIN_VTABLE_PATCHES(IInternetProtocol)
   VTABLE_PATCH_ENTRY(kInternetProtocolStartIndex, ProtocolSinkWrap::OnStart)
   VTABLE_PATCH_ENTRY(kInternetProtocolReadIndex, ProtocolSinkWrap::OnRead)
 END_VTABLE_PATCHES()
 
+BEGIN_VTABLE_PATCHES(IInternetProtocolSecure)
+  VTABLE_PATCH_ENTRY(kInternetProtocolStartIndex, ProtocolSinkWrap::OnStart)
+  VTABLE_PATCH_ENTRY(kInternetProtocolReadIndex, ProtocolSinkWrap::OnRead)
+END_VTABLE_PATCHES()
+
 BEGIN_VTABLE_PATCHES(IInternetProtocolEx)
   VTABLE_PATCH_ENTRY(kInternetProtocolStartIndex, ProtocolSinkWrap::OnStart)
   VTABLE_PATCH_ENTRY(kInternetProtocolReadIndex, ProtocolSinkWrap::OnRead)
   VTABLE_PATCH_ENTRY(kInternetProtocolStartExIndex, ProtocolSinkWrap::OnStartEx)
 END_VTABLE_PATCHES()
 
+BEGIN_VTABLE_PATCHES(IInternetProtocolExSecure)
+  VTABLE_PATCH_ENTRY(kInternetProtocolStartIndex, ProtocolSinkWrap::OnStart)
+  VTABLE_PATCH_ENTRY(kInternetProtocolReadIndex, ProtocolSinkWrap::OnRead)
+  VTABLE_PATCH_ENTRY(kInternetProtocolStartExIndex, ProtocolSinkWrap::OnStartEx)
+END_VTABLE_PATCHES()
+
 //
 // ProtocolSinkWrap implementation
 //
 
 // Static map initialization
 ProtocolSinkWrap::ProtocolSinkMap ProtocolSinkWrap::sink_map_;
 CComAutoCriticalSection ProtocolSinkWrap::sink_map_lock_;
 
 ProtocolSinkWrap::ProtocolSinkWrap()
     : protocol_(NULL), renderer_type_(UNDETERMINED),
       buffer_size_(0), buffer_pos_(0), is_saved_result_(false),
       result_code_(0), result_error_(0), report_data_recursiveness_(0) {
   memset(buffer_, 0, arraysize(buffer_));
 }
 
 ProtocolSinkWrap::~ProtocolSinkWrap() {
   CComCritSecLock<CComAutoCriticalSection> lock(sink_map_lock_);
   DCHECK(sink_map_.end() != sink_map_.find(protocol_));
   sink_map_.erase(protocol_);
   protocol_ = NULL;
   DLOG(INFO) << "ProtocolSinkWrap: active sinks: " << sink_map_.size();
 }
 
-bool ProtocolSinkWrap::PatchProtocolHandler(const wchar_t* dll,
-    const CLSID& handler_clsid) {
-  HMODULE module = ::GetModuleHandle(dll);
-  if (!module) {
-    NOTREACHED() << "urlmon is not yet loaded. Error: " << GetLastError();
+bool ProtocolSinkWrap::PatchProtocolHandlers() {
+  HRESULT hr = PatchProtocolMethods(CLSID_HttpProtocol,
+                                    IInternetProtocol_PatchInfo,
+                                    IInternetProtocolEx_PatchInfo);
+  if (FAILED(hr)) {
+    NOTREACHED() << "Failed to patch IInternetProtocol interface."
+        << " Error: " << hr;
     return false;
   }
 
+  hr = PatchProtocolMethods(CLSID_HttpSProtocol,
+                            IInternetProtocolSecure_PatchInfo,
+                            IInternetProtocolExSecure_PatchInfo);
+  if (FAILED(hr)) {
+    NOTREACHED() << "Failed to patch IInternetProtocol secure interface."
+        << " Error: " << hr;
+    return false;
+  }
+
+  return true;
+}
+
+void ProtocolSinkWrap::UnpatchProtocolHandlers() {
+  vtable_patch::UnpatchInterfaceMethods(IInternetProtocol_PatchInfo);
+  vtable_patch::UnpatchInterfaceMethods(IInternetProtocolEx_PatchInfo);
+  vtable_patch::UnpatchInterfaceMethods(IInternetProtocolSecure_PatchInfo);
+  vtable_patch::UnpatchInterfaceMethods(IInternetProtocolExSecure_PatchInfo);
+}
+
+HRESULT ProtocolSinkWrap::CreateProtocolHandlerInstance(
+    const CLSID& clsid, IInternetProtocol** protocol) {
+  if (!protocol) {
+    return E_INVALIDARG;
+  }
+
+  HMODULE module = ::GetModuleHandle(kUrlMonDllName);
+  if (!module) {
+    NOTREACHED() << "urlmon is not yet loaded. Error: " << GetLastError();
+    return E_FAIL;
+  }
+
   typedef HRESULT (WINAPI* DllGetClassObject_Fn)(REFCLSID, REFIID, LPVOID*);
   DllGetClassObject_Fn fn = reinterpret_cast<DllGetClassObject_Fn>(
       ::GetProcAddress(module, "DllGetClassObject"));
   if (!fn) {
     NOTREACHED() << "DllGetClassObject not found in urlmon.dll";
-    return false;
+    return E_FAIL;
   }
 
   ScopedComPtr<IClassFactory> protocol_class_factory;
-  HRESULT hr = fn(handler_clsid, IID_IClassFactory,
+  HRESULT hr = fn(clsid, IID_IClassFactory,
       reinterpret_cast<LPVOID*>(protocol_class_factory.Receive()));
   if (FAILED(hr)) {
     NOTREACHED() << "DllGetclassObject failed. Error: " << hr;
-    return false;
+    return hr;
   }
 
   ScopedComPtr<IInternetProtocol> handler_instance;
   hr = protocol_class_factory->CreateInstance(NULL, IID_IInternetProtocol,
       reinterpret_cast<void**>(handler_instance.Receive()));
   if (FAILED(hr)) {
     NOTREACHED() << "ClassFactory::CreateInstance failed for InternetProtocol."
         << " Error: " << hr;
+  } else {
+    *protocol = handler_instance.Detach();
+  }
+
+  return hr;
+}
+
+HRESULT ProtocolSinkWrap::PatchProtocolMethods(
+    const CLSID& clsid_protocol,
+    vtable_patch::MethodPatchInfo* protocol_patch_info,
+    vtable_patch::MethodPatchInfo* protocol_ex_patch_info) {
+  if (!protocol_patch_info || !protocol_ex_patch_info) {
+    return E_INVALIDARG;
+  }
+
+  ScopedComPtr<IInternetProtocol> http_protocol;
+  HRESULT hr = CreateProtocolHandlerInstance(clsid_protocol,
+                                             http_protocol.Receive());
+  if (FAILED(hr)) {
+    NOTREACHED() << "ClassFactory::CreateInstance failed for InternetProtocol."
+        << " Error: " << hr;
     return false;
   }
 
   ScopedComPtr<IInternetProtocolEx> ipex;
-  ipex.QueryFrom(handler_instance);
+  ipex.QueryFrom(http_protocol);
   if (ipex) {
-    vtable_patch::PatchInterfaceMethods(ipex, IInternetProtocolEx_PatchInfo);
+    hr = vtable_patch::PatchInterfaceMethods(ipex, protocol_ex_patch_info);
   } else {
-    vtable_patch::PatchInterfaceMethods(handler_instance,
-                                        IInternetProtocol_PatchInfo);
+    hr = vtable_patch::PatchInterfaceMethods(http_protocol,
+                                             protocol_patch_info);
   }
-
-  return true;
+  return hr;
 }
 
 // IInternetProtocol/Ex method implementation.
 HRESULT ProtocolSinkWrap::OnStart(InternetProtocol_Start_Fn orig_start,
     IInternetProtocol* protocol, LPCWSTR url, IInternetProtocolSink* prot_sink,
     IInternetBindInfo* bind_info, DWORD flags, HANDLE_PTR reserved) {
   DCHECK(orig_start);
   DLOG_IF(INFO, url != NULL) << "OnStart: " << url;
 
   ScopedComPtr<IInternetProtocolSink> sink_to_use(MaybeWrapSink(protocol,
@@ skipping 137 matching lines @@
   //                 |--> orig_prot->Read(...) - 2nd read S_FALSE, 0 bytes
   //
   // Inner call returns S_FALSE and no data. We try to make a determination
   // of render type then and incorrectly set it to 'OTHER' as we don't have
   // any data yet. However, we can make a determination in the context of
   // outer ReportData since the first read will return S_OK with data. Then
   // the next Read in the loop will return S_FALSE and we will enter the
   // determination logic.
 
   // NOTE: We use the report_data_recursiveness_ variable to detect situations
-  //    in which calls to ReportData are re-entrant (such as when the entire 
-  //    contents of a page fit inside a single packet). In these cases, we 
-  //    don't care about re-entrant calls beyond the second, and so we compare 
-  //    report_data_recursiveness_ inside the while loop, making sure we skip 
+  //    in which calls to ReportData are re-entrant (such as when the entire
+  //    contents of a page fit inside a single packet). In these cases, we
+  //    don't care about re-entrant calls beyond the second, and so we compare
+  //    report_data_recursiveness_ inside the while loop, making sure we skip
   //    what would otherwise be spurious calls to ReportProgress().
   report_data_recursiveness_++;
 
   HRESULT hr = S_OK;
   if (is_undetermined()) {
     HRESULT hr_read = S_OK;
     while (hr_read == S_OK) {
       ULONG size_read = 0;
       hr_read = protocol_->Read(buffer_ + buffer_size_,
           kMaxContentSniffLength - buffer_size_, &size_read);
@@ skipping 325 matching lines @@
     CComObject<ProtocolSinkWrap>* wrap = NULL;
     CComObject<ProtocolSinkWrap>::CreateInstance(&wrap);
     DCHECK(wrap);
     if (wrap->Initialize(protocol, prot_sink, url)) {
       sink_to_use = wrap;
     }
   }
 
   return sink_to_use;
 }