From 84251a6a5ed163d2f71cd72bb3e6935d2d983afc Mon Sep 17 00:00:00 2001
From: Eugene Wang <8755753+soukoku@users.noreply.github.com>
Date: Sun, 2 Apr 2023 21:45:48 -0400
Subject: [PATCH] Added winform/wpf reg methods

---
 samples/WinForm32/Form1.cs         |  7 ++---
 src/NTwain/TwainSession.Sources.cs |  2 +-
 src/NTwain/TwainSession.Windows.cs | 49 +++++++++++++++++++++++++++---
 3 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/samples/WinForm32/Form1.cs b/samples/WinForm32/Form1.cs
index cee6df7..ea2a00d 100644
--- a/samples/WinForm32/Form1.cs
+++ b/samples/WinForm32/Form1.cs
@@ -16,7 +16,7 @@ namespace WinForm32
     {
       InitializeComponent();
       Text += TwainPlatform.Is32bit ? " 32bit" : " 64bit";
-      TwainPlatform.PreferLegacyDSM = true;
+      TwainPlatform.PreferLegacyDSM = false;
 
       twain = new TwainSession(Assembly.GetExecutingAssembly().Location);
       twain.StateChanged += Twain_StateChanged;
@@ -46,7 +46,7 @@ namespace WinForm32
 
       var hwnd = this.Handle;
       var rc = twain.OpenDSM(hwnd);
-      Application.AddMessageFilter(twain);
+      twain.AddWinformFilter();
       Debug.WriteLine($"OpenDSM={rc}");
     }
 
@@ -54,8 +54,7 @@ namespace WinForm32
     {
       var finalState = twain.TryStepdown(STATE.S2);
       Debug.WriteLine($"Stepdown result state={finalState}");
-
-      Application.RemoveMessageFilter(twain);
+      twain.RemoveWinformFilter();
       base.OnClosing(e);
     }
 
diff --git a/src/NTwain/TwainSession.Sources.cs b/src/NTwain/TwainSession.Sources.cs
index 27e303c..fd18491 100644
--- a/src/NTwain/TwainSession.Sources.cs
+++ b/src/NTwain/TwainSession.Sources.cs
@@ -67,7 +67,7 @@ namespace NTwain
     /// <returns></returns>
     public STS SetDefaultSource(TW_IDENTITY_LEGACY source)
     {
-      // TODO: this doesn't work???
+      // this doesn't work on windows legacy twain_32.dll
 
       var rc = DGControl.Identity.Set(ref _appIdentity, ref source);
       if (rc == STS.SUCCESS)
diff --git a/src/NTwain/TwainSession.Windows.cs b/src/NTwain/TwainSession.Windows.cs
index 15b164b..5b450fc 100644
--- a/src/NTwain/TwainSession.Windows.cs
+++ b/src/NTwain/TwainSession.Windows.cs
@@ -7,6 +7,8 @@ using System.Diagnostics;
 using System.Runtime.InteropServices;
 using System.Threading;
 using System.Windows.Forms;
+using System.Windows.Interop;
+using MSG = NTwain.Data.MSG;
 
 namespace NTwain
 {
@@ -15,14 +17,53 @@ namespace NTwain
 
   partial class TwainSession : IMessageFilter
   {
-    // winform use with Application.AddMessageFilter(<TwainSession instance>)
+    private HwndSource? _wpfhook;
+
+    /// <summary>
+    /// Registers this session for use in a Winform UI thread.
+    /// </summary>
+    public void AddWinformFilter()
+    {
+      Application.AddMessageFilter(this);
+    }
+    /// <summary>
+    /// Unregisters this session if previously registered with <see cref="AddWinformFilter"/>.
+    /// </summary>
+    public void RemoveWinformFilter()
+    {
+      Application.RemoveMessageFilter(this);
+    }
+
+    /// <summary>
+    /// Registers this session for use in a WPF UI thread.
+    /// </summary>
+    public void AddWpfHook()
+    {
+      if (_wpfhook == null)
+      {
+        _wpfhook = HwndSource.FromHwnd(_hwnd);
+        _wpfhook.AddHook(WpfHook);
+      }
+    }
+    /// <summary>
+    /// Unregisters this session if previously registered with <see cref="AddWpfHook"/>.
+    /// </summary>
+    public void RemoveWpfHook()
+    {
+      if (_wpfhook != null)
+      {
+        _wpfhook.RemoveHook(WpfHook);
+        _wpfhook = null;
+      }
+    }
+
     bool System.Windows.Forms.IMessageFilter.PreFilterMessage(ref System.Windows.Forms.Message m)
     {
       return CheckIfTwainMessage(m.HWnd, m.Msg, m.WParam, m.LParam);
     }
 
-    // wpf use
-    IntPtr HwndSourceHook(IntPtr hwnd, int msg, IntPtr wParam, IntPtr lParam, ref bool handled)
+    // wpf use with HwndSource.FromHwnd(Handle).AddHook(
+    IntPtr WpfHook(IntPtr hwnd, int msg, IntPtr wParam, IntPtr lParam, ref bool handled)
     {
       handled = CheckIfTwainMessage(hwnd, msg, wParam, lParam);
       return IntPtr.Zero;
@@ -32,7 +73,7 @@ namespace NTwain
     {
       // this handles the message from a typical WndProc message loop and check if it's from the TWAIN source.
       bool handled = false;
-      if (_state >= Data.STATE.S5)
+      if (_state >= STATE.S5)
       {
         TW_EVENT evt = default;
         try