From 852017f7b295065a0ed96dd7b4c5bd4c8da9d676 Mon Sep 17 00:00:00 2001
From: John Yun <140559986+sfc-gh-ext-simba-jy@users.noreply.github.com>
Date: Mon, 30 Sep 2024 18:46:36 -0700
Subject: [PATCH] SNOW-1694143: Fix sso token authentication (#919)
 (Co-authored-by: supermacro)

---
 index.d.ts                   | 12 +++++++++++-
 lib/connection/connection.js |  3 ++-
 lib/services/sf.js           | 11 +++++------
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/index.d.ts b/index.d.ts
index deacf23c4..566ed877b 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -484,6 +484,16 @@ declare module 'snowflake-sdk' {
          */
         credentialCacheDir?: string;
 
+        /**
+         * The option to enable the MFA token. The default value is false.
+         */
+        clientRequestMFAToken?: boolean;
+
+        /**
+         * The option to enable the SSO token. The default value is false.
+         */
+        clientStoreTemporaryCredential?: boolean;
+
         /**
          *  The option to include the passcode from DUO into the password.
          */
@@ -492,7 +502,7 @@ declare module 'snowflake-sdk' {
         /**
          *  The option to pass passcode from DUO.
          */
-        passcode?: string
+        passcode?: string;
     }
 
     export interface Connection {
diff --git a/lib/connection/connection.js b/lib/connection/connection.js
index 55a30ca43..42815f775 100644
--- a/lib/connection/connection.js
+++ b/lib/connection/connection.js
@@ -305,7 +305,8 @@ function Connection(context) {
     }
    
     // Get authenticator to use
-    const auth = services.sf.getAuthenticator();
+    const auth = Authenticator.getAuthenticator(connectionConfig, context.getHttpClient());
+    services.sf.authenticator = auth;
 
     try {
       await initEasyLogging(connectionConfig.clientConfigFile);
diff --git a/lib/services/sf.js b/lib/services/sf.js
index 79ebe559a..f32ecb04f 100644
--- a/lib/services/sf.js
+++ b/lib/services/sf.js
@@ -1185,12 +1185,11 @@ StateConnecting.prototype.continue = function () {
               numRetries,
             };
 
-            auth.reauthenticate(context.options.json, retryOption).then(() => {
-              numRetries = retryOption.numRetries;
-              totalElapsedTime = retryOption.totalElapsedTime;
-              setTimeout(sendRequest, sleep * 1000);
-              return;
-            });
+            await auth.reauthenticate(context.options.json, retryOption);
+            numRetries = retryOption.numRetries;
+            totalElapsedTime = retryOption.totalElapsedTime;
+            setTimeout(sendRequest, sleep * 1000);
+            return;
           } else {
             if (auth instanceof AuthKeypair) {
               Logger.getInstance().debug('AuthKeyPair authentication requires token refresh.');