From 1dad81f08c2f0ebd7188d49353f33aca64174952 Mon Sep 17 00:00:00 2001
From: Giuseppe Lippolis <workgroupengineering@outlook.com>
Date: Wed, 4 Sep 2024 11:16:12 +0200
Subject: [PATCH] fix(Windows): WindowsInputPane memory leak

---
 .../Avalonia.Win32/Input/WindowsInputPane.cs  | 62 +++++++++++++++----
 1 file changed, 50 insertions(+), 12 deletions(-)

diff --git a/src/Windows/Avalonia.Win32/Input/WindowsInputPane.cs b/src/Windows/Avalonia.Win32/Input/WindowsInputPane.cs
index 73195ad500..879176f0cf 100644
--- a/src/Windows/Avalonia.Win32/Input/WindowsInputPane.cs
+++ b/src/Windows/Avalonia.Win32/Input/WindowsInputPane.cs
@@ -1,6 +1,5 @@
 using System;
 using Avalonia.Controls.Platform;
-using Avalonia.MicroCom;
 using Avalonia.Win32.Interop;
 using Avalonia.Win32.Win32Com;
 using Avalonia.Win32.WinRT;
@@ -11,16 +10,17 @@ namespace Avalonia.Win32.Input;
 internal unsafe class WindowsInputPane : InputPaneBase, IDisposable
 {
     private static readonly Lazy<bool> s_inputPaneSupported = new(() =>
-        WinRTApiInformation.IsTypePresent("Windows.UI.ViewManagement.InputPane")); 
+        WinRTApiInformation.IsTypePresent("Windows.UI.ViewManagement.InputPane"));
 
     // GUID: D5120AA3-46BA-44C5-822D-CA8092C1FC72
     private static readonly Guid CLSID_FrameworkInputPane = new(0xD5120AA3, 0x46BA, 0x44C5, 0x82, 0x2D, 0xCA, 0x80, 0x92, 0xC1, 0xFC, 0x72);
     // GUID: 5752238B-24F0-495A-82F1-2FD593056796
-    private static readonly Guid SID_IFrameworkInputPane  = new(0x5752238B, 0x24F0, 0x495A, 0x82, 0xF1, 0x2F, 0xD5, 0x93, 0x05, 0x67, 0x96);
+    private static readonly Guid SID_IFrameworkInputPane = new(0x5752238B, 0x24F0, 0x495A, 0x82, 0xF1, 0x2F, 0xD5, 0x93, 0x05, 0x67, 0x96);
 
-    private readonly WindowImpl _windowImpl;
+    private WindowImpl _windowImpl;
     private IFrameworkInputPane? _inputPane;
     private readonly uint _cookie;
+    private Handler _eventHandler;
 
     private WindowsInputPane(WindowImpl windowImpl)
     {
@@ -31,12 +31,11 @@ internal unsafe class WindowsInputPane : InputPaneBase, IDisposable
             _inputPane = inputPane.CloneReference();
         }
 
-        using (var handler = new Handler(this))
-        {
-            uint cookie = 0;
-            _inputPane.AdviseWithHWND(windowImpl.Handle.Handle, handler, &cookie);
-            _cookie = cookie;
-        }
+        _eventHandler = new Handler(this);
+        uint cookie = 0;
+        _inputPane.AdviseWithHWND(windowImpl.Handle.Handle, _eventHandler, &cookie);
+        _cookie = cookie;
+
     }
 
     public static WindowsInputPane? TryCreate(WindowImpl windowImpl)
@@ -72,6 +71,9 @@ internal unsafe class WindowsInputPane : InputPaneBase, IDisposable
 
     public void Dispose()
     {
+        _windowImpl = null!;
+        _eventHandler?.Dispose();
+        _eventHandler = null!;
         if (_inputPane is not null)
         {
             if (_cookie != 0)
@@ -84,9 +86,45 @@ internal unsafe class WindowsInputPane : InputPaneBase, IDisposable
         }
     }
 
-    private class Handler : CallbackBase, IFrameworkInputPaneHandler
+    private class Handler : IUnknown, IMicroComShadowContainer, IFrameworkInputPaneHandler
     {
-        private readonly WindowsInputPane _pane;
+        private WindowsInputPane _pane;
+
+        private readonly object _lock = new();
+        private bool _destroyed;
+        public bool IsDestroyed => _destroyed;
+
+        public void Dispose()
+        {
+            lock (_lock)
+            {
+                DestroyIfNeeded();
+            }
+        }
+
+        void DestroyIfNeeded()
+        {
+            if (_destroyed)
+                return;
+            _destroyed = true;
+            Shadow?.Dispose();
+            Shadow = null;
+            _pane = null!;
+        }
+
+        public MicroComShadow? Shadow { get; set; }
+        public void OnReferencedFromNative()
+        {
+
+        }
+
+        public void OnUnreferencedFromNative()
+        {
+            lock (_lock)
+            {
+                DestroyIfNeeded();
+            }
+        }
 
         public Handler(WindowsInputPane pane) => _pane = pane;
         public void Showing(UnmanagedMethods.RECT* rect, int _) => _pane.OnStateChanged(true, *rect);