ComWrappers fixes from NET 5 memory leak (#51661)

* Port over .NET 5 fix.

* Make .NET 5 fix more efficient on .NET 6.
Add tests for the controlled QueryInterface failure.

* Improve logging for external object wrapper creation.
Validate the returned value for ReleaseFromReferenceTracker API.

7 files changed, 94 insertions, 16 deletions

diff --git a/src/coreclr/debug/daccess/request.cpp b/src/coreclr/debug/daccess/request.cpp
index 0f51fb5b083..664f540bf62 100644
--- a/src/coreclr/debug/daccess/request.cpp
+++ b/src/coreclr/debug/daccess/request.cpp
@@ -4143,7 +4143,8 @@ BOOL ClrDataAccess::DACIsComWrappersCCW(CLRDATA_ADDRESS ccwPtr)
         return FALSE;
     }
 
-    return qiAddress == GetEEFuncEntryPoint(ManagedObjectWrapper_QueryInterface);
+    return (qiAddress == GetEEFuncEntryPoint(ManagedObjectWrapper_QueryInterface)
+        || qiAddress == GetEEFuncEntryPoint(TrackerTarget_QueryInterface));
 }
 
 TADDR ClrDataAccess::DACGetManagedObjectWrapperFromCCW(CLRDATA_ADDRESS ccwPtr)
diff --git a/src/coreclr/inc/daccess.h b/src/coreclr/inc/daccess.h
index b0269592c9e..399c5c57a74 100644
--- a/src/coreclr/inc/daccess.h
+++ b/src/coreclr/inc/daccess.h
@@ -630,6 +630,7 @@ typedef struct _DacGlobals
 #endif
 #ifdef FEATURE_COMWRAPPERS
     ULONG fn__ManagedObjectWrapper_QueryInterface;
+    ULONG fn__TrackerTarget_QueryInterface;
 #endif
 
     // Vtable pointer values for all classes that must
diff --git a/src/coreclr/interop/comwrappers.cpp b/src/coreclr/interop/comwrappers.cpp
index a386723cca6..8485125c8b7 100644
--- a/src/coreclr/interop/comwrappers.cpp
+++ b/src/coreclr/interop/comwrappers.cpp
@@ -3,6 +3,7 @@
 
 #include "comwrappers.hpp"
 #include <interoplibimports.h>
+#include <corerror.h>
 
 #include <new> // placement new
 
@@ -179,8 +180,7 @@ namespace ABI
 }
 
 // ManagedObjectWrapper_QueryInterface needs to be visible outside of this compilation unit
-// to support the DAC. See code:ClrDataAccess::DACTryGetComWrappersObjectFromCCW for the
-// usage in the DAC (look for the GetEEFuncEntryPoint call).
+// to support the DAC (look for the GetEEFuncEntryPoint call).
 HRESULT STDMETHODCALLTYPE ManagedObjectWrapper_QueryInterface(
     _In_ ABI::ComInterfaceDispatch* disp,
     /* [in] */ REFIID riid,
@@ -219,6 +219,36 @@ namespace
     static_assert(sizeof(ManagedObjectWrapper_IUnknownImpl) == (3 * sizeof(void*)), "Unexpected vtable size");
 }
 
