Keep delegate fields alive across a full native call (#20896)

* Repurpose CleanupWorkList to also preserve delegate references in structs across the full native call.

* Change CleanupWorkListElement to abstract base class instead of interface.

* Make CleanupWorkList a singlely linked list.

* PR Feedback.

* Remove CleanupWorkList and make CleanupWorkListElement be able to represent the full list.

* Add back throw in SafeHandle field marshalling.

* PR feedback.

6 files changed, 109 insertions, 69 deletions

diff --git a/src/System.Private.CoreLib/src/System/StubHelpers.cs b/src/System.Private.CoreLib/src/System/StubHelpers.cs
index ef2a952b32..9fa2952041 100644
--- a/src/System.Private.CoreLib/src/System/StubHelpers.cs
+++ b/src/System.Private.CoreLib/src/System/StubHelpers.cs
@@ -584,7 +584,7 @@ namespace System.StubHelpers
     internal static class ValueClassMarshaler
     {
         [MethodImplAttribute(MethodImplOptions.InternalCall)]
-        internal static extern void ConvertToNative(IntPtr dst, IntPtr src, IntPtr pMT, ref CleanupWorkList pCleanupWorkList);
+        internal static extern void ConvertToNative(IntPtr dst, IntPtr src, IntPtr pMT, ref CleanupWorkListElement pCleanupWorkList);
 
         [MethodImplAttribute(MethodImplOptions.InternalCall)]
         internal static extern void ConvertToManaged(IntPtr dst, IntPtr src, IntPtr pMT);
@@ -900,7 +900,7 @@ namespace System.StubHelpers
         private Type layoutType;
 
         // Cleanup list to be destroyed when clearing the native view (for layouts with SafeHandles).
-        private CleanupWorkList cleanupWorkList;
+        private CleanupWorkListElement cleanupWorkList;
 
         [Flags]
         internal enum AsAnyFlags
@@ -1453,42 +1453,82 @@ namespace System.StubHelpers
 #endif
     }  // struct NativeVariant
 
+    internal abstract class CleanupWorkListElement
+    {
+        private CleanupWorkListElement m_Next;
+        protected abstract void DestroyCore();
+
+        public void Destroy()
+        {
+            DestroyCore();
+            CleanupWorkListElement next = m_Next;
+            while (next != null)
+            {
+                next.DestroyCore();
+                next = next.m_Next;
+            }
+        }
+        
+        public static void AddToCleanupList(ref CleanupWorkListElement list, CleanupWorkListElement newElement)
+        {
+            if (list == null)
+            {
+                list = newElement;
+            }
+            else
+            {
+                newElement.m_Next = list;
+                list = newElement;
+            }
+        }
+    }
+
+    // Keeps a Delegate instance alive across the full Managed->Native call.
+    // This ensures that users don't have to call GC.KeepAlive after passing a struct or class
+    // that has a delegate field to native code.
+    internal sealed class DelegateCleanupWorkListElement : CleanupWorkListElement
+    {
+        public DelegateCleanupWorkListElement(Delegate del)
+        {
+            m_del = del;
+        }
+
+        private Delegate m_del;
+
+        protected override void DestroyCore()
+        {
+            GC.KeepAlive(m_del);
+        }
+    }
+
     // Aggregates SafeHandle and the "owned" bit which indicates whether the SafeHandle
     // has been successfully AddRef'ed. This allows us to do realiable cleanup (Release)
     // if and only if it is needed.
-    internal sealed class CleanupWorkListElement
+    internal sealed class SafeHandleCleanupWorkListElement : CleanupWorkListElement
     {
-        public CleanupWorkListElement(SafeHandle handle)
+        public SafeHandleCleanupWorkListElement(SafeHandle handle)
         {
             m_handle = handle;
         }
 
-        public SafeHandle m_handle;
+        private SafeHandle m_handle;
 
         // This field is passed by-ref to SafeHandle.DangerousAddRef.
-        // CleanupWorkList.Destroy ignores this element if m_owned is not set to true.
-        public bool m_owned;
-    }  // class CleanupWorkListElement
-
-    internal sealed class CleanupWorkList
-    {
-        private List<CleanupWorkListElement> m_list = new List<CleanupWorkListElement>();
+        // DestroyCore ignores this element if m_owned is not set to true.
+        private bool m_owned;
 
-        public void Add(CleanupWorkListElement elem)
+        protected override void DestroyCore()
         {
-            Debug.Assert(elem.m_owned == false, "m_owned is supposed to be false and set later by DangerousAddRef");
-            m_list.Add(elem);
+            if (m_owned)
+                StubHelpers.SafeHandleRelease(m_handle);
         }
 
-        public void Destroy()
+        public IntPtr AddRef()
         {
-            for (int i = m_list.Count - 1; i >= 0; i--)
-            {
-                if (m_list[i].m_owned)
-                    StubHelpers.SafeHandleRelease(m_list[i].m_handle);
-            }
+            // element.m_owned will be true iff the AddRef succeeded
+            return StubHelpers.SafeHandleAddRef(m_handle, ref m_owned);
         }
-    }  // class CleanupWorkList
+    }  // class CleanupWorkListElement
 
     internal static class StubHelpers
     {
@@ -1513,19 +1553,20 @@ namespace System.StubHelpers
         [MethodImplAttribute(MethodImplOptions.InternalCall)]
         internal static extern void ThrowInteropParamException(int resID, int paramIdx);
 
-        internal static IntPtr AddToCleanupList(ref CleanupWorkList pCleanupWorkList, SafeHandle handle)
+        internal static IntPtr AddToCleanupList(ref CleanupWorkListElement pCleanupWorkList, SafeHandle handle)
         {
-            if (pCleanupWorkList == null)
-                pCleanupWorkList = new CleanupWorkList();
-
-            CleanupWorkListElement element = new CleanupWorkListElement(handle);
-            pCleanupWorkList.Add(element);
+            SafeHandleCleanupWorkListElement element = new SafeHandleCleanupWorkListElement(handle);
+            CleanupWorkListElement.AddToCleanupList(ref pCleanupWorkList, element);
+            return element.AddRef();
+        }
 
-            // element.m_owned will be true iff the AddRef succeeded
-            return SafeHandleAddRef(handle, ref element.m_owned);
+        internal static void AddToCleanupList(ref CleanupWorkListElement pCleanupWorkList, Delegate del)
+        {
+            DelegateCleanupWorkListElement element = new DelegateCleanupWorkListElement(del);
+            CleanupWorkListElement.AddToCleanupList(ref pCleanupWorkList, element);
         }
 
-        internal static void DestroyCleanupList(ref CleanupWorkList pCleanupWorkList)
+        internal static void DestroyCleanupList(ref CleanupWorkListElement pCleanupWorkList)
         {
             if (pCleanupWorkList != null)
             {
@@ -1690,7 +1731,7 @@ namespace System.StubHelpers
         internal static extern unsafe int strlen(sbyte* ptr);
 
         [MethodImplAttribute(MethodImplOptions.InternalCall)]
-        internal static extern unsafe void FmtClassUpdateNativeInternal(object obj, byte* pNative, ref CleanupWorkList pCleanupWorkList);
+        internal static extern unsafe void FmtClassUpdateNativeInternal(object obj, byte* pNative, ref CleanupWorkListElement pCleanupWorkList);
         [MethodImplAttribute(MethodImplOptions.InternalCall)]
         internal static extern unsafe void FmtClassUpdateCLRInternal(object obj, byte* pNative);
         [MethodImplAttribute(MethodImplOptions.InternalCall)]
diff --git a/src/vm/dllimport.cpp b/src/vm/dllimport.cpp
index 6fd4366624..acb356a42b 100644
--- a/src/vm/dllimport.cpp
+++ b/src/vm/dllimport.cpp
@@ -2167,7 +2167,7 @@ void NDirectStubLinker::NeedsCleanupList()
         SetCleanupNeeded();
 
         // we setup a new local that will hold the cleanup work list
-        LocalDesc desc(MscorlibBinder::GetClass(CLASS__CLEANUP_WORK_LIST));
+        LocalDesc desc(MscorlibBinder::GetClass(CLASS__CLEANUP_WORK_LIST_ELEMENT));
         m_dwCleanupWorkListLocalNum = NewLocal(desc);
     }
 }
diff --git a/src/vm/fieldmarshaler.cpp b/src/vm/fieldmarshaler.cpp
index d5c16b0c3b..57f41c33ce 100644
--- a/src/vm/fieldmarshaler.cpp
+++ b/src/vm/fieldmarshaler.cpp
@@ -3911,6 +3911,24 @@ VOID FieldMarshaler_Delegate::UpdateNativeImpl(OBJECTREF* pCLRValue, LPVOID pNat
     CONTRACTL_END;
 
     LPVOID fnPtr = COMDelegate::ConvertToCallback(*pCLRValue);
+    
+    // If there is no CleanupWorkList (i.e. a call from Marshal.StructureToPtr), we don't use it to manage delegate lifetime.
+    // In that case, it falls on the user to manage the delegate lifetime. This is the cleanest way to do this since there is no well-defined
+    // object lifetime for the unmanaged memory that the structure would be marshalled to in the Marshal.StructureToPtr case.
+    if (*pCLRValue != NULL && ppCleanupWorkListOnStack != NULL)
+    {
+        // Call StubHelpers.AddToCleanupList to ensure the delegate is kept alive across the full native call.
+        MethodDescCallSite AddToCleanupList(METHOD__STUBHELPERS__ADD_TO_CLEANUP_LIST_DELEGATE);
+
+        ARG_SLOT args[] =
+        {
+            (ARG_SLOT)ppCleanupWorkListOnStack,
+            ObjToArgSlot(*pCLRValue)
+        };
+
+        AddToCleanupList.Call(args);
+    }
+
     MAYBE_UNALIGNED_WRITE(pNativeValue, _PTR, fnPtr);
 }
 
@@ -3956,14 +3974,14 @@ VOID FieldMarshaler_SafeHandle::UpdateNativeImpl(OBJECTREF* pCLRValue, LPVOID pN
     // A cleanup list MUST be specified in order for us to be able to marshal
     // the SafeHandle.
     if (ppCleanupWorkListOnStack == NULL)
-        COMPlusThrow(kInvalidOperationException, IDS_EE_SH_FIELD_INVALID_OPERATION);
+        COMPlusThrow(kInvalidOperationException, IDS_EE_SH_FIELD_INVALID_OPERATION); 
 
     if (*pSafeHandleObj == NULL)
         COMPlusThrow(kArgumentNullException, W("ArgumentNull_SafeHandle"));
 
     // Call StubHelpers.AddToCleanupList to AddRef and schedule Release on this SafeHandle
     // This is realiable, i.e. the cleanup will happen if and only if the SH was actually AddRef'ed.
-    MethodDescCallSite AddToCleanupList(METHOD__STUBHELPERS__ADD_TO_CLEANUP_LIST);
+    MethodDescCallSite AddToCleanupList(METHOD__STUBHELPERS__ADD_TO_CLEANUP_LIST_SAFEHANDLE);
 
     ARG_SLOT args[] =
     {
diff --git a/src/vm/ilmarshalers.cpp b/src/vm/ilmarshalers.cpp
index feb2030a3a..2dbf5d002c 100644
--- a/src/vm/ilmarshalers.cpp
+++ b/src/vm/ilmarshalers.cpp
@@ -1125,17 +1125,8 @@ void ILValueClassMarshaler::EmitConvertContentsCLRToNative(ILCodeStream* pslILEm
     pslILEmit->EmitLDTOKEN(managedVCToken); // pMT
     pslILEmit->EmitCALL(METHOD__RT_TYPE_HANDLE__GETVALUEINTERNAL, 1, 1); // Convert RTH to IntPtr
 
-    if (IsCLRToNative(m_dwMarshalFlags))
-    {
-        // this should only be needed in CLR-to-native scenarios for the SafeHandle field marshaler
-        m_pslNDirect->LoadCleanupWorkList(pslILEmit);
-    }
-    else
-    {
-        pslILEmit->EmitLoadNullPtr();
-    }
-
-    pslILEmit->EmitCALL(METHOD__VALUECLASSMARSHALER__CONVERT_TO_NATIVE, 4, 0);        // void ConvertToNative(IntPtr dst, IntPtr src, IntPtr pMT, ref CleanupWorkList pCleanupWorkList)
+    m_pslNDirect->LoadCleanupWorkList(pslILEmit);
+    pslILEmit->EmitCALL(METHOD__VALUECLASSMARSHALER__CONVERT_TO_NATIVE, 4, 0);        // void ConvertToNative(IntPtr dst, IntPtr src, IntPtr pMT, ref CleanupWorkListElement pCleanupWorkList)
 }
 
 void ILValueClassMarshaler::EmitConvertContentsNativeToCLR(ILCodeStream* pslILEmit)
@@ -2410,19 +2401,7 @@ void ILLayoutClassPtrMarshaler::EmitConvertContentsCLRToNative(ILCodeStream* psl
     EmitLoadManagedValue(pslILEmit);
     EmitLoadNativeValue(pslILEmit);
 
-    if (IsCLRToNative(m_dwMarshalFlags))
-    {
-        m_pslNDirect->LoadCleanupWorkList(pslILEmit);
-    }
-    else
-    {
-        //
-        // The assertion here is as follows:
-        //      1) the only field marshaler that requires the CleanupWorkList is FieldMarshaler_SafeHandle 
-        //      2) SafeHandle marshaling is disallowed in the native-to-CLR direction, so we'll never see it..
-        //
-        pslILEmit->EmitLDNULL();  // pass a NULL CleanupWorkList in the native-to-CLR case
-    }
+    m_pslNDirect->LoadCleanupWorkList(pslILEmit);
 
     // static void FmtClassUpdateNativeInternal(object obj, byte* pNative, IntPtr pOptionalCleanupList);
 
@@ -2778,7 +2757,7 @@ MarshalerOverrideStatus ILSafeHandleMarshaler::ArgumentOverride(NDirectStubLinke
                 pslIL->EmitLDLOC(dwInputHandleLocal);
 
                 // This is realiable, i.e. the cleanup will happen if and only if the SH was actually AddRef'ed.
-                pslIL->EmitCALL(METHOD__STUBHELPERS__ADD_TO_CLEANUP_LIST, 2, 1);
+                pslIL->EmitCALL(METHOD__STUBHELPERS__ADD_TO_CLEANUP_LIST_SAFEHANDLE, 2, 1);
 
                 pslIL->EmitSTLOC(dwNativeHandleLocal);
 
diff --git a/src/vm/metasig.h b/src/vm/metasig.h
index e826526976..b258f441ce 100644
--- a/src/vm/metasig.h
+++ b/src/vm/metasig.h
@@ -212,7 +212,7 @@ DEFINE_METASIG(SM(RetBool, _, F))
 DEFINE_METASIG(SM(IntPtr_RetStr, I, s))
 DEFINE_METASIG(SM(IntPtr_RetBool, I, F))
 DEFINE_METASIG(SM(IntPtrIntPtrIntPtr_RetVoid, I I I, v))
-DEFINE_METASIG_T(SM(IntPtrIntPtrIntPtr_RefCleanupWorkList_RetVoid, I I I r(C(CLEANUP_WORK_LIST)), v))
+DEFINE_METASIG_T(SM(IntPtrIntPtrIntPtr_RefCleanupWorkListElement_RetVoid, I I I r(C(CLEANUP_WORK_LIST_ELEMENT)), v))
 DEFINE_METASIG_T(SM(RuntimeType_RuntimeMethodHandleInternal_RetMethodBase, C(CLASS) g(METHOD_HANDLE_INTERNAL), C(METHOD_BASE) ))
 DEFINE_METASIG_T(SM(RuntimeType_IRuntimeFieldInfo_RetFieldInfo, C(CLASS) C(I_RT_FIELD_INFO), C(FIELD_INFO) ))
 DEFINE_METASIG_T(SM(RuntimeType_Int_RetPropertyInfo, C(CLASS) i, C(PROPERTY_INFO) ))
@@ -220,7 +220,7 @@ DEFINE_METASIG(SM(Char_Bool_Bool_RetByte, u F F, b))
 DEFINE_METASIG(SM(Byte_RetChar, b, u))
 DEFINE_METASIG(SM(Str_Bool_Bool_RefInt_RetIntPtr, s F F r(i), I))
 DEFINE_METASIG(SM(IntPtr_Int_RetStr, I i, s))
-DEFINE_METASIG_T(SM(Obj_PtrByte_RefCleanupWorkList_RetVoid, j P(b) r(C(CLEANUP_WORK_LIST)), v))
+DEFINE_METASIG_T(SM(Obj_PtrByte_RefCleanupWorkListElement_RetVoid, j P(b) r(C(CLEANUP_WORK_LIST_ELEMENT)), v))
 DEFINE_METASIG(SM(Obj_PtrByte_RetVoid, j P(b), v))
 DEFINE_METASIG(SM(PtrByte_IntPtr_RetVoid, P(b) I, v))
 DEFINE_METASIG(SM(Str_Bool_Bool_RefInt_RetArrByte, s F F r(i), a(b) ))
@@ -581,8 +581,9 @@ DEFINE_METASIG_T(SM(RefDec_RetVoid, r(g(DECIMAL)), v))
 DEFINE_METASIG(GM(RefT_T_T_RetT, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, r(M(0)) M(0) M(0), M(0)))
 DEFINE_METASIG(SM(RefObject_Object_Object_RetObject, r(j) j j, j))
 
-DEFINE_METASIG_T(SM(RefCleanupWorkList_RetVoid, r(C(CLEANUP_WORK_LIST)), v))
-DEFINE_METASIG_T(SM(RefCleanupWorkList_SafeHandle_RetIntPtr, r(C(CLEANUP_WORK_LIST)) C(SAFE_HANDLE), I))
+DEFINE_METASIG_T(SM(RefCleanupWorkListElement_RetVoid, r(C(CLEANUP_WORK_LIST_ELEMENT)), v))
+DEFINE_METASIG_T(SM(RefCleanupWorkListElement_SafeHandle_RetIntPtr, r(C(CLEANUP_WORK_LIST_ELEMENT)) C(SAFE_HANDLE), I))
+DEFINE_METASIG_T(SM(RefCleanupWorkListElement_Delegate_RetVoid, r(C(CLEANUP_WORK_LIST_ELEMENT)) C(DELEGATE), v))
 
 #ifdef FEATURE_ICASTABLE
 DEFINE_METASIG_T(SM(ICastable_RtType_RefException_RetBool, C(ICASTABLE) C(CLASS) r(C(EXCEPTION)), F))
diff --git a/src/vm/mscorlib.h b/src/vm/mscorlib.h
index d9fb6dabb5..3cabfc8a52 100644
--- a/src/vm/mscorlib.h
+++ b/src/vm/mscorlib.h
@@ -95,7 +95,7 @@ DEFINE_METHOD(APP_DOMAIN,           ON_DESIGNER_NAMESPACE_RESOLVE, OnDesignerNam
 DEFINE_METHOD(APP_DOMAIN,           SETUP_DOMAIN,           SetupDomain,                IM_Bool_Str_Str_ArrStr_ArrStr_RetVoid)
 DEFINE_METHOD(APP_DOMAIN,           INITIALIZE_COMPATIBILITY_FLAGS, InitializeCompatibilityFlags,  IM_RetVoid)
 
-DEFINE_CLASS(CLEANUP_WORK_LIST,     StubHelpers,            CleanupWorkList)
+DEFINE_CLASS(CLEANUP_WORK_LIST_ELEMENT,     StubHelpers,            CleanupWorkListElement)
 
 #ifdef FEATURE_COMINTEROP
 // Define earlier in mscorlib.h to avoid BinderClassID to const BYTE truncation warning
@@ -1017,14 +1017,15 @@ DEFINE_METHOD(STUBHELPERS,          SET_LAST_ERROR,         SetLastError,
 DEFINE_METHOD(STUBHELPERS,          CLEAR_LAST_ERROR,       ClearLastError,             SM_RetVoid)
 
 DEFINE_METHOD(STUBHELPERS,          THROW_INTEROP_PARAM_EXCEPTION, ThrowInteropParamException,   SM_Int_Int_RetVoid)
-DEFINE_METHOD(STUBHELPERS,          ADD_TO_CLEANUP_LIST,    AddToCleanupList,           SM_RefCleanupWorkList_SafeHandle_RetIntPtr)
-DEFINE_METHOD(STUBHELPERS,          DESTROY_CLEANUP_LIST,   DestroyCleanupList,         SM_RefCleanupWorkList_RetVoid)
+DEFINE_METHOD(STUBHELPERS,          ADD_TO_CLEANUP_LIST_SAFEHANDLE,    AddToCleanupList,           SM_RefCleanupWorkListElement_SafeHandle_RetIntPtr)
+DEFINE_METHOD(STUBHELPERS,          ADD_TO_CLEANUP_LIST_DELEGATE,    AddToCleanupList,             SM_RefCleanupWorkListElement_Delegate_RetVoid)
+DEFINE_METHOD(STUBHELPERS,          DESTROY_CLEANUP_LIST,   DestroyCleanupList,         SM_RefCleanupWorkListElement_RetVoid)
 DEFINE_METHOD(STUBHELPERS,          GET_HR_EXCEPTION_OBJECT, GetHRExceptionObject,      SM_Int_RetException)
 DEFINE_METHOD(STUBHELPERS,          CREATE_CUSTOM_MARSHALER_HELPER, CreateCustomMarshalerHelper, SM_IntPtr_Int_IntPtr_RetIntPtr)
 
 DEFINE_METHOD(STUBHELPERS,          CHECK_STRING_LENGTH,    CheckStringLength,          SM_Int_RetVoid)
 
-DEFINE_METHOD(STUBHELPERS,          FMT_CLASS_UPDATE_NATIVE_INTERNAL,   FmtClassUpdateNativeInternal,   SM_Obj_PtrByte_RefCleanupWorkList_RetVoid)
+DEFINE_METHOD(STUBHELPERS,          FMT_CLASS_UPDATE_NATIVE_INTERNAL,   FmtClassUpdateNativeInternal,   SM_Obj_PtrByte_RefCleanupWorkListElement_RetVoid)
 DEFINE_METHOD(STUBHELPERS,          FMT_CLASS_UPDATE_CLR_INTERNAL,      FmtClassUpdateCLRInternal,      SM_Obj_PtrByte_RetVoid)
 DEFINE_METHOD(STUBHELPERS,          LAYOUT_DESTROY_NATIVE_INTERNAL,     LayoutDestroyNativeInternal,    SM_PtrByte_IntPtr_RetVoid)
 DEFINE_METHOD(STUBHELPERS,          ALLOCATE_INTERNAL,                  AllocateInternal,               SM_IntPtr_RetObj)
@@ -1161,7 +1162,7 @@ DEFINE_METHOD(HRESULTEXCEPTIONMARSHALER,  CONVERT_TO_MANAGED,    ConvertToManage
 #endif // FEATURE_COMINTEROP
 
 DEFINE_CLASS(VALUECLASSMARSHALER,   StubHelpers,            ValueClassMarshaler)
-DEFINE_METHOD(VALUECLASSMARSHALER,  CONVERT_TO_NATIVE,      ConvertToNative,            SM_IntPtrIntPtrIntPtr_RefCleanupWorkList_RetVoid)
+DEFINE_METHOD(VALUECLASSMARSHALER,  CONVERT_TO_NATIVE,      ConvertToNative,            SM_IntPtrIntPtrIntPtr_RefCleanupWorkListElement_RetVoid)
 DEFINE_METHOD(VALUECLASSMARSHALER,  CONVERT_TO_MANAGED,     ConvertToManaged,           SM_IntPtrIntPtrIntPtr_RetVoid)
 DEFINE_METHOD(VALUECLASSMARSHALER,  CLEAR_NATIVE,           ClearNative,                SM_IntPtr_IntPtr_RetVoid)