Fix issue with unsubscribing from events on non-agile objects (#2079)

* Fix issue where we maybe calling an add / remove handler function for a non agile object and thereby call the wrong vtable pointer when on a different context

* Fix offset calculation

* Simplify

* Add new constructor and refactor

* Use new constructor for events

* Add test

* Cleanup

Author: manodasanW

src/Tests/TestComponentCSharp/NonAgileClass.cpp

@@ -91,4 +91,23 @@ namespace winrt::TestComponentCSharp::implementation
         *put_abi(handler) = make<NonAgileDelegate>().detach();               
         vector.VectorChanged(handler);
     }
+
+    winrt::event_token NonAgileClass::CanExecuteChanged(winrt::Windows::Foundation::EventHandler<winrt::Windows::Foundation::IInspectable> const& handler)
+    {
+        return _event.add(handler);
+    }
+
+    void NonAgileClass::CanExecuteChanged(winrt::event_token const& token) noexcept
+    {
+        _event.remove(token);
+    }
+
+    bool NonAgileClass::CanExecute(winrt::Windows::Foundation::IInspectable const& parameter)
+    {
+        return true;
+    }
+
+    void NonAgileClass::Execute(winrt::Windows::Foundation::IInspectable const& parameter)
+    {
+    }
 }

src/Tests/TestComponentCSharp/NonAgileClass.h

@@ -5,10 +5,23 @@ namespace winrt::TestComponentCSharp::implementation
 {
     struct NonAgileClass : NonAgileClassT<NonAgileClass, winrt::non_agile>
     {
+        winrt::event<Windows::Foundation::EventHandler<winrt::Windows::Foundation::IInspectable>> _event;
     public:
         NonAgileClass();
         void Observe(Microsoft::UI::Xaml::Interop::IBindableObservableVector vector);
         void VectorChanged(Microsoft::UI::Xaml::Interop::IBindableObservableVector vector, Windows::Foundation::IInspectable e);
+
+        winrt::event_token CanExecuteChanged(winrt::Windows::Foundation::EventHandler<winrt::Windows::Foundation::IInspectable> const& handler);
+        void CanExecuteChanged(winrt::event_token const& token) noexcept;
+        bool CanExecute(winrt::Windows::Foundation::IInspectable const& parameter);
+        void Execute(winrt::Windows::Foundation::IInspectable const& parameter);
+
+        // Overriding runtime class name to allow marshaling as just the ICommand interface during proxy calls
+        // to avoid defining proxy stubs for this dll.
+        hstring GetRuntimeClassName() const
+        {
+            return L"Windows.UI.Xaml.Input.ICommand";
+        }
     };
 }
 namespace winrt::TestComponentCSharp::factory_implementation

src/Tests/TestComponentCSharp/TestComponentCSharp.idl

@@ -507,7 +507,7 @@ namespace TestComponentCSharp
     }
 
     [threading(sta), marshaling_behavior(standard)]
-    runtimeclass NonAgileClass
+    runtimeclass NonAgileClass : Microsoft.UI.Xaml.Input.ICommand
     {
         NonAgileClass();
         void Observe(Microsoft.UI.Xaml.Interop.IBindableObservableVector vector);

src/Tests/UnitTest/TestComponentCSharp_Tests.cs

@@ -3000,6 +3000,9 @@ class NonAgileClassCaller
             public void AcquireObject()
             {
                 Assert.Equal(ApartmentState.STA, Thread.CurrentThread.GetApartmentState());
+                nonAgileClass = new NonAgileClass();
+                nonAgileClass.CanExecuteChanged += NonAgileClass_Event;
+
                 nonAgileObject = new Windows.UI.Popups.PopupMenu();
                 nonAgileObject.Commands.Add(new Windows.UI.Popups.UICommand("test"));
                 nonAgileObject.Commands.Add(new Windows.UI.Popups.UICommand("test2"));
@@ -3011,6 +3014,7 @@ public void AcquireObject()
 
                 // Object gets proxied to the apartment.
                 Assert.Equal(2, proxyObject.Commands.Count);
+
                 agileReference.Dispose();
             }
 
@@ -3021,6 +3025,10 @@ public void CheckValue()
                 proxyObject = agileReference.Get();
                 Assert.Equal(2, proxyObject.Commands.Count);
 
+                // Remove the event handler from a different context from what it was initially added in
+                // to make sure we can unsubscribe from the event via the proxy in non agile scenarios.
+                nonAgileClass.CanExecuteChanged -= NonAgileClass_Event;
+
                 valueAcquired.Set();
             }
 
