From ef21429ee711835526d42c1827d44aa10e439eb3 Mon Sep 17 00:00:00 2001
From: Johan Laanstra <jlaans@microsoft.com>
Date: Thu, 1 Feb 2024 17:24:21 -0800
Subject: [PATCH] Ensure the FactoryObjectReference also holds a pointer to the
 context callback. (#1492)

---
 src/WinRT.Runtime/Context.cs | 15 ---------------
 src/cswinrt/strings/WinRT.cs | 29 ++++++++++++++++++++++++++++-
 2 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/src/WinRT.Runtime/Context.cs b/src/WinRT.Runtime/Context.cs
index ba138e313..355b908d6 100644
--- a/src/WinRT.Runtime/Context.cs
+++ b/src/WinRT.Runtime/Context.cs
@@ -9,18 +9,8 @@ namespace WinRT
 {
     static partial class Context
     {
-        [DllImport("api-ms-win-core-com-l1-1-0.dll")]
-        private static extern int CoGetObjectContext(ref Guid riid, out IntPtr ppv);
-
         private static readonly Guid IID_ICallbackWithNoReentrancyToApplicationSTA = new(0x0A299774, 0x3E4E, 0xFC42, 0x1D, 0x9D, 0x72, 0xCE, 0xE1, 0x05, 0xCA, 0x57);
 
-        public static IntPtr GetContextCallback()
-        {
-            Guid riid = ABI.WinRT.Interop.IContextCallback.IID;
-            Marshal.ThrowExceptionForHR(CoGetObjectContext(ref riid, out IntPtr contextCallbackPtr));
-            return contextCallbackPtr;
-        }
-
         // Calls the given callback in the right context.
         // On any exception, calls onFail callback if any set.
         // If not set, exception is handled due to today we don't
@@ -50,10 +40,5 @@ public unsafe static void CallInContext(IntPtr contextCallbackPtr, IntPtr contex
                 onFailCallback?.Invoke();
             }
         }
-
-        public static void DisposeContextCallback(IntPtr contextCallbackPtr)
-        {
-            MarshalInspectable<object>.DisposeAbi(contextCallbackPtr);
-        }
     }
 }
diff --git a/src/cswinrt/strings/WinRT.cs b/src/cswinrt/strings/WinRT.cs
index 216e1cc42..f94bb3655 100644
--- a/src/cswinrt/strings/WinRT.cs
+++ b/src/cswinrt/strings/WinRT.cs
@@ -267,12 +267,30 @@ internal static partial class Context
         [DllImport("api-ms-win-core-com-l1-1-0.dll")]
         private static extern unsafe int CoGetContextToken(IntPtr* contextToken);
 
+        [DllImport("api-ms-win-core-com-l1-1-0.dll")]
+        private static extern int CoGetObjectContext(ref Guid riid, out IntPtr ppv);
+
         public unsafe static IntPtr GetContextToken()
         {
             IntPtr contextToken;
             Marshal.ThrowExceptionForHR(CoGetContextToken(&contextToken));
             return contextToken;
         }
+
+        public static IntPtr GetContextCallback()
+        {
+            Guid riid = InterfaceIIDs.IContextCallback_IID;
+            Marshal.ThrowExceptionForHR(CoGetObjectContext(ref riid, out IntPtr contextCallbackPtr));
+            return contextCallbackPtr;
+        }
+
+        public static void DisposeContextCallback(IntPtr contextCallbackPtr)
+        {
+            if (contextCallbackPtr != IntPtr.Zero)
+            {
+                MarshalInspectable<object>.DisposeAbi(contextCallbackPtr);
+            }
+        }
     }
 
     internal unsafe sealed class DllModule
@@ -459,7 +477,8 @@ internal sealed class FactoryObjectReference<
 #endif
         T> : IObjectReference
     {
-        private readonly IntPtr _contextToken;
+        private readonly IntPtr _contextToken;
+        private readonly IntPtr _contextCallback;
 
         public static FactoryObjectReference<T> Attach(ref IntPtr thisPtr)
         {
@@ -478,6 +497,7 @@ internal FactoryObjectReference(IntPtr thisPtr) :
             if (!IsFreeThreaded(this))
             {
                 _contextToken = Context.GetContextToken();
+                _contextCallback = Context.GetContextCallback();
             }
         }
 
@@ -498,6 +518,13 @@ public static unsafe FactoryObjectReference<T> FromAbi(IntPtr thisPtr)
             return obj;
         }
 
+        protected override unsafe void Release()
+        {
+            base.Release();
+
+            Context.DisposeContextCallback(_contextCallback);
+        }
+
         public bool IsObjectInContext()
         {
             return _contextToken == IntPtr.Zero || _contextToken == Context.GetContextToken();