Bug 1323069: mscom changes to facilitate resolution of remote client thread ID; r=jimm

MozReview-Commit-ID: 9x9wZr7sRFI

ipc/mscom/MainThreadClientInfo.cpp

@@ -0,0 +1,146 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "MainThreadClientInfo.h"
+
+#include "MainThreadUtils.h"
+#include "mozilla/Assertions.h"
+
+#include <objbase.h>
+
+namespace mozilla {
+namespace mscom {
+
+/* static */
+HRESULT
+MainThreadClientInfo::Create(MainThreadClientInfo** aOutObj)
+{
+  MOZ_ASSERT(aOutObj && NS_IsMainThread());
+  *aOutObj = nullptr;
+
+  RefPtr<MainThreadClientInfo> obj(new MainThreadClientInfo());
+
+  RefPtr<IMessageFilter> prevFilter;
+  HRESULT hr = ::CoRegisterMessageFilter(obj.get(),
+                                         getter_AddRefs(prevFilter));
+  if (FAILED(hr)) {
+    return hr;
+  }
+
+  obj->mPrevFilter = prevFilter.forget();
+
+  obj.forget(aOutObj);
+  return S_OK;
+}
+
+DWORD
+MainThreadClientInfo::GetLastRemoteCallThreadId() const
+{
+  MOZ_ASSERT(NS_IsMainThread());
+  return mLastRemoteCallTid;
+}
+
+HRESULT
+MainThreadClientInfo::QueryInterface(REFIID aIid, void** aOutInterface)
+{
+  MOZ_ASSERT(NS_IsMainThread());
+
+  if (!aOutInterface) {
+    return E_INVALIDARG;
+  }
+
+  if (aIid == IID_IUnknown || aIid == IID_IMessageFilter) {
+    RefPtr<IMessageFilter> filter(this);
+    filter.forget(aOutInterface);
+    return S_OK;
+  }
+
+  return E_NOINTERFACE;
+}
+
+ULONG
+MainThreadClientInfo::AddRef()
+{
+  MOZ_ASSERT(NS_IsMainThread());
+
+  return ++mRefCnt;
+}
+
+ULONG
+MainThreadClientInfo::Release()
+{
+  MOZ_ASSERT(NS_IsMainThread());
+
+  ULONG newCount = --mRefCnt;
+  if (!newCount) {
+    delete this;
+  }
+  return newCount;
+}
+
+DWORD
+MainThreadClientInfo::HandleInComingCall(DWORD aCallType, HTASK aCallerTid,
+                                         DWORD aTickCount,
+                                         LPINTERFACEINFO aInterfaceInfo)
+{
+  MOZ_ASSERT(NS_IsMainThread());
+
+  // aCallerTid is an HTASK for historical reasons but is actually just a
+  // regular DWORD Thread ID.
+  mLastRemoteCallTid =
+    static_cast<DWORD>(reinterpret_cast<uintptr_t>(aCallerTid));
+
+  if (!mPrevFilter) {
+    return SERVERCALL_ISHANDLED;
+  }
+
+  return mPrevFilter->HandleInComingCall(aCallType, aCallerTid, aTickCount,
+                                         aInterfaceInfo);
+}
+
+DWORD
+MainThreadClientInfo::RetryRejectedCall(HTASK aCalleeTid, DWORD aTickCount,
+                                        DWORD aRejectType)
+{
+  MOZ_ASSERT(NS_IsMainThread());
+
+  if (!mPrevFilter) {
+    return 0;
+  }
+
+  return mPrevFilter->RetryRejectedCall(aCalleeTid, aTickCount, aRejectType);
+}
+
+DWORD
+MainThreadClientInfo::MessagePending(HTASK aCalleeTid, DWORD aTickCount,
+                                     DWORD aPendingType)
+{
+  MOZ_ASSERT(NS_IsMainThread());
+
+  if (!mPrevFilter) {
+    return PENDINGMSG_WAITNOPROCESS;
+  }
+
+  return mPrevFilter->MessagePending(aCalleeTid, aTickCount, aPendingType);
+}
+
+MainThreadClientInfo::MainThreadClientInfo()
+  : mRefCnt(0)
+  , mLastRemoteCallTid(0)
+{
+  MOZ_ASSERT(NS_IsMainThread());
+}
+
+void
+MainThreadClientInfo::Detach()
+{
+  MOZ_ASSERT(NS_IsMainThread());
+  ::CoRegisterMessageFilter(mPrevFilter, nullptr);
+  mPrevFilter = nullptr;
+}
+
+} // namespace mscom
+} // namespace mozilla

ipc/mscom/MainThreadClientInfo.h

