[release/5.0] WinRT Reference Tracker memory fixes

src/coreclr/src/interop/comwrappers.cpp

@@ -4,6 +4,7 @@
 #include "comwrappers.hpp"
 #include <interoplibabi.h>
 #include <interoplibimports.h>
+#include <corerror.h>
 
 #include <new> // placement new
 
@@ -188,7 +189,7 @@ HRESULT STDMETHODCALLTYPE ManagedObjectWrapper_QueryInterface(
     /* [iid_is][out] */ _COM_Outptr_ void __RPC_FAR* __RPC_FAR* ppvObject)
 {
     ManagedObjectWrapper* wrapper = ABI::ToManagedObjectWrapper(disp);
-    return wrapper->QueryInterface(riid, ppvObject);
+    return wrapper->QueryInterface(disp, riid, ppvObject);
 }
 
 namespace
@@ -238,6 +239,11 @@ namespace
         return static_cast<ULONG>(c & ComRefCountMask);
     }
 
+    constexpr bool IsMarkedToDestroy(_In_ ULONGLONG c)
+    {
+        return (c & DestroySentinel) != 0;
+    }
+
     ULONG STDMETHODCALLTYPE TrackerTarget_AddRefFromReferenceTracker(_In_ ABI::ComInterfaceDispatch* disp)
     {
         _ASSERTE(disp != nullptr && disp->vtable != nullptr);
@@ -458,9 +464,14 @@ ManagedObjectWrapper::ManagedObjectWrapper(
     , _dispatches{ dispatches }
     , _refCount{ 1 }
     , _flags{ flags }
+    , _refTrackerDispatch{ nullptr }
 {
     bool wasSet = TrySetObjectHandle(objectHandle);
     _ASSERTE(wasSet);
+
+    // Retain the dispatch entry for the IReferenceTrackerTarget so we
+    // can branch quickly during a QI from that interface.
+    _refTrackerDispatch = (const ABI::ComInterfaceDispatch*)AsRuntimeDefined(__uuidof(IReferenceTrackerTarget));
 }
 
 ManagedObjectWrapper::~ManagedObjectWrapper()
@@ -542,6 +553,11 @@ bool ManagedObjectWrapper::IsRooted() const
     return rooted;
 }
 
+bool ManagedObjectWrapper::IsMarkedToDestroy() const
+{
+    return ::IsMarkedToDestroy(_refCount);
+}
+
 ULONG ManagedObjectWrapper::AddRefFromReferenceTracker()
 {
     LONGLONG prev;
@@ -574,10 +590,7 @@ ULONG ManagedObjectWrapper::ReleaseFromReferenceTracker()
     // If we observe the destroy sentinel, then this release
     // must destroy the wrapper.
     if (refCount == DestroySentinel)
-    {
-        _ASSERTE(!IsSet(CreateComInterfaceFlagsEx::IsPegged));
         Destroy(this);
-    }
 
     return GetTrackerCount(refCount);
 }
@@ -595,12 +608,37 @@ HRESULT ManagedObjectWrapper::Unpeg()
 }
 
 HRESULT ManagedObjectWrapper::QueryInterface(
+    _In_ const ABI::ComInterfaceDispatch* dispatch,
     /* [in] */ REFIID riid,
     /* [iid_is][out] */ _COM_Outptr_ void __RPC_FAR* __RPC_FAR* ppvObject)
 {
     if (ppvObject == nullptr)
         return E_POINTER;
 
+    // Check if this is a QI from an IReferenceTrackerTarget.
+    ComHolder<ManagedObjectWrapper> releaseMaybe;
+    if (dispatch == _refTrackerDispatch)
+    {
+        // AddRef to keep the managed object alive.
+        // AddRef is "safe" at this point because if it is a MOW with outstanding
+        // Reference Tracker reference, we know for sure the MOW is not claimed yet
+        // but the managed object could be.
+        AddRef();
+
+        // We are taking an extra AddRef() that now must be released.
+        releaseMaybe.Attach(this);
+
+        // For MOWs that have outstanding Reference Tracker reference, they could be either:
+        //  1. Marked to Destroy - in this case it is unsafe to touch wrapper.
+        //  2. Object Handle target has been NULLed out by GC.
+        if (IsMarkedToDestroy() || !InteropLibImports::HasValidTarget(Target))
+        {
+            // It is unsafe to proceed with a QueryInterface call. The MOW has been
+            // marked destroyed or the associated managed object has been collected.
+            return COR_E_ACCESSING_CCW;
+        }
+    }
+
     // Find target interface
     *ppvObject = AsRuntimeDefined(riid);
     if (*ppvObject == nullptr)
@@ -653,13 +691,11 @@ HRESULT ManagedObjectWrapper::QueryInterface(
 
 ULONG ManagedObjectWrapper::AddRef(void)
 {
-    _ASSERTE((_refCount & DestroySentinel) == 0);
     return GetComCount(::InterlockedIncrement64(&_refCount));
 }
 
 ULONG ManagedObjectWrapper::Release(void)
 {
-    _ASSERTE((_refCount & DestroySentinel) == 0);
     if (GetComCount(_refCount) == 0)
     {
         _ASSERTE(!"Over release of MOW - COM");

src/coreclr/src/interop/comwrappers.hpp

@@ -47,6 +47,7 @@ class ManagedObjectWrapper
 
     LONGLONG _refCount;
     Volatile<CreateComInterfaceFlagsEx> _flags;
+    const ABI::ComInterfaceDispatch* _refTrackerDispatch;
 
 public: // static
     // Get the implementation for IUnknown.
@@ -101,6 +102,9 @@ class ManagedObjectWrapper
     // Indicate if the wrapper should be considered a GC root.
     bool IsRooted() const;
 
+    // Check if the wrapper has been marked to be destroyed.
+    bool IsMarkedToDestroy() const;
+
 public: // IReferenceTrackerTarget
     ULONG AddRefFromReferenceTracker();
     ULONG ReleaseFromReferenceTracker();
@@ -109,6 +113,7 @@ class ManagedObjectWrapper
 
 public: // Lifetime
     HRESULT QueryInterface(
+        _In_ const ABI::ComInterfaceDispatch* dispatch,
         /* [in] */ REFIID riid,
         /* [iid_is][out] */ _COM_Outptr_ void __RPC_FAR * __RPC_FAR * ppvObject);
     ULONG AddRef(void);

src/coreclr/src/interop/inc/interoplibimports.h

@@ -44,6 +44,9 @@ namespace InteropLibImports
     // Delete Object instance handle.
     void DeleteObjectInstanceHandle(_In_ InteropLib::OBJECTHANDLE handle) noexcept;
 
+    // Check if Object instance handle still points at an Object.
+    bool HasValidTarget(_In_ InteropLib::OBJECTHANDLE handle) noexcept;
+
     // Get the current global pegging state.
     bool GetGlobalPeggingState() noexcept;
 

src/coreclr/src/interop/interoplib.cpp

@@ -73,8 +73,6 @@ namespace InteropLib
             if (mow == nullptr)
                 return E_INVALIDARG;
 
-            (void)mow->AddRef();
-
             *object = mow->Target;
             return S_OK;
         }

src/coreclr/src/interop/trackerobjectmanager.cpp

@@ -176,8 +176,8 @@ namespace
 
             ManagedObjectWrapper* mow = ManagedObjectWrapper::MapFromIUnknown(target);
 
-            // Not a target we implemented.
-            if (mow == nullptr)
+            // Not a target we implemented or wrapper is marked to be destroyed.
+            if (mow == nullptr || mow->IsMarkedToDestroy())
                 return S_OK;
 
             // Notify the runtime a reference path was found.
@@ -331,17 +331,17 @@ HRESULT TrackerObjectManager::BeginReferenceTracking(_In_ RuntimeCallContext* cx
 
     s_HasTrackingStarted = TRUE;
 
-    // From this point, the tracker runtime decides whether a target
-    // should be pegged or not as the global pegging flag is now off.
-    InteropLibImports::SetGlobalPeggingState(false);
-
     // Let the tracker runtime know we are about to walk external objects so that
     // they can lock their reference cache. Note that the tracker runtime doesn't need to
     // unpeg all external objects at this point and they can do the pegging/unpegging.
     // in FindTrackerTargetsCompleted.
     _ASSERTE(s_TrackerManager != nullptr);
     RETURN_IF_FAILED(s_TrackerManager->ReferenceTrackingStarted());
 
+    // From this point, the tracker runtime decides whether a target
+    // should be pegged or not as the global pegging flag is now off.
+    InteropLibImports::SetGlobalPeggingState(false);
+
     // Time to walk the external objects
     RETURN_IF_FAILED(WalkExternalTrackerObjects(cxt));
 

src/coreclr/src/vm/interoplibinterface.cpp

@@ -699,6 +699,8 @@ namespace
         ::ZeroMemory(&gc, sizeof(gc));
         GCPROTECT_BEGIN(gc);
 
+        STRESS_LOG4(LF_INTEROP, LL_INFO1000, "Get or Create EOC: (Identity: 0x%p) (Flags: %x) (Maybe: 0x%p) (ID: %lld)\n", identity, flags, OBJECTREFToObject(wrapperMaybe), wrapperId);
+
         gc.implRef = impl;
         gc.wrapperMaybeRef = wrapperMaybe;
 
@@ -731,6 +733,8 @@ namespace
             }
         }
 
+        STRESS_LOG2(LF_INTEROP, LL_INFO1000, "EOC: 0x%p or Handle: 0x%p\n", extObjCxt, handle);
+
         if (extObjCxt != NULL)
         {
             gc.objRefMaybe = extObjCxt->GetObjectRef();
@@ -800,6 +804,8 @@ namespace
                     extObjCxt = cache->FindOrAdd(cacheKey, resultHolder.GetContext());
                 }
 
+                STRESS_LOG2(LF_INTEROP, LL_INFO100, "EOC cache insert: 0x%p == 0x%p\n", extObjCxt, resultHolder.GetContext());
+
                 // If the returned context matches the new context it means the
                 // new context was inserted or a unique instance was requested.
                 if (extObjCxt == resultHolder.GetContext())
@@ -843,6 +849,8 @@ namespace
             }
         }
 
