SNOW-923398: improve robustness in handling authentication response (#…

…1742)

DESCRIPTION.md

@@ -10,10 +10,11 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 - v3.2.1(September 26,2023)
 
-  - Fixed a bug where url port and path were ignore in private link oscp retry.
+  - Fixed a bug where url port and path were ignored in private link oscp retry.
   - Added thread safety in telemetry when instantiating multiple connections concurrently.
   - Bumped platformdirs dependency from >=2.6.0,<3.9.0 to >=2.6.0,<4.0.0.0 and made necessary changes to allow this.
   - Removed the deprecation warning from the vendored urllib3 about urllib3.contrib.pyopenssl deprecation.
+  - Improved robustness in handling authentication response.
 
 - v3.2.0(September 06,2023)
 

src/snowflake/connector/auth/_auth.py

@@ -282,11 +282,11 @@ def authenticate(
             )
 
         # waiting for MFA authentication
-        if ret["data"].get("nextAction") in (
+        if ret["data"] and ret["data"].get("nextAction") in (
             "EXT_AUTHN_DUO_ALL",
             "EXT_AUTHN_DUO_PUSH_N_PASSCODE",
         ):
-            body["inFlightCtx"] = ret["data"]["inFlightCtx"]
+            body["inFlightCtx"] = ret["data"].get("inFlightCtx")
             body["data"]["EXT_AUTHN_DUO_METHOD"] = "push"
             self.ret = {"message": "Timeout", "data": {}}
 
@@ -310,9 +310,13 @@ def post_request_wrapper(self, url, headers, body) -> None:
                 t.join(timeout=timeout)
 
             ret = self.ret
-            if ret and ret["data"].get("nextAction") == "EXT_AUTHN_SUCCESS":
+            if (
+                ret
+                and ret["data"]
+                and ret["data"].get("nextAction") == "EXT_AUTHN_SUCCESS"
+            ):
                 body = copy.deepcopy(body_template)
-                body["inFlightCtx"] = ret["data"]["inFlightCtx"]
+                body["inFlightCtx"] = ret["data"].get("inFlightCtx")
                 # final request to get tokens
                 ret = self._rest._post_request(
                     url,
@@ -321,7 +325,7 @@ def post_request_wrapper(self, url, headers, body) -> None:
                     timeout=self._rest._connection.login_timeout,
                     socket_timeout=self._rest._connection.login_timeout,
                 )
-            elif not ret or not ret["data"].get("token"):
+            elif not ret or not ret["data"] or not ret["data"].get("token"):
                 # not token is returned.
                 Error.errorhandler_wrapper(
                     self._rest._connection,
@@ -343,10 +347,10 @@ def post_request_wrapper(self, url, headers, body) -> None:
                 )
                 return session_parameters  # required for unit test
 
-        elif ret["data"].get("nextAction") == "PWD_CHANGE":
+        elif ret["data"] and ret["data"].get("nextAction") == "PWD_CHANGE":
             if callable(password_callback):
                 body = copy.deepcopy(body_template)
-                body["inFlightCtx"] = ret["data"]["inFlightCtx"]
+                body["inFlightCtx"] = ret["data"].get("inFlightCtx")
                 body["data"]["LOGIN_NAME"] = user
                 body["data"]["PASSWORD"] = (
                     auth_instance.password
@@ -411,41 +415,59 @@ def post_request_wrapper(self, url, headers, body) -> None:
             )
         else:
             logger.debug(
-                "token = %s", "******" if ret["data"]["token"] is not None else "NULL"
+                "token = %s",
+                "******"
+                if ret["data"] and ret["data"].get("token") is not None
+                else "NULL",
             )
             logger.debug(
                 "master_token = %s",
-                "******" if ret["data"]["masterToken"] is not None else "NULL",
+                "******"
+                if ret["data"] and ret["data"].get("masterToken") is not None
+                else "NULL",
             )
             logger.debug(
                 "id_token = %s",
-                "******" if ret["data"].get("idToken") is not None else "NULL",
+                "******"
+                if ret["data"] and ret["data"].get("idToken") is not None
+                else "NULL",
             )
             logger.debug(
                 "mfa_token = %s",
-                "******" if ret["data"].get("mfaToken") is not None else "NULL",
+                "******"
+                if ret["data"] and ret["data"].get("mfaToken") is not None
+                else "NULL",
             )
+            if not ret["data"]:
+                Error.errorhandler_wrapper(
+                    None,
+                    None,
+                    Error,
+                    {
+                        "msg": "There is no data in the returning response, please retry the operation."
+                    },
+                )
             self._rest.update_tokens(
-                ret["data"]["token"],
-                ret["data"]["masterToken"],
+                ret["data"].get("token"),
+                ret["data"].get("masterToken"),
                 master_validity_in_seconds=ret["data"].get("masterValidityInSeconds"),
                 id_token=ret["data"].get("idToken"),
                 mfa_token=ret["data"].get("mfaToken"),
             )
             self.write_temporary_credentials(
                 self._rest._host, user, session_parameters, ret
             )
-            if "sessionId" in ret["data"]:
-                self._rest._connection._session_id = ret["data"]["sessionId"]
-            if "sessionInfo" in ret["data"]:
-                session_info = ret["data"]["sessionInfo"]
+            if ret["data"] and "sessionId" in ret["data"]:
+                self._rest._connection._session_id = ret["data"].get("sessionId")
+            if ret["data"] and "sessionInfo" in ret["data"]:
+                session_info = ret["data"].get("sessionInfo")
                 self._rest._connection._database = session_info.get("databaseName")
                 self._rest._connection._schema = session_info.get("schemaName")
                 self._rest._connection._warehouse = session_info.get("warehouseName")
                 self._rest._connection._role = session_info.get("roleName")
-            if "parameters" in ret["data"]:
+            if ret["data"] and "parameters" in ret["data"]:
                 session_parameters.update(
-                    {p["name"]: p["value"] for p in ret["data"]["parameters"]}
+                    {p["name"]: p["value"] for p in ret["data"].get("parameters")}
                 )
             self._rest._connection._update_parameters(session_parameters)
             return session_parameters

test/unit/test_auth.py

@@ -11,6 +11,7 @@
 
 import pytest
 
+import snowflake.connector.errors
 from snowflake.connector.constants import OCSPMode
 from snowflake.connector.description import CLIENT_NAME, CLIENT_VERSION
 from snowflake.connector.network import SnowflakeRestful
@@ -102,7 +103,12 @@ def _mock_auth_mfa_rest_response_failure(url, headers, body, **kwargs):
                 "inFlightCtx": "inFlightCtx",
             },
         }
-
+    elif mock_cnt == 2:
+        ret = {
+            "success": True,
+            "message": None,
+            "data": None,
+        }
     mock_cnt += 1
     return ret
 
@@ -126,6 +132,12 @@ def _mock_auth_mfa_rest_response_timeout(url, headers, body, **kwargs):
     elif mock_cnt == 1:
         time.sleep(10)  # should timeout while here
         ret = {}
+    elif mock_cnt == 2:
+        ret = {
+            "success": True,
+            "message": None,
+            "data": None,
+        }
 
     mock_cnt += 1
     return ret
@@ -168,6 +180,14 @@ def test_auth_mfa(next_action: str):
     auth.authenticate(auth_instance, account, user, timeout=1)
     assert rest._connection.errorhandler.called  # error
 
+    # ret["data"] is none
+    with pytest.raises(snowflake.connector.errors.Error):
+        mock_cnt = 2
+        rest = _init_rest(application, _mock_auth_mfa_rest_response_timeout)
+        auth = Auth(rest)
+        auth_instance = AuthByDefault(password)
+        auth.authenticate(auth_instance, account, user)
+
 
 def _mock_auth_password_change_rest_response(url, headers, body, **kwargs):
     """Test successful case."""