Smoother OAuth Cancellation Handling + EndSession Prefab_ConnectWalle…

…t Option (#187)
thirdweb-dev · May 17, 2024 · 890e617 · 890e617
1 parent 3c17230
commit 890e617
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 14 deletions.
diff --git a/Assets/Thirdweb/Core/Scripts/Browser/AndroidBrowser.cs b/Assets/Thirdweb/Core/Scripts/Browser/AndroidBrowser.cs
@@ -39,6 +39,14 @@ public async Task<BrowserResult> Login(string loginUrl, string customScheme, Can
                     return new BrowserResult(BrowserStatus.Timeout, null, "The operation timed out.");
                 }
             }
+            catch (TaskCanceledException)
+            {
+                return new BrowserResult(BrowserStatus.UserCanceled, null, "The operation was cancelled.");
+            }
+            catch (Exception ex)
+            {
+                return new BrowserResult(BrowserStatus.UnknownError, null, $"An error occurred: {ex.Message}");
+            }
             finally
             {
                 Application.deepLinkActivated -= OnDeepLinkActivated;
@@ -47,17 +55,17 @@ public async Task<BrowserResult> Login(string loginUrl, string customScheme, Can
 
         private void OpenURL(string url)
         {
-            AndroidJavaClass thirdwebActivityClass = new("com.unity3d.player.UnityPlayer");
-            AndroidJavaObject thirdwebActivity = thirdwebActivityClass.GetStatic<AndroidJavaObject>("currentActivity");
-            thirdwebActivity.Call("OpenCustomTab", url);
+            AndroidJavaClass unityPlayer = new("com.unity3d.player.UnityPlayer");
+            AndroidJavaObject activity = unityPlayer.GetStatic<AndroidJavaObject>("currentActivity");
+            activity.Call("OpenCustomTab", url);
         }
 
         private void OnDeepLinkActivated(string url)
         {
-            if (!url.StartsWith(_customScheme))
+            if (_taskCompletionSource.Task.IsCanceled || !url.StartsWith(_customScheme))
                 return;
 
-            _taskCompletionSource.SetResult(new BrowserResult(BrowserStatus.Success, url));
+            _taskCompletionSource.TrySetResult(new BrowserResult(BrowserStatus.Success, url));
         }
     }
 }

