[async] Add WIF impersonation path length as data sent to Snowflake backend (#2521)

Author: sfc-gh-pczajka

test/helpers.py

@@ -340,3 +340,9 @@ def apply_auth_class_update_body(auth_class, req_body_before):
     req_body_after = copy.deepcopy(req_body_before)
     auth_class.update_body(req_body_after)
     return req_body_after
+
+
+async def apply_auth_class_update_body_async(auth_class, req_body_before):
+    req_body_after = copy.deepcopy(req_body_before)
+    await auth_class.update_body(req_body_after)
+    return req_body_after

test/unit/aio/test_auth_async.py

@@ -8,6 +8,7 @@
 import asyncio
 import inspect
 import sys
+from test.helpers import apply_auth_class_update_body_async, create_mock_auth_body
 from test.unit.aio.mock_utils import mock_connection
 from unittest.mock import Mock, PropertyMock
 
@@ -340,3 +341,21 @@ def test_mro():
     assert AuthByDefault.mro().index(AuthByPluginAsync) < AuthByDefault.mro().index(
         AuthByPluginSync
     )
+
+
+async def test_auth_by_default_prepare_body_does_not_overwrite_client_environment_fields():
+    password = "testpassword"
+    auth_class = AuthByDefault(password)
+
+    req_body_before = create_mock_auth_body()
+    req_body_after = await apply_auth_class_update_body_async(
+        auth_class, req_body_before
+    )
+
+    assert all(
+        [
+            req_body_before["data"]["CLIENT_ENVIRONMENT"][k]
+            == req_body_after["data"]["CLIENT_ENVIRONMENT"][k]
+            for k in req_body_before["data"]["CLIENT_ENVIRONMENT"]
+        ]
+    )

test/unit/aio/test_auth_keypair_async.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+from test.helpers import apply_auth_class_update_body_async, create_mock_auth_body
 from test.unit.aio.mock_utils import mock_connection
 from unittest.mock import Mock, PropertyMock, patch
 
@@ -61,6 +62,24 @@ async def test_auth_keypair(authenticator):
     assert rest.master_token == "MASTER_TOKEN"
 
 
+async def test_auth_prepare_body_does_not_overwrite_client_environment_fields():
+    private_key_der, _ = generate_key_pair(2048)
+    auth_class = AuthByKeyPair(private_key=private_key_der)
+
+    req_body_before = create_mock_auth_body()
+    req_body_after = await apply_auth_class_update_body_async(
+        auth_class, req_body_before
+    )
+
+    assert all(
+        [
+            req_body_before["data"]["CLIENT_ENVIRONMENT"][k]
+            == req_body_after["data"]["CLIENT_ENVIRONMENT"][k]
+            for k in req_body_before["data"]["CLIENT_ENVIRONMENT"]
+        ]
+    )
+
+
 async def test_auth_keypair_abc():
     """Simple Key Pair test using abstraction layer."""
     private_key_der, public_key_der_encoded = generate_key_pair(2048)

test/unit/aio/test_auth_oauth_async.py

@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+from test.helpers import apply_auth_class_update_body_async, create_mock_auth_body
+
 import pytest
 
 from snowflake.connector.aio.auth import AuthByOAuth
@@ -20,6 +22,24 @@ async def test_auth_oauth():
     assert body["data"]["AUTHENTICATOR"] == "OAUTH", body
 
 
+async def test_auth_prepare_body_does_not_overwrite_client_environment_fields():
+    token = "oAuthToken"
+    auth_class = AuthByOAuth(token)
+
+    req_body_before = create_mock_auth_body()
+    req_body_after = await apply_auth_class_update_body_async(
+        auth_class, req_body_before
+    )
+
+    assert all(
+        [
+            req_body_before["data"]["CLIENT_ENVIRONMENT"][k]
+            == req_body_after["data"]["CLIENT_ENVIRONMENT"][k]
+            for k in req_body_before["data"]["CLIENT_ENVIRONMENT"]
+        ]
+    )
+
+
 @pytest.mark.parametrize("authenticator", ["oauth", "OAUTH"])
 async def test_oauth_authenticator_is_case_insensitive(monkeypatch, authenticator):
     """Test that oauth authenticator is case insensitive."""

test/unit/aio/test_auth_oauth_auth_code_async.py

@@ -4,6 +4,7 @@
 #
 
 import unittest.mock as mock
+from test.helpers import apply_auth_class_update_body_async, create_mock_auth_body
 from unittest.mock import patch
 
 import pytest
@@ -44,6 +45,34 @@ async def test_auth_oauth_auth_code_oauth_type(omit_oauth_urls_check):
     )
 
 