@@ -3031,11 +3039,16 @@ public void CallProxyObject()
                 Assert.ThrowsAny<System.Exception>(() => proxyObject.Commands);
             }
 
+            public void NonAgileClass_Event(object sender, object e)
+            {
+            }
+
             private Windows.UI.Popups.PopupMenu nonAgileObject;
             private Windows.UI.Popups.PopupMenu proxyObject;
-            private AgileReference<Windows.UI.Popups.PopupMenu> agileReference, agileReference2;
+            private AgileReference<Windows.UI.Popups.PopupMenu> agileReference;
             private readonly AutoResetEvent objectAcquired = new AutoResetEvent(false);
             private readonly AutoResetEvent valueAcquired = new AutoResetEvent(false);
+            private NonAgileClass nonAgileClass;
         }
 
 

src/WinRT.Runtime/Interop/EventHandlerEventSource.cs

@@ -34,6 +34,13 @@ public EventHandlerEventSource(
         {
         }
 
+        public EventHandlerEventSource(
+            IObjectReference objectReference,
+            int vtableIndexForAddHandler)
+            : base(objectReference, vtableIndexForAddHandler)
+        {
+        }
+
         /// <inheritdoc/>
         protected override ObjectReferenceValue CreateMarshaler(EventHandler del) =>
             ABI.System.EventHandler.CreateMarshaler2(del);

src/WinRT.Runtime/Interop/EventHandlerEventSource{T}.cs

@@ -35,6 +35,11 @@ public EventHandlerEventSource(
         {
         }
 
+        public EventHandlerEventSource(IObjectReference objectReference, int vtableIndexForAddHandler)
+            : base(objectReference, vtableIndexForAddHandler)
+        {
+        }
+
         /// <inheritdoc/>
         protected override ObjectReferenceValue CreateMarshaler(EventHandler<T> del) =>
             ABI.System.EventHandler<T>.CreateMarshaler2(del);

src/WinRT.Runtime/Interop/EventSource{TDelegate}.cs

@@ -28,15 +28,57 @@ abstract unsafe class EventSource<TDelegate>
         where TDelegate : class, MulticastDelegate
     {
         private readonly IObjectReference _objectReference;
-        private readonly int _index;
-#if NET
-        private readonly delegate* unmanaged[Stdcall]<IntPtr, IntPtr, EventRegistrationToken*, int> _addHandler;
-#else
-        private readonly delegate* unmanaged[Stdcall]<IntPtr, IntPtr, out EventRegistrationToken, int> _addHandler;
-#endif
-        private readonly delegate* unmanaged[Stdcall]<IntPtr, EventRegistrationToken, int> _removeHandler;
-        private global::System.WeakReference<object>? _state;
-
+        private readonly int _index;
+        private readonly long _vtableIndexForAddHandler;
+#if NET
+        private readonly delegate* unmanaged[Stdcall]<IntPtr, IntPtr, EventRegistrationToken*, int> _addHandler;
+#else
+        private readonly delegate* unmanaged[Stdcall]<IntPtr, IntPtr, out EventRegistrationToken, int> _addHandler;
+#endif
+        private readonly delegate* unmanaged[Stdcall]<IntPtr, EventRegistrationToken, int> _removeHandler;
+        private global::System.WeakReference<object>? _state;
+
+        // The add / remove handlers given to us can be for an object which is not agile meaning we may have
+        // a proxy which we need to call through.  Due to this, we store the offset of the add handler
+        // we are given rather than directly caching it. Then we use that offset, to determine the add and
+        // remove handler to call based on the pointer from the current context.
+#if NET
+        private delegate* unmanaged[Stdcall]<IntPtr, IntPtr, EventRegistrationToken*, int> AddHandler
+#else
+        private delegate* unmanaged[Stdcall]<IntPtr, IntPtr, out EventRegistrationToken, int> AddHandler
+#endif
+        {
+            get
+            {
+                if (_addHandler is not null)
+                {
+                    return _addHandler;
+                }
+
+                var thisPtr = _objectReference.ThisPtr;
+#if NET
+                return (*(delegate* unmanaged[Stdcall]<IntPtr, IntPtr, global::WinRT.EventRegistrationToken*, int>**)thisPtr)[_vtableIndexForAddHandler];
+#else
+                return (*(delegate* unmanaged[Stdcall]<IntPtr, IntPtr, out EventRegistrationToken, int>**)thisPtr)[_vtableIndexForAddHandler];
+#endif
+            }
+        }
+
+        private delegate* unmanaged[Stdcall]<IntPtr, EventRegistrationToken, int> RemoveHandler
+        {
+            get
+            {
+                if (_removeHandler is not null)
+                {
+                    return _removeHandler;
+                }
+
+                var thisPtr = _objectReference.ThisPtr;
+                // Add 1 to the offset to get remove handler from add handler offset.
+                return (*(delegate* unmanaged[Stdcall]<IntPtr, EventRegistrationToken, int>**)thisPtr)[_vtableIndexForAddHandler  + 1];
+            }
+        }
+
         /// <summary>
         /// Creates a new <see cref="EventSource{TDelegate}"/> instance with the specified parameters.
         /// </summary>
@@ -55,10 +97,42 @@ protected EventSource(
             int index = 0)
         {
             _objectReference = objectReference;
-            _addHandler = addHandler;
-            _removeHandler = removeHandler;
             _index = index;
-            _state = EventSourceCache.GetState(objectReference, index);
+            _state = EventSourceCache.GetState(objectReference, index);
+
+            // If this isn't a free threaded object, we can't cache the handlers due to we
+            // might be accessing it from a different context which would have its own add handler address
+            // for the proxy. So caching it would end up calling the wrong vtable.
+            // We instead use it to calculate the vtable offset for the add handler to use later.
+            if (objectReference is IObjectReferenceWithContext)
+            {
+                int vtableIndexForAddHandler = 0;
+                while ((*(void***)objectReference.ThisPtr)[vtableIndexForAddHandler] != addHandler)
+                {
+                    vtableIndexForAddHandler++;
+                }
+                _vtableIndexForAddHandler = vtableIndexForAddHandler;
+            }
+            else
+            {
+                _addHandler = addHandler;
+                _removeHandler = removeHandler;
+            }
+        }
+
+        /// <summary>
+        /// Creates a new <see cref="EventSource{TDelegate}"/> instance with the specified parameters.
+        /// </summary>
+        /// <param name="objectReference">The <see cref="IObjectReference"/> instance holding the event.</param>
+        /// <param name="vtableIndexForAddHandler">The vtable index for the add handler of the event being managed.</param>
+        protected EventSource(
+            IObjectReference objectReference,
+            int vtableIndexForAddHandler)
+        {
+            _objectReference = objectReference;
+            _index = vtableIndexForAddHandler;
+            _state = EventSourceCache.GetState(objectReference, vtableIndexForAddHandler);
+            _vtableIndexForAddHandler = vtableIndexForAddHandler;
         }
 
         /// <summary>
@@ -116,9 +190,9 @@ public void Subscribe(TDelegate handler)
 
                         EventRegistrationToken token;
 #if NET
-                        ExceptionHelpers.ThrowExceptionForHR(_addHandler(_objectReference.ThisPtr, nativeDelegate, &token));
+                        ExceptionHelpers.ThrowExceptionForHR(AddHandler(_objectReference.ThisPtr, nativeDelegate, &token));
 #else
-                        ExceptionHelpers.ThrowExceptionForHR(_addHandler(_objectReference.ThisPtr, nativeDelegate, out token));
+                        ExceptionHelpers.ThrowExceptionForHR(AddHandler(_objectReference.ThisPtr, nativeDelegate, out token));
 #endif
                         global::System.GC.KeepAlive(_objectReference);
                         state.token = token;
@@ -157,7 +231,7 @@ public void Unsubscribe(TDelegate handler)
 
         private void UnsubscribeFromNative(EventSourceState<TDelegate> state)
         {
-            ExceptionHelpers.ThrowExceptionForHR(_removeHandler(_objectReference.ThisPtr, state.token));
+            ExceptionHelpers.ThrowExceptionForHR(RemoveHandler(_objectReference.ThisPtr, state.token));
             global::System.GC.KeepAlive(_objectReference);
             state.Dispose();
             _state = null;

src/WinRT.Runtime/MatchingRefApiCompatBaseline.net8.0.txt

@@ -11,4 +11,7 @@ TypesMustExist : Type 'WinRT.DynamicWindowsRuntimeCastAttribute' does not exist
 MembersMustExist : Member 'public System.Guid WinRT.Interop.IID.IID_IWeakReferenceSource.get()' does not exist in the reference but it does exist in the implementation.
 CannotRemoveAttribute : Attribute 'System.ComponentModel.EditorBrowsableAttribute' exists on 'WinRT.Marshaler<T>' in the implementation but not the reference.
 TypesMustExist : Type 'WinRT.MarshalGenericHelper<T>' does not exist in the reference but it does exist in the implementation.
-Total Issues: 12
+MembersMustExist : Member 'protected void ABI.WinRT.Interop.EventSource<TDelegate>..ctor(WinRT.IObjectReference, System.Int32)' does not exist in the reference but it does exist in the implementation.
+MembersMustExist : Member 'public void ABI.WinRT.Interop.EventHandlerEventSource<T>..ctor(WinRT.IObjectReference, System.Int32)' does not exist in the reference but it does exist in the implementation.
+MembersMustExist : Member 'public void ABI.WinRT.Interop.EventHandlerEventSource..ctor(WinRT.IObjectReference, System.Int32)' does not exist in the reference but it does exist in the implementation.
+Total Issues: 15

src/WinRT.Runtime/Projections/ICommand.net5.cs

@@ -67,11 +67,7 @@ public static unsafe EventHandlerEventSource Get_CanExecuteChanged2(IObjectRefer
         {
             return CanExecuteChanged.GetValue(thisObj, (key) =>
             {
-                var ThisPtr = obj.ThisPtr;
-
-                return new EventHandlerEventSource(obj,
-                    (*(delegate* unmanaged[Stdcall]<IntPtr, IntPtr, global::WinRT.EventRegistrationToken*, int>**)ThisPtr)[6],
-                    (*(delegate* unmanaged[Stdcall]<IntPtr, global::WinRT.EventRegistrationToken, int>**)ThisPtr)[7]);
+                return new EventHandlerEventSource(obj, 6);
             });
         }
     }

src/WinRT.Runtime/Projections/ICommand.netstandard2.0.cs

@@ -146,11 +146,7 @@ public ICommand(IObjectReference obj) : this(obj.As<Vftbl>()) { }
         public ICommand(ObjectReference<Vftbl> obj)
         {
             _obj = obj;
-
-            _CanExecuteChanged =
-                new EventHandlerEventSource(_obj,
-                _obj.Vftbl.add_CanExecuteChanged_0,
-                _obj.Vftbl.remove_CanExecuteChanged_1);
+            _CanExecuteChanged = new EventHandlerEventSource(_obj, 6);
         }