+        STRESS_LOG3(LF_INTEROP, LL_INFO1000, "EOC: 0x%p, 0x%p => 0x%p\n", extObjCxt, identity, OBJECTREFToObject(gc.objRefMaybe));
+
         GCPROTECT_END();
 
         *objRef = gc.objRefMaybe;
@@ -1024,6 +1032,29 @@ namespace InteropLibImports
         DestroyHandleCommon(static_cast<::OBJECTHANDLE>(handle), InstanceHandleType);
     }
 
+    bool HasValidTarget(_In_ InteropLib::OBJECTHANDLE handle) noexcept
+    {
+        CONTRACTL
+        {
+            NOTHROW;
+            GC_NOTRIGGER;
+            MODE_ANY;
+            PRECONDITION(handle != NULL);
+        }
+        CONTRACTL_END;
+
+        bool isValid = false;
+        ::OBJECTHANDLE objectHandle = static_cast<::OBJECTHANDLE>(handle);
+
+        {
+            // Switch to cooperative mode so the handle can be safely inspected.
+            GCX_COOP_THREAD_EXISTS(GET_THREAD());
+            isValid = ObjectFromHandle(objectHandle) != NULL;
+        }
+
+        return isValid;
+    }
+
     bool GetGlobalPeggingState() noexcept
     {
         CONTRACTL

src/tests/Interop/COM/ComWrappers/MockReferenceTrackerRuntime/ReferenceTrackerRuntime.cpp

@@ -203,9 +203,17 @@ namespace
             {
                 assert(c != nullptr && id != nullptr);
 
+                ComSmartPtr<API::IReferenceTrackerTarget> mowMaybe;
+                if (S_OK == c->QueryInterface(&mowMaybe))
+                {
+                    (void)mowMaybe->AddRefFromReferenceTracker();
+                    c = mowMaybe.p;
+                }
+
                 try
                 {
                     *id = _elementId;
+
                     if (!_elements.insert(std::make_pair(*id, ComSmartPtr<IUnknown>{ c })).second)
                         return S_FALSE;
 
@@ -216,10 +224,6 @@ namespace
                     return E_OUTOFMEMORY;
                 }
 
-                ComSmartPtr<API::IReferenceTrackerTarget> mowMaybe;
-                if (S_OK == c->QueryInterface(&mowMaybe))
-                    (void)mowMaybe->AddRefFromReferenceTracker();
-
                 return S_OK;
             }