+// TrackerTarget_QueryInterface needs to be visible outside of this compilation unit
+// to support the DAC (look for the GetEEFuncEntryPoint call).
+HRESULT STDMETHODCALLTYPE TrackerTarget_QueryInterface(
+    _In_ ABI::ComInterfaceDispatch* disp,
+    /* [in] */ REFIID riid,
+    /* [iid_is][out] */ _COM_Outptr_ void __RPC_FAR* __RPC_FAR* ppvObject)
+{
+    ManagedObjectWrapper* wrapper = ABI::ToManagedObjectWrapper(disp);
+
+    // AddRef is "safe" at this point because since it is a MOW with an outstanding
+    // Reference Tracker reference, we know for sure the MOW is not claimed yet
+    // but the managed object could be. If the managed object is alive at this
+    // moment the AddRef will ensure it remains alive for the duration of the
+    // QueryInterface.
+    ComHolder<ManagedObjectWrapper> ensureStableLifetime{ wrapper };
+
+    // 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 (wrapper->IsMarkedToDestroy()
+        || !InteropLibImports::HasValidTarget(wrapper->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;
+    }
+
+    return wrapper->QueryInterface(riid, ppvObject);
+}
+
 namespace
 {
     const int32_t TrackerRefShift = 32;
@@ -237,6 +267,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);
@@ -280,13 +315,17 @@ namespace
     // Hard-coded IReferenceTrackerTarget vtable
     const struct
     {
-        decltype(ManagedObjectWrapper_IUnknownImpl) IUnknownImpl;
+        decltype(&TrackerTarget_QueryInterface) QueryInterface;
+        decltype(&ManagedObjectWrapper_AddRef) AddRef;
+        decltype(&ManagedObjectWrapper_Release) Release;
         decltype(&TrackerTarget_AddRefFromReferenceTracker) AddRefFromReferenceTracker;
         decltype(&TrackerTarget_ReleaseFromReferenceTracker) ReleaseFromReferenceTracker;
         decltype(&TrackerTarget_Peg) Peg;
         decltype(&TrackerTarget_Unpeg) Unpeg;
     } ManagedObjectWrapper_IReferenceTrackerTargetImpl {
-        ManagedObjectWrapper_IUnknownImpl,
+        &TrackerTarget_QueryInterface,
+        &ManagedObjectWrapper_AddRef,
+        &ManagedObjectWrapper_Release,
         &TrackerTarget_AddRefFromReferenceTracker,
         &TrackerTarget_ReleaseFromReferenceTracker,
         &TrackerTarget_Peg,
@@ -541,6 +580,11 @@ bool ManagedObjectWrapper::IsRooted() const
     return rooted;
 }
 
+bool ManagedObjectWrapper::IsMarkedToDestroy() const
+{
+    return ::IsMarkedToDestroy(_refCount);
+}
+
 ULONG ManagedObjectWrapper::AddRefFromReferenceTracker()
 {
     LONGLONG prev;
@@ -573,10 +617,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);
 }
@@ -652,13 +693,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");
diff --git a/src/coreclr/interop/comwrappers.hpp b/src/coreclr/interop/comwrappers.hpp
index dc743651e7f..31b26e0ea34 100644
--- a/src/coreclr/interop/comwrappers.hpp
+++ b/src/coreclr/interop/comwrappers.hpp
@@ -103,6 +103,9 @@ public:
     // 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();
diff --git a/src/coreclr/interop/inc/interoplibimports.h b/src/coreclr/interop/inc/interoplibimports.h
index deb3f196f96..57824c36d78 100644
--- a/src/coreclr/interop/inc/interoplibimports.h
+++ b/src/coreclr/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;
 
diff --git a/src/coreclr/interop/trackerobjectmanager.cpp b/src/coreclr/interop/trackerobjectmanager.cpp
index f205484d3b0..ff5fcdb05c9 100644
--- a/src/coreclr/interop/trackerobjectmanager.cpp
+++ b/src/coreclr/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,10 +331,6 @@ 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.
@@ -342,6 +338,10 @@ HRESULT TrackerObjectManager::BeginReferenceTracking(_In_ RuntimeCallContext* cx
     _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));
 
diff --git a/src/coreclr/vm/interoplibinterface.cpp b/src/coreclr/vm/interoplibinterface.cpp
index d11b35ebf0d..cc45bd7a520 100644
--- a/src/coreclr/vm/interoplibinterface.cpp
+++ b/src/coreclr/vm/interoplibinterface.cpp
@@ -693,6 +693,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;
 
@@ -725,6 +727,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();
@@ -794,6 +798,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())
@@ -837,6 +843,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;
@@ -1018,6 +1026,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