@@ -0,0 +1,52 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#ifndef mozilla_mscom_MainThreadClientInfo
+#define mozilla_mscom_MainThreadClientInfo
+
+#include "mozilla/RefPtr.h"
+#include "nsString.h"
+
+#include <objidl.h>
+
+namespace mozilla {
+namespace mscom {
+
+class MainThreadClientInfo final : public IMessageFilter
+{
+public:
+  static HRESULT Create(MainThreadClientInfo** aOutObj);
+
+  DWORD GetLastRemoteCallThreadId() const;
+  void Detach();
+
+  STDMETHODIMP QueryInterface(REFIID aIid, void** aOutInterface) override;
+  STDMETHODIMP_(ULONG) AddRef() override;
+  STDMETHODIMP_(ULONG) Release() override;
+
+  STDMETHODIMP_(DWORD) HandleInComingCall(DWORD aCallType, HTASK aCallerTid,
+                                          DWORD aTickCount,
+                                          LPINTERFACEINFO aInterfaceInfo) override;
+  STDMETHODIMP_(DWORD) RetryRejectedCall(HTASK aCalleeTid, DWORD aTickCount,
+                                         DWORD aRejectType) override;
+  STDMETHODIMP_(DWORD) MessagePending(HTASK aCalleeTid, DWORD aTickCount,
+                                      DWORD aPendingType) override;
+
+private:
+  MainThreadClientInfo();
+  ~MainThreadClientInfo() = default;
+
+private:
+  ULONG mRefCnt;
+  RefPtr<IMessageFilter> mPrevFilter;
+  DWORD mLastRemoteCallTid;
+};
+
+} // namespace mscom
+} // namespace mozilla
+
+#endif // mozilla_mscom_MainThreadClientInfo
+

ipc/mscom/MainThreadRuntime.cpp

@@ -12,6 +12,7 @@
 #include "mozilla/UniquePtr.h"
 #include "nsDebug.h"
 #include "nsWindowsHelpers.h"
+#include "nsXULAppAPI.h"
 
 #include <accctrl.h>
 #include <aclapi.h>
@@ -36,6 +37,8 @@ extern "C" void __cdecl SetOaNoCache(void);
 namespace mozilla {
 namespace mscom {
 
+MainThreadRuntime* MainThreadRuntime::sInstance = nullptr;
+
 MainThreadRuntime::MainThreadRuntime()
   : mInitResult(E_UNEXPECTED)
 {
@@ -69,6 +72,61 @@ MainThreadRuntime::MainThreadRuntime()
 
   // Disable the BSTR cache (as it never invalidates, thus leaking memory)
   ::SetOaNoCache();
+
+  if (FAILED(mInitResult)) {
+    return;
+  }
+
+  if (XRE_IsParentProcess()) {
+    MainThreadClientInfo::Create(getter_AddRefs(mClientInfo));
+  }
+
+  MOZ_ASSERT(!sInstance);
+  sInstance = this;
+}
+
+MainThreadRuntime::~MainThreadRuntime()
+{
+  if (mClientInfo) {
+    mClientInfo->Detach();
+  }
+
+  MOZ_ASSERT(sInstance == this);
+  if (sInstance == this) {
+    sInstance = nullptr;
+  }
+}
+
+/* static */
+DWORD
+MainThreadRuntime::GetClientThreadId()
+{
+  MOZ_ASSERT(NS_IsMainThread());
+  MOZ_ASSERT(XRE_IsParentProcess(), "Unsupported outside of parent process");
+  if (!XRE_IsParentProcess()) {
+    return 0;
+  }
+
+  // Don't check for a calling executable if the caller is in-process.
+  // We verify this by asking COM for a call context. If none exists, then
+  // we must be a local call.
+  RefPtr<IServerSecurity> serverSecurity;
+  if (FAILED(::CoGetCallContext(IID_IServerSecurity,
+                                getter_AddRefs(serverSecurity)))) {
+    return 0;
+  }
+
+  MOZ_ASSERT(sInstance);
+  if (!sInstance) {
+    return 0;
+  }
+
+  MOZ_ASSERT(sInstance->mClientInfo);
+  if (!sInstance->mClientInfo) {
+    return 0;
+  }
+
+  return sInstance->mClientInfo->GetLastRemoteCallThreadId();
 }
 
 HRESULT

ipc/mscom/MainThreadRuntime.h

@@ -9,6 +9,8 @@
 
 #include "mozilla/Attributes.h"
 #include "mozilla/mscom/COMApartmentRegion.h"
+#include "mozilla/mscom/MainThreadClientInfo.h"
+#include "mozilla/RefPtr.h"
 
 namespace mozilla {
 namespace mscom {
@@ -17,6 +19,7 @@ class MOZ_NON_TEMPORARY_CLASS MainThreadRuntime
 {
 public:
   MainThreadRuntime();
+  ~MainThreadRuntime();
 
   explicit operator bool() const
   {
@@ -28,11 +31,21 @@ public:
   MainThreadRuntime& operator=(MainThreadRuntime&) = delete;
   MainThreadRuntime& operator=(MainThreadRuntime&&) = delete;
 
+  /**
+   * @return 0 if call is in-process or resolving the calling thread failed,
+   *         otherwise contains the thread id of the calling thread.
+   */
+  static DWORD GetClientThreadId();
+
 private:
   HRESULT InitializeSecurity();
 
   STARegion mStaRegion;
   HRESULT mInitResult;
+
+  RefPtr<MainThreadClientInfo>  mClientInfo;
+
+  static MainThreadRuntime* sInstance;
 };
 
 } // namespace mscom

ipc/mscom/moz.build

@@ -11,6 +11,7 @@ EXPORTS.mozilla.mscom += [
     'COMApartmentRegion.h',
     'COMPtrHolder.h',
     'EnsureMTA.h',
+    'MainThreadClientInfo.h',
     'MainThreadRuntime.h',
     'Objref.h',
     'ProxyStream.h',
@@ -21,6 +22,7 @@ EXPORTS.mozilla.mscom += [
 UNIFIED_SOURCES += [
     'AgileReference.cpp',
     'EnsureMTA.cpp',
+    'MainThreadClientInfo.cpp',
     'MainThreadRuntime.cpp',
     'Objref.cpp',
     'ProxyStream.cpp',