Fix resolving IWeakReference from different context (microsoft#1301)

* Fix issue with weak reference when marshaled across contexts where we can end up treating IUnknown as the IWeakReference interface * Slight refactor
dongle-the-gadget · Jul 29, 2023 · 5c60130 · 5c60130
1 parent 7cc7f24
commit 5c60130
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 4 deletions.
diff --git a/src/Tests/UnitTest/TestComponentCSharp_Tests.cs b/src/Tests/UnitTest/TestComponentCSharp_Tests.cs
@@ -9,7 +9,6 @@
 
 using Windows.Foundation;
 using Windows.UI;
-using Windows.Security.Credentials.UI;
 using Windows.Storage;
 using Windows.Storage.Streams;
 using Microsoft.UI.Xaml;
@@ -30,6 +29,7 @@
 using System.Reflection;
 using Windows.Devices.Enumeration.Pnp;
 using System.Diagnostics;
+using Windows.Devices.Enumeration;
 
 #if NET
 using WeakRefNS = System;
@@ -2936,5 +2936,47 @@ private void TestExperimentAttribute()
             CustomExperimentClass custom = new CustomExperimentClass();
             custom.f();
         }
+
+        void OnDeviceAdded(DeviceWatcher sender, DeviceInformation args)
+        {
+        }
+
+        void OnDeviceUpdated(DeviceWatcher sender, DeviceInformationUpdate args)
+        {
+        }
+
+        [Fact]
+        public void TestWeakReferenceEventsFromMultipleContexts()
+        {
+            SemaphoreSlim semaphore = new SemaphoreSlim(0);
+            DeviceWatcher watcher = null;
+
+            Thread staThread = new Thread(() =>
+            {
+                Assert.True(Thread.CurrentThread.GetApartmentState() == ApartmentState.STA);
+
+                watcher = DeviceInformation.CreateWatcher();
+                var exception = Record.Exception(() => { 
+                    watcher.Added += OnDeviceAdded; 
+                });
+                Assert.Null(exception);
+
+                Thread mtaThread = new Thread(() =>
+                {
+                    Assert.True(Thread.CurrentThread.GetApartmentState() == ApartmentState.MTA);
+
+                    exception = Record.Exception(() => { 
+                        watcher.Updated += OnDeviceUpdated; 
+                    });
+                    Assert.Null(exception);
+                });
+                mtaThread.SetApartmentState(ApartmentState.MTA);
+                mtaThread.Start();
+                mtaThread.Join();
+            });
+            staThread.SetApartmentState(ApartmentState.STA);
+            staThread.Start();
+            staThread.Join();
+        }
     }
 }
diff --git a/src/WinRT.Runtime/ComWrappersSupport.cs b/src/WinRT.Runtime/ComWrappersSupport.cs
@@ -125,14 +125,28 @@ public static ObjectReference<T> GetObjectReferenceForInterface<T>(IntPtr extern
         }
 
         public static ObjectReference<T> GetObjectReferenceForInterface<T>(IntPtr externalComObject, Guid iid)
+        {
+            return GetObjectReferenceForInterface<T>(externalComObject, iid, true);
+        }
+
+        internal static ObjectReference<T> GetObjectReferenceForInterface<T>(IntPtr externalComObject, Guid iid, bool requireQI)
         {
             if (externalComObject == IntPtr.Zero)
             {
                 return null;
             }
 
-            Marshal.ThrowExceptionForHR(Marshal.QueryInterface(externalComObject, ref iid, out IntPtr ptr));
-            ObjectReference<T> objRef = ObjectReference<T>.Attach(ref ptr);
+            ObjectReference<T> objRef;
+            if (requireQI)
+            {
+                Marshal.ThrowExceptionForHR(Marshal.QueryInterface(externalComObject, ref iid, out IntPtr ptr));
+                objRef = ObjectReference<T>.Attach(ref ptr);
+            }
+            else
+            {
+                objRef = ObjectReference<T>.FromAbi(externalComObject);
+            }
+
             if (IsFreeThreaded(objRef))
             {
                 return objRef;

diff --git a/src/WinRT.Runtime/ComWrappersSupport.net5.cs b/src/WinRT.Runtime/ComWrappersSupport.net5.cs
@@ -539,7 +539,7 @@ private static object CreateObject(IntPtr externalComObject)
                 {
                     // IWeakReference is IUnknown-based, so implementations of it may not (and likely won't) implement
                     // IInspectable. As a result, we need to check for them explicitly.
-                    var iunknownObjRef = ComWrappersSupport.GetObjectReferenceForInterface<IUnknownVftbl>(ptr);
+                    var iunknownObjRef = ComWrappersSupport.GetObjectReferenceForInterface<IUnknownVftbl>(ptr, weakReferenceIID, false);
                     ComWrappersHelper.Init(iunknownObjRef);
 
                     return new SingleInterfaceOptimizedObject(typeof(IWeakReference), iunknownObjRef, false);