diff --git a/Assets/Thirdweb/Core/Scripts/Browser/StandaloneBrowser.cs b/Assets/Thirdweb/Core/Scripts/Browser/StandaloneBrowser.cs
@@ -12,6 +12,8 @@ public class StandaloneBrowser : IThirdwebBrowser
     {
         private TaskCompletionSource<BrowserResult> _taskCompletionSource;
 
+        private readonly HttpListener httpListener = new();
+
         private readonly string closePageResponse =
             @"
             <html>
@@ -57,14 +59,17 @@ public async Task<BrowserResult> Login(string loginUrl, string redirectUrl, Canc
             cancellationToken.Register(() =>
             {
                 _taskCompletionSource?.TrySetCanceled();
+                StopHttpListener();
             });
 
-            using var httpListener = new HttpListener();
-
             try
             {
                 redirectUrl = AddForwardSlashIfNecessary(redirectUrl);
-                httpListener.Prefixes.Add(redirectUrl);
+                if (httpListener.Prefixes.Count == 0 || !httpListener.Prefixes.Contains(redirectUrl))
+                {
+                    httpListener.Prefixes.Clear();
+                    httpListener.Prefixes.Add(redirectUrl);
+                }
                 httpListener.Start();
                 httpListener.BeginGetContext(IncomingHttpRequest, httpListener);
 
@@ -80,15 +85,35 @@ public async Task<BrowserResult> Login(string loginUrl, string redirectUrl, Canc
                     return new BrowserResult(BrowserStatus.Timeout, null, "The operation timed out.");
                 }
             }
+            catch (TaskCanceledException)
+            {
+                return new BrowserResult(BrowserStatus.UserCanceled, null, "The operation was cancelled.");
+            }
+            catch (Exception ex)
+            {
+                return new BrowserResult(BrowserStatus.UnknownError, null, $"An error occurred: {ex.Message}");
+            }
             finally
+            {
+                StopHttpListener();
+            }
+        }
+
+        private void StopHttpListener()
+        {
+            if (httpListener != null && httpListener.IsListening)
             {
                 httpListener.Stop();
+                httpListener.Close();
             }
         }
 
         private void IncomingHttpRequest(IAsyncResult result)
         {
             var httpListener = (HttpListener)result.AsyncState;
+            if (!httpListener.IsListening)
+                return;
+
             var httpContext = httpListener.EndGetContext(result);
             var httpRequest = httpContext.Request;
             var httpResponse = httpContext.Response;
@@ -104,10 +129,9 @@ private void IncomingHttpRequest(IAsyncResult result)
 
         private string AddForwardSlashIfNecessary(string url)
         {
-            string forwardSlash = "/";
-            if (!url.EndsWith(forwardSlash))
+            if (!url.EndsWith("/"))
             {
-                url += forwardSlash;
+                url += "/";
             }
             return url;
         }

diff --git a/Assets/Thirdweb/Core/Scripts/WalletsUI/InAppWalletUI.cs b/Assets/Thirdweb/Core/Scripts/WalletsUI/InAppWalletUI.cs
@@ -12,6 +12,7 @@
 using System.Linq;
 using System.Collections.Generic;
 using UnityEngine.Events;
+using System.Threading;
 
 namespace Thirdweb.Wallets
 {
@@ -37,6 +38,7 @@ public class InAppWalletUI : MonoBehaviour
         protected Exception _exception;
         protected string _callbackUrl;
         protected string _customScheme;
+        protected CancellationTokenSource _cancellationTokenSource;
 
         #endregion
 
@@ -141,8 +143,16 @@ public virtual async Task<User> Connect(EmbeddedWallet embeddedWallet, string em
             return _user;
         }
 
+        [ContextMenu("Cancel")]
         public virtual void Cancel()
         {
+            if (_cancellationTokenSource != null && !_cancellationTokenSource.IsCancellationRequested)
+            {
+                _cancellationTokenSource.Cancel();
+                _cancellationTokenSource.Dispose();
+                _cancellationTokenSource = null;
+            }
+
             _exception = new UnityException("User cancelled");
         }
 
@@ -285,11 +295,27 @@ public virtual async Task LoginWithOauth(string authProviderStr)
 
             string redirectUrl = Application.isMobilePlatform ? _customScheme : "http://localhost:8789/";
             CrossPlatformBrowser browser = new();
-            var browserResult = await browser.Login(loginUrl, redirectUrl);
-            if (browserResult.status != BrowserStatus.Success)
+            _cancellationTokenSource?.Cancel();
+            _cancellationTokenSource?.Dispose();
+            _cancellationTokenSource = new CancellationTokenSource();
+            var browserResult = await browser.Login(loginUrl, redirectUrl, _cancellationTokenSource.Token);
+
+            if (browserResult.status == BrowserStatus.UserCanceled)
+            {
+                _cancellationTokenSource?.Dispose();
+                _cancellationTokenSource = null;
+                ThirdwebDebug.LogWarning("User cancelled login");
+                return;
+            }
+            else if (browserResult.status != BrowserStatus.Success)
+            {
                 _exception = new UnityException($"Failed to login with {authProviderStr}: {browserResult.status} | {browserResult.error}");
+                return;
+            }
             else
+            {
                 _callbackUrl = browserResult.callbackUrl;
+            }
 
             await new WaitUntil(() => _callbackUrl != null);
 

diff --git a/Assets/Thirdweb/Examples/Prefabs/Prefab_ConnectWallet.prefab b/Assets/Thirdweb/Examples/Prefabs/Prefab_ConnectWallet.prefab
@@ -4788,6 +4788,7 @@ MonoBehaviour:
   m_EditorClassIdentifier: 
   enabledWalletProviders: 07000000000000000100000002000000060000000300000004000000
   useSmartWallets: 0
+  endSessionOnDisconnect: 0
   onStart:
     m_PersistentCalls:
       m_Calls:

diff --git a/Assets/Thirdweb/Examples/Scripts/Prefabs/Prefab_ConnectWallet.cs b/Assets/Thirdweb/Examples/Scripts/Prefabs/Prefab_ConnectWallet.cs
@@ -36,6 +36,9 @@ public class WalletProviderUIDictionary : SerializableDictionaryBase<WalletProvi
         [Header("Use ERC-4337 (Account Abstraction) compatible smart wallets.\nEnabling this will connect user to the associated smart wallet as per your ThirwebManager settings.")]
         public bool useSmartWallets = false;
 
+        [Header("End session on disconnect. If enabled, user will have to re-authenticate on next connect.")]
+        public bool endSessionOnDisconnect = false;
+
         [Header("Events")]
         public UnityEvent onStart;
         public UnityEvent<WalletConnection> onConnectionRequested;
@@ -165,7 +168,7 @@ public async void Disconnect()
             {
                 _address = null;
                 _password = null;
-                await ThirdwebManager.Instance.SDK.Wallet.Disconnect(endSession: false);
+                await ThirdwebManager.Instance.SDK.Wallet.Disconnect(endSession: endSessionOnDisconnect);
                 onDisconnected.Invoke();
             }
             catch (System.Exception e)