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.

4 files changed, 60 insertions, 15 deletions

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));