Bug 1395329: Modify mscom::Interceptor::Create to properly initialize IUnknown interfaces; r=jimm

MozReview-Commit-ID: 7ZLzpH030Of

--HG--
extra : rebase_source : 24fd48eb72bcee18bbc75a670d0035362dbeb92d

ipc/mscom/Interceptor.cpp

@@ -305,6 +305,7 @@ Interceptor::MapEntry*
 Interceptor::Lookup(REFIID aIid)
 {
   mMutex.AssertCurrentThreadOwns();
+
   for (uint32_t index = 0, len = mInterceptorMap.Length(); index < len; ++index) {
     if (mInterceptorMap[index].mIID == aIid) {
       return &mInterceptorMap[index];
@@ -367,15 +368,57 @@ Interceptor::CreateInterceptor(REFIID aIid, IUnknown* aOuter, IUnknown** aOutput
 }
 
 HRESULT
-Interceptor::GetInitialInterceptorForIID(detail::LiveSetAutoLock& aLock,
+Interceptor::PublishTarget(detail::LiveSetAutoLock& aLiveSetLock,
+                           RefPtr<IUnknown> aInterceptor,
+                           REFIID aTargetIid,
+                           STAUniquePtr<IUnknown> aTarget)
+{
+  RefPtr<IWeakReference> weakRef;
+  HRESULT hr = GetWeakReference(getter_AddRefs(weakRef));
+  if (FAILED(hr)) {
+    return hr;
+  }
+
+  // mTarget is a weak reference to aTarget. This is safe because we transfer
+  // ownership of aTarget into mInterceptorMap which remains live for the
+  // lifetime of this Interceptor.
+  mTarget = ToInterceptorTargetPtr(aTarget);
+  GetLiveSet().Put(mTarget.get(), weakRef.forget());
+
+  // Now we transfer aTarget's ownership into mInterceptorMap.
+  mInterceptorMap.AppendElement(MapEntry(aTargetIid,
+                                         aInterceptor,
+                                         aTarget.release()));
+
+  // Release the live set lock because subsequent operations may post work to
+  // the main thread, creating potential for deadlocks.
+  aLiveSetLock.Unlock();
+  return S_OK;
+}
+
+HRESULT
+Interceptor::GetInitialInterceptorForIID(detail::LiveSetAutoLock& aLiveSetLock,
                                          REFIID aTargetIid,
                                          STAUniquePtr<IUnknown> aTarget,
                                          void** aOutInterceptor)
 {
   MOZ_ASSERT(aOutInterceptor);
-  MOZ_ASSERT(aTargetIid != IID_IUnknown && aTargetIid != IID_IMarshal);
+  MOZ_ASSERT(aTargetIid != IID_IMarshal);
   MOZ_ASSERT(!IsProxy(aTarget.get()));
 
+  if (aTargetIid == IID_IUnknown) {
+    // We must lock ourselves so that nothing can race with us once we have been
+    // published to the live set.
+    AutoLock lock(*this);
+
+    HRESULT hr = PublishTarget(aLiveSetLock, nullptr, aTargetIid, Move(aTarget));
+    if (FAILED(hr)) {
+      return hr;
+    }
+
+    return QueryInterface(aTargetIid, aOutInterceptor);
+  }
+
   // Raise the refcount for stabilization purposes during aggregation
   RefPtr<IUnknown> kungFuDeathGrip(static_cast<IUnknown*>(
         static_cast<WeakReferenceSupport*>(this)));
@@ -399,27 +442,15 @@ Interceptor::GetInitialInterceptorForIID(detail::LiveSetAutoLock& aLock,
     return hr;
   }
 
-  RefPtr<IWeakReference> weakRef;
-  hr = GetWeakReference(getter_AddRefs(weakRef));
+  // We must lock ourselves so that nothing can race with us once we have been
+  // published to the live set.
+  AutoLock lock(*this);
+
+  hr = PublishTarget(aLiveSetLock, unkInterceptor, aTargetIid, Move(aTarget));
   if (FAILED(hr)) {
     return hr;
   }
 
-  // mTarget is a weak reference to aTarget. This is safe because we transfer
-  // ownership of aTarget into mInterceptorMap which remains live for the
-  // lifetime of this Interceptor.
-  mTarget = ToInterceptorTargetPtr(aTarget);
-  GetLiveSet().Put(mTarget.get(), weakRef.forget());
-
-  // Release the live set lock because GetInterceptorForIID will post work to
-  // the main thread, creating potential for deadlocks.
-  aLock.Unlock();
-
-  // Now we transfer aTarget's ownership into mInterceptorMap.
-  mInterceptorMap.AppendElement(MapEntry(aTargetIid,
-                                         unkInterceptor,
-                                         aTarget.release()));
-
   if (mEventSink->MarshalAs(aTargetIid) == aTargetIid) {
     return unkInterceptor->QueryInterface(aTargetIid, aOutInterceptor);
   }

ipc/mscom/Interceptor.h

@@ -120,7 +120,7 @@ private:
 private:
   explicit Interceptor(IInterceptorSink* aSink);
   ~Interceptor();
-  HRESULT GetInitialInterceptorForIID(detail::LiveSetAutoLock& aLock,
+  HRESULT GetInitialInterceptorForIID(detail::LiveSetAutoLock& aLiveSetLock,
                                       REFIID aTargetIid,
                                       STAUniquePtr<IUnknown> aTarget,
                                       void** aOutInterface);
@@ -129,6 +129,10 @@ private:
   HRESULT ThreadSafeQueryInterface(REFIID aIid,
                                    IUnknown** aOutInterface) override;
   HRESULT CreateInterceptor(REFIID aIid, IUnknown* aOuter, IUnknown** aOutput);
+  HRESULT PublishTarget(detail::LiveSetAutoLock& aLiveSetLock,
+                        RefPtr<IUnknown> aInterceptor,
+                        REFIID aTargetIid,
+                        STAUniquePtr<IUnknown> aTarget);
 
 private:
   InterceptorTargetPtr<IUnknown>  mTarget;

ipc/mscom/WeakRef.cpp

@@ -109,6 +109,18 @@ WeakReferenceSupport::~WeakReferenceSupport()
   ::DeleteCriticalSection(&mCSForQI);
 }
 
+void
+WeakReferenceSupport::Lock()
+{
+  ::EnterCriticalSection(&mCSForQI);
+}
+
+void
+WeakReferenceSupport::Unlock()
+{
+  ::LeaveCriticalSection(&mCSForQI);
+}
+
 HRESULT
 WeakReferenceSupport::QueryInterface(REFIID riid, void** ppv)
 {

ipc/mscom/WeakRef.h

@@ -12,6 +12,7 @@
 
 #include "mozilla/Assertions.h"
 #include "mozilla/Atomics.h"
+#include "mozilla/Mutex.h"
 #include "mozilla/RefPtr.h"
 #include "nsISupportsImpl.h"
 
@@ -101,6 +102,12 @@ protected:
   virtual HRESULT ThreadSafeQueryInterface(REFIID aIid,
                                            IUnknown** aOutInterface) = 0;
 
+  void Lock();
+  void Unlock();
+
+  typedef BaseAutoLock<WeakReferenceSupport> AutoLock;
+  friend class AutoLock;
+
 private:
   RefPtr<detail::SharedRef> mSharedRef;
   ULONG                     mRefCnt;