2 files changed, 68 insertions, 9 deletions

diff --git a/chrome_frame/bind_context_info.cc b/chrome_frame/bind_context_info.cc
index 3eb2554..07ba1f3 100644
--- a/chrome_frame/bind_context_info.cc
+++ b/chrome_frame/bind_context_info.cc
@@ -15,6 +15,22 @@ BindContextInfo::BindContextInfo()
       is_switching_(false) {
 }
 
+BindContextInfo::~BindContextInfo() {
+}
+
+HRESULT BindContextInfo::Initialize(IBindCtx* bind_ctx) {
+  DCHECK(bind_ctx);
+  DCHECK(!ftm_);
+  HRESULT hr = CoCreateFreeThreadedMarshaler(GetUnknown(), ftm_.Receive());
+  DCHECK(ftm_);
+  if (SUCCEEDED(hr)) {
+    hr = bind_ctx->RegisterObjectParam(kBindContextInfo, GetUnknown());
+  }
+
+  DCHECK(SUCCEEDED(hr)) << "Failed to initialize BindContextInfo";
+  return hr;
+}
+
 BindContextInfo* BindContextInfo::FromBindContext(IBindCtx* bind_context) {
   if (!bind_context) {
     NOTREACHED();
@@ -24,14 +40,31 @@ BindContextInfo* BindContextInfo::FromBindContext(IBindCtx* bind_context) {
   ScopedComPtr<IUnknown> context;
   bind_context->GetObjectParam(kBindContextInfo, context.Receive());
   if (context) {
-    return static_cast<BindContextInfo*>(context.get());
+    ScopedComPtr<IBindContextInfoInternal> internal;
+    HRESULT hr = internal.QueryFrom(context);
+    DCHECK(SUCCEEDED(hr));
+    if (SUCCEEDED(hr)) {
+      BindContextInfo* ret = NULL;
+      internal->GetCppObject(reinterpret_cast<void**>(&ret));
+      DCHECK(ret);
+      DLOG_IF(ERROR, reinterpret_cast<void*>(ret) !=
+                     reinterpret_cast<void*>(internal.get()))
+          << "marshalling took place!";
+      return ret;
+    }
   }
 
   CComObject<BindContextInfo>* bind_context_info = NULL;
-  CComObject<BindContextInfo>::CreateInstance(&bind_context_info);
+  HRESULT hr = CComObject<BindContextInfo>::CreateInstance(&bind_context_info);
   DCHECK(bind_context_info != NULL);
+  if (bind_context_info) {
+    bind_context_info->AddRef();
+    hr = bind_context_info->Initialize(bind_context);
+    bind_context_info->Release();
+    if (FAILED(hr))
+      bind_context_info = NULL;
+  }
 
-  bind_context->RegisterObjectParam(kBindContextInfo, bind_context_info);
   return bind_context_info;
 }
 
diff --git a/chrome_frame/bind_context_info.h b/chrome_frame/bind_context_info.h
index b6dca88..d41c440 100644
--- a/chrome_frame/bind_context_info.h
+++ b/chrome_frame/bind_context_info.h
@@ -11,15 +11,24 @@
 #include "base/scoped_bstr_win.h"
 #include "base/scoped_comptr_win.h"
 
+class __declspec(uuid("71CC3EC7-7E8A-457f-93BC-1090CF31CC18"))
+IBindContextInfoInternal : public IUnknown {
+ public:
+  STDMETHOD(GetCppObject)(void** me) = 0;
+};
+
 // This class maintains contextual information used by ChromeFrame.
 // This information is maintained in the bind context.
-class BindContextInfo : public IUnknown, public CComObjectRoot {
+class BindContextInfo
+  : public IBindContextInfoInternal,
+    public CComObjectRoot {
  public:
   BindContextInfo();
-  virtual ~BindContextInfo() {}
+  ~BindContextInfo();
 
   BEGIN_COM_MAP(BindContextInfo)
-    COM_INTERFACE_ENTRY(IUnknown)
+    COM_INTERFACE_ENTRY(IBindContextInfoInternal)
+    COM_INTERFACE_ENTRY_AGGREGATE(IID_IMarshal, ftm_)
   END_COM_MAP()
 
   // Returns the BindContextInfo instance associated with the bind
@@ -52,20 +61,37 @@ class BindContextInfo : public IUnknown, public CComObjectRoot {
     return cache_;
   }
 
-  void set_url(const std::wstring& url) {
-    url_ = url;
+  // Accept a const wchar_t* to ensure that we don't have a reference
+  // to someone else's buffer.
+  void set_url(const wchar_t* url) {
+    DCHECK(url);
+    if (url) {
+      url_ = url;
+    } else {
+      url_.clear();
+    }
   }
 
-  const std::wstring url() const {
+  const std::wstring& url() const {
     return url_;
   }
 
+ protected:
+  STDMETHOD(GetCppObject)(void** me) {
+    DCHECK(me);
+    *me = static_cast<BindContextInfo*>(this);
+    return S_OK;
+  }
+
+  HRESULT Initialize(IBindCtx* bind_ctx);
+
  private:
   ScopedComPtr<IStream> cache_;
   bool no_cache_;
   bool chrome_request_;
   bool is_switching_;
   std::wstring url_;
+  ScopedComPtr<IUnknown> ftm_;
 
   DISALLOW_COPY_AND_ASSIGN(BindContextInfo);
 };