From 8d30e9066f0e850a72aa404115ac8529c7960799 Mon Sep 17 00:00:00 2001
From: soukoku <yinchun.wang@gmail.com>
Date: Thu, 8 Jan 2015 07:05:22 -0500
Subject: [PATCH] Changed when callbacks are registered.

---
 NTwain/DataSource.cs                      |  2 +
 NTwain/Internals/ITwainSessionInternal.cs |  2 +
 NTwain/TwainSessionInternal.cs            | 82 ++++++++++++-----------
 3 files changed, 46 insertions(+), 40 deletions(-)

diff --git a/NTwain/DataSource.cs b/NTwain/DataSource.cs
index 7309a20..9355a2e 100644
--- a/NTwain/DataSource.cs
+++ b/NTwain/DataSource.cs
@@ -38,6 +38,7 @@ namespace NTwain
                 Debug.WriteLine(string.Format(CultureInfo.InvariantCulture, "Thread {0}: OpenSource.", Thread.CurrentThread.ManagedThreadId));
 
                 rc = _session.DGControl.Identity.OpenDS(this);
+                _session.UpdateCallback();
             });
             return rc;
         }
@@ -58,6 +59,7 @@ namespace NTwain
                 {
                     SupportedCaps = null;
                 }
+                _session.UpdateCallback();
             });
             return rc;
         }
diff --git a/NTwain/Internals/ITwainSessionInternal.cs b/NTwain/Internals/ITwainSessionInternal.cs
index f296b83..aca2372 100644
--- a/NTwain/Internals/ITwainSessionInternal.cs
+++ b/NTwain/Internals/ITwainSessionInternal.cs
@@ -36,6 +36,8 @@ namespace NTwain.Internals
 
         void ChangeCurrentSource(DataSource source);
 
+        void UpdateCallback();
+
         ReturnCode DisableSource();
 
         void SafeSyncableRaiseEvent(DataTransferredEventArgs e);
diff --git a/NTwain/TwainSessionInternal.cs b/NTwain/TwainSessionInternal.cs
index 36fcd59..b182de8 100644
--- a/NTwain/TwainSessionInternal.cs
+++ b/NTwain/TwainSessionInternal.cs
@@ -27,6 +27,44 @@ namespace NTwain
         /// <value>The app id.</value>
         TWIdentity ITwainSessionInternal.AppId { get { return _appId; } }
 
+        void ITwainSessionInternal.UpdateCallback()
+        {
+            if (State < 4)
+            {
+                _callbackObj = null;
+            }
+            else
+            {
+                ReturnCode rc = ReturnCode.Failure;
+                // per the spec (8-10) apps for 2.2 or higher uses callback2 so try this first
+                if (_appId.ProtocolMajor > 2 || (_appId.ProtocolMajor >= 2 && _appId.ProtocolMinor >= 2))
+                {
+                    var cb = new TWCallback2(HandleCallback);
+                    rc = ((ITwainSessionInternal)this).DGControl.Callback2.RegisterCallback(cb);
+
+                    if (rc == ReturnCode.Success)
+                    {
+                        Debug.WriteLine("Registered callback2 OK.");
+                        _callbackObj = cb;
+                    }
+                }
+
+                if (rc != ReturnCode.Success)
+                {
+                    // always try old callback
+                    var cb = new TWCallback(HandleCallback);
+
+                    rc = ((ITwainSessionInternal)this).DGControl.Callback.RegisterCallback(cb);
+
+                    if (rc == ReturnCode.Success)
+                    {
+                        Debug.WriteLine("Registered callback OK.");
+                        _callbackObj = cb;
+                    }
+                }
+            }
+        }
+
         void ITwainSessionInternal.ChangeState(int newState, bool notifyChange)
         {
             _state = newState;
@@ -122,54 +160,19 @@ namespace NTwain
             {
                 Debug.WriteLine(string.Format(CultureInfo.InvariantCulture, "Thread {0}: EnableSource with {1}.", Thread.CurrentThread.ManagedThreadId, mode));
 
-                // per the spec (8-10) app v2.2 or higher uses callback2
-                if (_appId.ProtocolMajor > 2 || (_appId.ProtocolMajor >= 2 && _appId.ProtocolMinor >= 2))
-                {
-                    var cb = new TWCallback2(HandleCallback);
-                    var rc2 = ((ITwainSessionInternal)this).DGControl.Callback2.RegisterCallback(cb);
-
-                    if (rc2 == ReturnCode.Success)
-                    {
-                        Debug.WriteLine("Registered callback2 OK.");
-                        _callbackObj = cb;
-                    }
-                }
-                else
-                {
-                    // all else try old callback
-                    var cb = new TWCallback(HandleCallback);
-
-                    var rc2 = ((ITwainSessionInternal)this).DGControl.Callback.RegisterCallback(cb);
-
-                    if (rc2 == ReturnCode.Success)
-                    {
-                        Debug.WriteLine("Registered callback OK.");
-                        _callbackObj = cb;
-                    }
-                }
 
                 _twui = new TWUserInterface();
                 _twui.ShowUI = mode == SourceEnableMode.ShowUI;
                 _twui.ModalUI = modal;
                 _twui.hParent = windowHandle;
 
-                try
+                if (mode == SourceEnableMode.ShowUIOnly)
                 {
-                    if (mode == SourceEnableMode.ShowUIOnly)
-                    {
-                        rc = ((ITwainSessionInternal)this).DGControl.UserInterface.EnableDSUIOnly(_twui);
-                    }
-                    else
-                    {
-                        rc = ((ITwainSessionInternal)this).DGControl.UserInterface.EnableDS(_twui);
-                    }
+                    rc = ((ITwainSessionInternal)this).DGControl.UserInterface.EnableDSUIOnly(_twui);
                 }
-                finally
+                else
                 {
-                    if (rc != ReturnCode.Success)
-                    {
-                        _callbackObj = null;
-                    }
+                    rc = ((ITwainSessionInternal)this).DGControl.UserInterface.EnableDS(_twui);
                 }
             });
             return rc;
@@ -191,7 +194,6 @@ namespace NTwain
                         rc = ((ITwainSessionInternal)this).DGControl.UserInterface.DisableDS(_twui);
                         if (rc == ReturnCode.Success)
                         {
-                            _callbackObj = null;
                             SafeAsyncSyncableRaiseOnEvent(OnSourceDisabled, SourceDisabled);
                         }
                     });