SNOW-1226666: allow disabling saml url check in okta authentication (#1961)

Author: sfc-gh-aling

DESCRIPTION.md

@@ -12,6 +12,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
   - Added support for `token_file_path` connection parameter to read an OAuth token from a file when connecting to Snowflake.
   - Added support for `debug_arrow_chunk` connection parameter to allow debugging raw arrow data in case of arrow data parsing failure.
+  - Added support for `disable_saml_url_check` connection parameter to disable SAML URL check in OKTA authentication.
   - Fixed a bug that OCSP certificate signed using SHA384 algorithm cannot be verified.
   - Fixed a bug that status code shown as uploaded when PUT command failed with 400 error.
 

src/snowflake/connector/auth/okta.py

@@ -315,7 +315,9 @@ def _step5(
             host=conn._rest._host,
             port=conn._rest._port,
         )
-        if not _is_prefix_equal(post_back_url, full_url):
+        if not getattr(conn, "_disable_saml_url_check", False) and not _is_prefix_equal(
+            post_back_url, full_url
+        ):
             Error.errorhandler_wrapper(
                 conn._rest._connection,
                 None,

src/snowflake/connector/connection.py

@@ -285,6 +285,10 @@ def _get_private_bytes_from_file(
         False,
         bool,
     ),  # log raw arrow chunk for debugging purpuse in case there is malformed arrow data
+    "disable_saml_url_check": (
+        False,
+        bool,
+    ),  # disable saml url check in okta authentication
 }
 
 APPLICATION_RE = re.compile(r"[\w\d_]+")

test/unit/mock_utils.py

@@ -32,6 +32,7 @@ def mock_connection(
     network_timeout=None,
     socket_timeout=None,
     backoff_policy=DEFAULT_BACKOFF_POLICY,
+    disable_saml_url_check=False,
 ):
     return MagicMock(
         _login_timeout=login_timeout,
@@ -42,6 +43,7 @@ def mock_connection(
         socket_timeout=socket_timeout,
         _backoff_policy=backoff_policy,
         backoff_policy=backoff_policy,
+        _disable_saml_url_check=disable_saml_url_check,
     )
 
 

test/unit/test_auth_okta.py

@@ -248,7 +248,8 @@ def mock_session_request(*args, **kwargs):
         assert not rest._connection.errorhandler.called
 
 
-def test_auth_okta_step5_negative():
+@pytest.mark.parametrize("disable_saml_url_check", [True, False])
+def test_auth_okta_step5_negative(disable_saml_url_check):
     """Authentication by OKTA step5 negative test case."""
     authenticator = "https://testsso.snowflake.net/"
     application = "testapplication"
@@ -259,7 +260,9 @@ def test_auth_okta_step5_negative():
 
     ref_sso_url = "https://testsso.snowflake.net/sso"
     ref_token_url = "https://testsso.snowflake.net/token"
-    rest = _init_rest(ref_sso_url, ref_token_url)
+    rest = _init_rest(
+        ref_sso_url, ref_token_url, disable_saml_url_check=disable_saml_url_check
+    )
 
     auth = AuthByOkta(application)
     # step 1
@@ -306,10 +309,12 @@ def get_one_time_token():
     rest._host = f"{account}.snowflakecomputing.com"
     rest._port = 443
     auth._step5(rest._connection, ref_response_html)
-    assert rest._connection.errorhandler.called  # error
+    assert disable_saml_url_check ^ rest._connection.errorhandler.called  # error
 
 
-def _init_rest(ref_sso_url, ref_token_url, success=True, message=None):
+def _init_rest(
+    ref_sso_url, ref_token_url, success=True, message=None, disable_saml_url_check=False
+):
     def post_request(url, headers, body, **kwargs):
         _ = url
         _ = headers
@@ -324,7 +329,7 @@ def post_request(url, headers, body, **kwargs):
             },
         }
 
-    connection = mock_connection()
+    connection = mock_connection(disable_saml_url_check=disable_saml_url_check)
     connection.errorhandler = Mock(return_value=None)
     connection._ocsp_mode = Mock(return_value=OCSPMode.FAIL_OPEN)
     type(connection).application = PropertyMock(return_value=CLIENT_NAME)

test/unit/test_connection.py

@@ -21,6 +21,7 @@
 from cryptography.hazmat.primitives.asymmetric import rsa
 
 import snowflake.connector
+from snowflake.connector.connection import DEFAULT_CONFIGURATION
 from snowflake.connector.errors import (
     Error,
     ForbiddenError,
@@ -513,3 +514,23 @@ def test_expired_detection():
             with pytest.raises(ProgrammingError):
                 cur.execute("select 1;")
     assert conn.expired
+
+
+@pytest.mark.skipolddriver
+def test_disable_saml_url_check_config():
+    with mock.patch(
+        "snowflake.connector.network.SnowflakeRestful._post_request",
+        return_value={
+            "data": {
+                "serverVersion": "a.b.c",
+            },
+            "code": None,
+            "message": None,
+            "success": True,
+        },
+    ):
+        conn = fake_connector()
+        assert (
+            conn._disable_saml_url_check
+            == DEFAULT_CONFIGURATION.get("disable_saml_url_check")[0]
+        )