+async def test_auth_prepare_body_does_not_overwrite_client_environment_fields(
+    omit_oauth_urls_check,
+):
+    auth_class = AuthByOauthCode(
+        "app",
+        "clientId",
+        "clientSecret",
+        "auth_url",
+        "tokenRequestUrl",
+        "redirectUri:{port}",
+        "scope",
+        "host",
+    )
+
+    req_body_before = create_mock_auth_body()
+    req_body_after = await apply_auth_class_update_body_async(
+        auth_class, req_body_before
+    )
+
+    assert all(
+        [
+            req_body_before["data"]["CLIENT_ENVIRONMENT"][k]
+            == req_body_after["data"]["CLIENT_ENVIRONMENT"][k]
+            for k in req_body_before["data"]["CLIENT_ENVIRONMENT"]
+        ]
+    )
+
+
 @pytest.mark.parametrize("rtr_enabled", [True, False])
 async def test_auth_oauth_auth_code_single_use_refresh_tokens(
     rtr_enabled: bool, omit_oauth_urls_check

test/unit/aio/test_auth_oauth_credentials_async.py

@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+from test.helpers import apply_auth_class_update_body_async, create_mock_auth_body
+
 import pytest
 
 from snowflake.connector.aio.auth import AuthByOauthCredentials
@@ -27,6 +29,29 @@ async def test_auth_oauth_credentials_oauth_type():
     )
 
 
