Bug 1409545: Add the ability to query a handler provider for the expected interface for an IUnknown outparam; r=jimm

MozReview-Commit-ID: 7eZ2giTfVG1

--HG--
extra : rebase_source : cc36fbd875ef798c41f20951b8354880138c0037

ipc/mscom/IHandlerProvider.h

@@ -26,6 +26,8 @@ struct HandlerProvider
 struct IHandlerProvider : public IUnknown
                         , public HandlerProvider
 {
+  virtual STDMETHODIMP_(REFIID) GetEffectiveOutParamIid(REFIID aCallIid,
+                                                        ULONG aCallMethod) = 0;
   virtual STDMETHODIMP NewInstance(REFIID aIid,
                                    InterceptorTargetPtr<IUnknown> aTarget,
                                    NotNull<IHandlerProvider**> aOutNewPayload) = 0;

ipc/mscom/MainThreadHandoff.cpp

@@ -14,6 +14,7 @@
 #include "mozilla/mscom/Utils.h"
 #include "mozilla/Assertions.h"
 #include "mozilla/DebugOnly.h"
+#include "mozilla/ThreadLocal.h"
 #include "nsThreadUtils.h"
 #include "nsProxyRelease.h"
 
@@ -154,6 +155,52 @@ private:
   HRESULT       mResult;
 };
 
+class MOZ_RAII SavedCallFrame final
+{
+public:
+  explicit SavedCallFrame(mozilla::NotNull<ICallFrame*> aFrame)
+    : mCallFrame(aFrame)
+  {
+    static const bool sIsInit = tlsFrame.init();
+    MOZ_ASSERT(sIsInit);
+    MOZ_ASSERT(!tlsFrame.get());
+    tlsFrame.set(this);
+  }
+
+  ~SavedCallFrame()
+  {
+    MOZ_ASSERT(tlsFrame.get());
+    tlsFrame.set(nullptr);
+  }
+
+  HRESULT GetIidAndMethod(mozilla::NotNull<IID*> aIid,
+                          mozilla::NotNull<ULONG*> aMethod) const
+  {
+    return mCallFrame->GetIIDAndMethod(aIid, aMethod);
+  }
+
+  static const SavedCallFrame& Get()
+  {
+    SavedCallFrame* saved = tlsFrame.get();
+    MOZ_ASSERT(saved);
+
+    return *saved;
+  }
+
+  SavedCallFrame(const SavedCallFrame&) = delete;
+  SavedCallFrame(SavedCallFrame&&) = delete;
+  SavedCallFrame& operator=(const SavedCallFrame&) = delete;
+  SavedCallFrame& operator=(SavedCallFrame&&) = delete;
+
+private:
+  ICallFrame* mCallFrame;
+
+private:
+  static MOZ_THREAD_LOCAL(SavedCallFrame*) tlsFrame;
+};
+
+MOZ_THREAD_LOCAL(SavedCallFrame*) SavedCallFrame::tlsFrame;
+
 } // anonymous namespace
 
 namespace mozilla {
@@ -338,6 +385,8 @@ MainThreadHandoff::OnCall(ICallFrame* aFrame)
       return hr;
     }
   } else {
+    SavedCallFrame savedFrame(WrapNotNull(aFrame));
+
     // (7) Scan the outputs looking for any outparam interfaces that need wrapping.
     // NB: WalkFrame does not correctly handle array outparams. It processes the
     // first element of an array but not the remaining elements (if any).
@@ -566,9 +615,26 @@ MainThreadHandoff::OnWalkInterface(REFIID aIid, PVOID* aInterface,
     }
   }
 
+  IID effectiveIid = aIid;
+
   RefPtr<IHandlerProvider> payload;
   if (mHandlerProvider) {
-    hr = mHandlerProvider->NewInstance(aIid,
+    if (aIid == IID_IUnknown) {
+      const SavedCallFrame& curFrame = SavedCallFrame::Get();
+
+      IID callIid;
+      ULONG callMethod;
+      hr = curFrame.GetIidAndMethod(WrapNotNull(&callIid),
+                                    WrapNotNull(&callMethod));
+      if (FAILED(hr)) {
+        return hr;
+      }
+
+      effectiveIid = mHandlerProvider->GetEffectiveOutParamIid(callIid,
+                                                               callMethod);
+    }
+
+    hr = mHandlerProvider->NewInstance(effectiveIid,
                                        ToInterceptorTargetPtr(origInterface),
                                        WrapNotNull((IHandlerProvider**)getter_AddRefs(payload)));
     MOZ_ASSERT(SUCCEEDED(hr));
@@ -585,7 +651,8 @@ MainThreadHandoff::OnWalkInterface(REFIID aIid, PVOID* aInterface,
     return hr;
   }
 
-  REFIID interceptorIid = payload ? payload->MarshalAs(aIid) : aIid;
+  REFIID interceptorIid = payload ? payload->MarshalAs(effectiveIid) :
+                                    effectiveIid;
 
   RefPtr<IUnknown> wrapped;
   hr = Interceptor::Create(Move(origInterface), handoff, interceptorIid,