Support WIF Impersonation on GCP workloads (#2496)

Co-authored-by: Peter Mansour <[email protected]>

Author: 2 people

src/snowflake/connector/auth/workload_identity.py

@@ -55,12 +55,14 @@ def __init__(
         provider: AttestationProvider,
         token: str | None = None,
         entra_resource: str | None = None,
+        impersonation_path: list[str] | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
         self.provider = provider
         self.token = token
         self.entra_resource = entra_resource
+        self.impersonation_path = impersonation_path
 
         self.attestation: WorkloadIdentityAttestation | None = None
 
@@ -85,7 +87,10 @@ def prepare(
             self.provider,
             self.entra_resource,
             self.token,
-            session_manager=conn._session_manager.clone() if conn else None,
+            self.impersonation_path,
+            session_manager=(
+                conn._session_manager.clone(max_retries=0) if conn else None
+            ),
         )
 
     def reauthenticate(self, **kwargs: typing.Any) -> dict[str, bool]:

src/snowflake/connector/connection.py

@@ -214,6 +214,7 @@ def _get_private_bytes_from_file(
     "authenticator": (DEFAULT_AUTHENTICATOR, (type(None), str)),
     "workload_identity_provider": (None, (type(None), AttestationProvider)),
     "workload_identity_entra_resource": (None, (type(None), str)),
+    "workload_identity_impersonation_path": (None, (type(None), list[str])),
     "mfa_callback": (None, (type(None), Callable)),
     "password_callback": (None, (type(None), Callable)),
     "auth_class": (None, (type(None), AuthByPlugin)),
@@ -1355,10 +1356,24 @@ def __open_connection(self):
                             "errno": ER_INVALID_WIF_SETTINGS,
                         },
                     )
+                if (
+                    self._workload_identity_impersonation_path
+                    and self._workload_identity_provider != AttestationProvider.GCP
+                ):
+                    Error.errorhandler_wrapper(
+                        self,
+                        None,
+                        ProgrammingError,
+                        {
+                            "msg": "workload_identity_impersonation_path is currently only supported for GCP.",
+                            "errno": ER_INVALID_WIF_SETTINGS,
+                        },
+                    )
                 self.auth_class = AuthByWorkloadIdentity(
                     provider=self._workload_identity_provider,
                     token=self._token,
                     entra_resource=self._workload_identity_entra_resource,
+                    impersonation_path=self._workload_identity_impersonation_path,
                 )
             else:
                 # okta URL, e.g., https://<account>.okta.com/
@@ -1531,6 +1546,7 @@ def __config(self, **kwargs):
             workload_identity_dependent_options = [
                 "workload_identity_provider",
                 "workload_identity_entra_resource",
+                "workload_identity_impersonation_path",
             ]
             for dependent_option in workload_identity_dependent_options:
                 if (

src/snowflake/connector/wif_util.py

@@ -20,6 +20,9 @@
 logger = logging.getLogger(__name__)
 SNOWFLAKE_AUDIENCE = "snowflakecomputing.com"
 DEFAULT_ENTRA_SNOWFLAKE_RESOURCE = "api://fd3f753b-eed3-462c-b6a7-a4b5bb650aad"
+GCP_METADATA_SERVICE_ACCOUNT_BASE_URL = (
+    "http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default"
+)
 
 
 @unique
@@ -184,29 +187,103 @@ def create_aws_attestation(
     )
 
 
-def create_gcp_attestation(
-    session_manager: SessionManager | None = None,
-) -> WorkloadIdentityAttestation:
-    """Tries to create a workload identity attestation for GCP.
+def get_gcp_access_token(session_manager: SessionManager) -> str:
+    """Gets a GCP access token from the metadata server.
+
+    If the application isn't running on GCP or no credentials were found, raises an error.
+    """
+    try:
+        res = session_manager.request(
+            method="GET",
+            url=f"{GCP_METADATA_SERVICE_ACCOUNT_BASE_URL}/token",
+            headers={
+                "Metadata-Flavor": "Google",
+            },
+        )
+        res.raise_for_status()
+        return res.json()["access_token"]
+    except Exception as e:
+        raise ProgrammingError(
+            msg=f"Error fetching GCP access token: {e}. Ensure the application is running on GCP.",
+            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+        )
+
+
+def get_gcp_identity_token_via_impersonation(
+    impersonation_path: list[str], session_manager: SessionManager
+) -> str:
+    """Gets a GCP identity token from the metadata server.
+
+    If the application isn't running on GCP or no credentials were found, raises an error.
+    """
+    if not impersonation_path:
+        raise ProgrammingError(
+            msg="Error: impersonation_path cannot be empty.",
+            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+        )
+
+    current_sa_token = get_gcp_access_token(session_manager)
+    impersonation_path = [
+        f"projects/-/serviceAccounts/{client_id}" for client_id in impersonation_path
+    ]
+    try:
+        res = session_manager.post(
+            url=f"https://iamcredentials.googleapis.com/v1/{impersonation_path[-1]}:generateIdToken",
+            headers={
+                "Authorization": f"Bearer {current_sa_token}",
+                "Content-Type": "application/json",
+            },
+            json={
+                "delegates": impersonation_path[:-1],
+                "audience": SNOWFLAKE_AUDIENCE,
+            },
+        )
+        res.raise_for_status()
+        return res.json()["token"]
+    except Exception as e:
+        raise ProgrammingError(
+            msg=f"Error fetching GCP identity token for impersonated GCP service account '{impersonation_path[-1]}': {e}. Ensure the application is running on GCP.",
+            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+        )
+
+
+def get_gcp_identity_token(session_manager: SessionManager) -> str:
+    """Gets a GCP identity token from the metadata server.
 
     If the application isn't running on GCP or no credentials were found, raises an error.
     """
     try:
         res = session_manager.request(
             method="GET",
-            url=f"http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/identity?audience={SNOWFLAKE_AUDIENCE}",
+            url=f"{GCP_METADATA_SERVICE_ACCOUNT_BASE_URL}/identity?audience={SNOWFLAKE_AUDIENCE}",
             headers={
                 "Metadata-Flavor": "Google",
             },
         )
         res.raise_for_status()
+        return res.content.decode("utf-8")
     except Exception as e:
         raise ProgrammingError(
-            msg=f"Error fetching GCP metadata: {e}. Ensure the application is running on GCP.",
+            msg=f"Error fetching GCP identity token: {e}. Ensure the application is running on GCP.",
             errno=ER_WIF_CREDENTIALS_NOT_FOUND,
         )
 
-    jwt_str = res.content.decode("utf-8")
+
+def create_gcp_attestation(
+    session_manager: SessionManager,
+    impersonation_path: list[str] | None = None,
+) -> WorkloadIdentityAttestation:
+    """Tries to create a workload identity attestation for GCP.
+
+    If the application isn't running on GCP or no credentials were found, raises an error.
+    """
+    if impersonation_path:
+        jwt_str = get_gcp_identity_token_via_impersonation(
+            impersonation_path, session_manager
+        )
+    else:
+        jwt_str = get_gcp_identity_token(session_manager)
+
     _, subject = extract_iss_and_sub_without_signature_verification(jwt_str)
     return WorkloadIdentityAttestation(
         AttestationProvider.GCP, jwt_str, {"sub": subject}
@@ -295,6 +372,7 @@ def create_attestation(
     provider: AttestationProvider,
     entra_resource: str | None = None,
     token: str | None = None,
+    impersonation_path: list[str] | None = None,
     session_manager: SessionManager | None = None,
 ) -> WorkloadIdentityAttestation:
     """Entry point to create an attestation using the given provider.
@@ -311,7 +389,7 @@ def create_attestation(
     elif provider == AttestationProvider.AZURE:
         return create_azure_attestation(entra_resource, session_manager)
     elif provider == AttestationProvider.GCP:
-        return create_gcp_attestation(session_manager)
+        return create_gcp_attestation(session_manager, impersonation_path)
     elif provider == AttestationProvider.OIDC:
         return create_oidc_attestation(token)
     else:

test/csp_helpers.py

@@ -40,6 +40,13 @@ def gen_dummy_id_token(
     )
 
 
+def gen_dummy_access_token(sub="test-subject") -> str:
+    """Generates a dummy access token using the given subject."""
+    key = "secret"
+    logger.debug(f"Generating dummy access token for subject {sub}")
+    return (sub + key).encode("utf-8").hex()
+
+
 def build_response(content: bytes, status_code: int = 200, headers=None) -> Response:
     """Builds a requests.Response object with the given status code and content."""
     response = Response()
@@ -285,6 +292,19 @@ def handle_request(self, method, parsed_url, headers, timeout):
             audience = query_string["audience"][0]
             self.token = gen_dummy_id_token(sub=self.sub, iss=self.iss, aud=audience)
             return build_response(self.token.encode("utf-8"))
+        elif (
+            method == "GET"
+            and parsed_url.path
+            == "/computeMetadata/v1/instance/service-accounts/default/token"
+            and headers.get("Metadata-Flavor") == "Google"
+        ):
+            self.token = gen_dummy_access_token(sub=self.sub)
+            ret = {
+                "access_token": self.token,
+                "expires_in": 3599,
+                "token_type": "Bearer",
+            }
+            return build_response(json.dumps(ret).encode("utf-8"))
         else:
             # Reject malformed requests.
             raise HTTPError()

test/unit/test_auth_workload_identity.py

@@ -17,7 +17,13 @@
 )
 from snowflake.connector.wif_util import AttestationProvider, get_aws_sts_hostname
 
-from ..csp_helpers import FakeAwsEnvironment, FakeGceMetadataService, gen_dummy_id_token
+from ..csp_helpers import (
+    FakeAwsEnvironment,
+    FakeGceMetadataService,
+    build_response,
+    gen_dummy_access_token,
+    gen_dummy_id_token,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -289,7 +295,7 @@ def test_explicit_gcp_metadata_server_error_bubbles_up(exception):
         with pytest.raises(ProgrammingError) as excinfo:
             auth_class.prepare(conn=None)
 
-    assert "Error fetching GCP metadata:" in str(excinfo.value)
+    assert "Error fetching GCP identity token:" in str(excinfo.value)
     assert "Ensure the application is running on GCP." in str(excinfo.value)
 
 
@@ -317,6 +323,44 @@ def test_explicit_gcp_generates_unique_assertion_content(
     assert auth_class.assertion_content == '{"_provider":"GCP","sub":"123456"}'
 
 
+@mock.patch("snowflake.connector.session_manager.SessionManager.post")
+def test_gcp_calls_correct_apis_and_populates_auth_data_for_final_sa(
+    mock_post_request, fake_gce_metadata_service: FakeGceMetadataService
+):
+    fake_gce_metadata_service.sub = "sa1"
+    impersonation_path = ["sa2", "sa3"]
+    sa1_access_token = gen_dummy_access_token("sa1")
+    sa3_id_token = gen_dummy_id_token("sa3")
+
+    mock_post_request.return_value = build_response(
+        json.dumps({"token": sa3_id_token}).encode("utf-8")
+    )
+
+    auth_class = AuthByWorkloadIdentity(
+        provider=AttestationProvider.GCP, impersonation_path=impersonation_path
+    )
+    auth_class.prepare(conn=None)
+
+    mock_post_request.assert_called_once_with(
+        url="https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/sa3:generateIdToken",
+        headers={
+            "Authorization": f"Bearer {sa1_access_token}",
+            "Content-Type": "application/json",
+        },
+        json={
+            "delegates": ["projects/-/serviceAccounts/sa2"],
+            "audience": "snowflakecomputing.com",
+        },
+    )
+
+    assert auth_class.assertion_content == '{"_provider":"GCP","sub":"sa3"}'
+    assert extract_api_data(auth_class) == {
+        "AUTHENTICATOR": "WORKLOAD_IDENTITY",
+        "PROVIDER": "GCP",
+        "TOKEN": sa3_id_token,
+    }
+
+
 # -- Azure Tests --
 
 

test/unit/test_connection.py

@@ -631,6 +631,7 @@ def test_otel_error_message(caplog, mock_post_requests):
             "workload_identity_entra_resource",
             "api://0b2f151f-09a2-46eb-ad5a-39d5ebef917b",
         ),
+        ("workload_identity_impersonation_path", ["subject-b", "subject-c"]),
     ],
 )
 def test_cannot_set_dependent_params_without_wlid_authenticator(
@@ -679,6 +680,71 @@ def test_workload_identity_provider_is_required_for_wif_authenticator(
         assert expected_error_msg in str(excinfo.value)
 
 
+@pytest.mark.parametrize(
+    "provider_param",
+    [
+        # Strongly-typed values.
+        AttestationProvider.AWS,
+        AttestationProvider.AZURE,
+        AttestationProvider.OIDC,
+        # String values.
+        "AWS",
+        "AZURE",
+        "OIDC",
+    ],
+)
+def test_workload_identity_impersonation_path_unsupported_for_non_gcp_providers(
+    monkeypatch, provider_param
+):
+    with monkeypatch.context() as m:
+        m.setattr(
+            "snowflake.connector.SnowflakeConnection._authenticate", lambda *_: None
+        )
+
+        with pytest.raises(ProgrammingError) as excinfo:
+            snowflake.connector.connect(
+                account="account",
+                authenticator="WORKLOAD_IDENTITY",
+                workload_identity_provider=provider_param,
+                workload_identity_impersonation_path=[
+                    "[email protected]"
+                ],
+            )
+        assert (
+            "workload_identity_impersonation_path is currently only supported for GCP."
+            in str(excinfo.value)
+        )
+
+
+@pytest.mark.parametrize(
+    "provider_param",
+    [
+        AttestationProvider.GCP,
+        "GCP",
+    ],
+)
+def test_workload_identity_impersonation_path_supported_for_gcp_provider(
+    monkeypatch, provider_param
+):
+    with monkeypatch.context() as m:
+        m.setattr(
+            "snowflake.connector.SnowflakeConnection._authenticate", lambda *_: None
+        )
+
+        conn = snowflake.connector.connect(
+            account="account",
+            authenticator="WORKLOAD_IDENTITY",
+            workload_identity_provider=provider_param,
+            workload_identity_impersonation_path=[
+                "[email protected]"
+            ],
+        )
+        assert conn.auth_class.provider == AttestationProvider.GCP
+        assert conn.auth_class.impersonation_path == [
+            "[email protected]"
+        ]
+
+
 @pytest.mark.parametrize(
     "provider_param, parsed_provider",
     [

test/wif/test_wif.py

@@ -59,7 +59,6 @@ def test_should_authenticate_using_oidc():
 
 
 @pytest.mark.wif
-@pytest.mark.skip("Impersonation is still being developed")
 def test_should_authenticate_with_impersonation():
     if not isinstance(IMPERSONATION_PATH, str) or not IMPERSONATION_PATH:
         pytest.skip("Skipping test - IMPERSONATION_PATH is not set")