+async def test_auth_prepare_body_does_not_overwrite_client_environment_fields():
+    auth_class = AuthByOauthCredentials(
+        "app",
+        "clientId",
+        "clientSecret",
+        "https://example.com/oauth/token",
+        "scope",
+    )
+
+    req_body_before = create_mock_auth_body()
+    req_body_after = await apply_auth_class_update_body_async(
+        auth_class, req_body_before
+    )
+
+    assert all(
+        [
+            req_body_before["data"]["CLIENT_ENVIRONMENT"][k]
+            == req_body_after["data"]["CLIENT_ENVIRONMENT"][k]
+            for k in req_body_before["data"]["CLIENT_ENVIRONMENT"]
+        ]
+    )
+
+
 @pytest.mark.parametrize(
     "authenticator, oauth_credentials_in_body",
     [

test/unit/aio/test_auth_okta_async.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import logging
+from test.helpers import apply_auth_class_update_body_async, create_mock_auth_body
 from test.unit.aio.mock_utils import mock_connection
 from unittest.mock import MagicMock, Mock, PropertyMock, patch
 
@@ -18,6 +19,24 @@
 from snowflake.connector.description import CLIENT_NAME, CLIENT_VERSION
 
 
+async def test_auth_prepare_body_does_not_overwrite_client_environment_fields():
+    application = "testapplication"
+    auth_class = AuthByOkta(application)
+
+    req_body_before = create_mock_auth_body()
+    req_body_after = await apply_auth_class_update_body_async(
+        auth_class, req_body_before
+    )
+
+    assert all(
+        [
+            req_body_before["data"]["CLIENT_ENVIRONMENT"][k]
+            == req_body_after["data"]["CLIENT_ENVIRONMENT"][k]
+            for k in req_body_before["data"]["CLIENT_ENVIRONMENT"]
+        ]
+    )
+
+
 async def test_auth_okta():
     """Authentication by OKTA positive test case."""
     authenticator = "https://testsso.snowflake.net/"

test/unit/aio/test_auth_pat_async.py

@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+from test.helpers import apply_auth_class_update_body_async, create_mock_auth_body
+
 import pytest
 
 from snowflake.connector.aio.auth import AuthByPAT
@@ -27,6 +29,24 @@ async def test_auth_pat():
     assert auth.assertion_content is None
 
 
+async def test_pat_prepare_body_does_not_overwrite_client_environment_fields():
+    token = "patToken"
+    auth_class = AuthByPAT(token)
+
+    req_body_before = create_mock_auth_body()
+    req_body_after = await apply_auth_class_update_body_async(
+        auth_class, req_body_before
+    )
+
+    assert all(
+        [
+            req_body_before["data"]["CLIENT_ENVIRONMENT"][k]
+            == req_body_after["data"]["CLIENT_ENVIRONMENT"][k]
+            for k in req_body_before["data"]["CLIENT_ENVIRONMENT"]
+        ]
+    )
+
+
 async def test_auth_pat_reauthenticate():
     """Test PAT reauthenticate."""
     token = "patToken"

test/unit/aio/test_auth_webbrowser_async.py

@@ -8,6 +8,7 @@
 import asyncio
 import base64
 import socket
+from test.helpers import apply_auth_class_update_body_async, create_mock_auth_body
 from test.unit.aio.mock_utils import mock_connection
 from unittest import mock
 from unittest.mock import MagicMock, Mock, PropertyMock, patch
@@ -1008,6 +1009,22 @@ async def mock_webbrowser_auth_prepare(
     await conn.close()
 
 
+async def test_auth_prepare_body_does_not_overwrite_client_environment_fields():
+    auth_class = AuthByWebBrowser(application=APPLICATION)
+    req_body_before = create_mock_auth_body()
+    req_body_after = await apply_auth_class_update_body_async(
+        auth_class, req_body_before
+    )
+
+    assert all(
+        [
+            req_body_before["data"]["CLIENT_ENVIRONMENT"][k]
+            == req_body_after["data"]["CLIENT_ENVIRONMENT"][k]
+            for k in req_body_before["data"]["CLIENT_ENVIRONMENT"]
+        ]
+    )
+
+
 def test_mro():
     """Ensure that methods from AuthByPluginAsync override those from AuthByPlugin."""
     from snowflake.connector.aio.auth import AuthByPlugin as AuthByPluginAsync

test/unit/aio/test_auth_workload_identity_async.py

@@ -14,11 +14,15 @@
 import jwt
 import pytest
 
-from snowflake.connector.aio._wif_util import AttestationProvider
+from snowflake.connector.aio._wif_util import (
+    AttestationProvider,
+    WorkloadIdentityAttestation,
+)
 from snowflake.connector.aio.auth import AuthByWorkloadIdentity
 from snowflake.connector.errors import ProgrammingError
 
 from ...csp_helpers import gen_dummy_access_token, gen_dummy_id_token
+from ...helpers import apply_auth_class_update_body_async, create_mock_auth_body
 from .csp_helpers_async import FakeAwsEnvironmentAsync, FakeGceMetadataServiceAsync
 
 logger = logging.getLogger(__name__)
@@ -137,6 +141,42 @@ async def mock_post(*args, **kwargs):
     await connection.close()
 
 
+@pytest.mark.parametrize(
+    "provider,additional_args",
+    [
+        (AttestationProvider.AWS, {}),
+        (AttestationProvider.GCP, {}),
+        (AttestationProvider.AZURE, {}),
+        (
+            AttestationProvider.OIDC,
+            {"token": gen_dummy_id_token(sub="service-1", iss="issuer-1")},
+        ),
+    ],
+)
+async def test_auth_prepare_body_does_not_overwrite_client_environment_fields(
+    provider, additional_args
+):
+    auth_class = AuthByWorkloadIdentity(provider=provider, **additional_args)
+    auth_class.attestation = WorkloadIdentityAttestation(
+        provider=AttestationProvider.GCP,
+        credential=None,
+        user_identifier_components=None,
+    )
+
+    req_body_before = create_mock_auth_body()
+    req_body_after = await apply_auth_class_update_body_async(
+        auth_class, req_body_before
+    )
+
+    assert all(
+        [
+            req_body_before["data"]["CLIENT_ENVIRONMENT"][k]
+            == req_body_after["data"]["CLIENT_ENVIRONMENT"][k]
+            for k in req_body_before["data"]["CLIENT_ENVIRONMENT"]
+        ]
+    )
+
+
 # -- OIDC Tests --
 
 
@@ -151,6 +191,7 @@ async def test_explicit_oidc_valid_inline_token_plumbed_to_api():
         "AUTHENTICATOR": "WORKLOAD_IDENTITY",
         "PROVIDER": "OIDC",
         "TOKEN": dummy_token,
+        "CLIENT_ENVIRONMENT": {"WORKLOAD_IDENTITY_IMPERSONATION_PATH_LENGTH": 0},
     }
 
 
@@ -208,6 +249,9 @@ async def test_explicit_aws_encodes_audience_host_signature_to_api(
     data = await extract_api_data(auth_class)
     assert data["AUTHENTICATOR"] == "WORKLOAD_IDENTITY"
     assert data["PROVIDER"] == "AWS"
+    assert (
+        data["CLIENT_ENVIRONMENT"]["WORKLOAD_IDENTITY_IMPERSONATION_PATH_LENGTH"] == 0
+    )
     verify_aws_token(data["TOKEN"], fake_aws_environment.region)
 
 
@@ -309,6 +353,7 @@ async def test_explicit_gcp_plumbs_token_to_api(
         "AUTHENTICATOR": "WORKLOAD_IDENTITY",
         "PROVIDER": "GCP",
         "TOKEN": fake_gce_metadata_service.token,
+        "CLIENT_ENVIRONMENT": {"WORKLOAD_IDENTITY_IMPERSONATION_PATH_LENGTH": 0},
     }
 
 
@@ -365,6 +410,7 @@ def __init__(self, content):
         "AUTHENTICATOR": "WORKLOAD_IDENTITY",
         "PROVIDER": "GCP",
         "TOKEN": sa3_id_token,
+        "CLIENT_ENVIRONMENT": {"WORKLOAD_IDENTITY_IMPERSONATION_PATH_LENGTH": 2},
     }
 
 
@@ -419,6 +465,7 @@ async def test_explicit_azure_plumbs_token_to_api(fake_azure_metadata_service):
         "AUTHENTICATOR": "WORKLOAD_IDENTITY",
         "PROVIDER": "AZURE",
         "TOKEN": fake_azure_metadata_service.token,
+        "CLIENT_ENVIRONMENT": {"WORKLOAD_IDENTITY_IMPERSONATION_PATH_LENGTH": 0},
     }