SNOW-2111644 Support sovereign clouds for WIF (#2367)

Author: sfc-gh-xizhao

src/snowflake/connector/wif_util.py

@@ -22,6 +22,19 @@
 SNOWFLAKE_AUDIENCE = "snowflakecomputing.com"
 DEFAULT_ENTRA_SNOWFLAKE_RESOURCE = "api://fd3f753b-eed3-462c-b6a7-a4b5bb650aad"
 
+"""
+References:
+- https://learn.microsoft.com/en-us/entra/identity-platform/authentication-national-cloud#microsoft-entra-authentication-endpoints
+- https://learn.microsoft.com/en-us/answers/questions/1190472/what-are-the-token-issuers-for-the-sovereign-cloud
+"""
+AZURE_ISSUER_PREFIXES = [
+    "https://sts.windows.net/",  # Public and USGov (v1 issuer)
+    "https://sts.chinacloudapi.cn/",  # Mooncake (v1 issuer)
+    "https://login.microsoftonline.com/",  # Public (v2 issuer)
+    "https://login.microsoftonline.us/",  # USGov (v2 issuer)
+    "https://login.partner.microsoftonline.cn/",  # Mooncake (v2 issuer)
+]
+
 
 @unique
 class AttestationProvider(Enum):
@@ -108,6 +121,70 @@ def get_aws_arn() -> str | None:
     return caller_identity["Arn"]
 
 
+def get_aws_partition(arn: str) -> str | None:
+    """Get the current AWS partition from ARN, if any.
+
+    Args:
+        arn (str): The Amazon Resource Name (ARN) string.
+
+    Returns:
+        str | None: The AWS partition (e.g., 'aws', 'aws-cn', 'aws-us-gov')
+                    if found, otherwise None.
+
+    Reference: https://docs.aws.amazon.com/IAM/latest/UserGuide/reference-arns.html.
+    """
+    if not arn or not isinstance(arn, str):
+        return None
+    parts = arn.split(":")
+    if len(parts) > 1 and parts[0] == "arn" and parts[1]:
+        return parts[1]
+    logger.warning("Invalid AWS ARN: %s", arn)
+    return None
+
+
+def get_aws_sts_hostname(region: str, partition: str) -> str | None:
+    """Constructs the AWS STS hostname for a given region and partition.
+
+    Args:
+        region (str): The AWS region (e.g., 'us-east-1', 'cn-north-1').
+        partition (str): The AWS partition (e.g., 'aws', 'aws-cn', 'aws-us-gov').
+
+    Returns:
+        str | None: The AWS STS hostname (e.g., 'sts.us-east-1.amazonaws.com')
+                    if a valid hostname can be constructed, otherwise None.
+
+    References:
+    - https://docs.aws.amazon.com/sdkref/latest/guide/feature-sts-regionalized-endpoints.html
+    - https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_temp_region-endpoints.html
+    - https://docs.aws.amazon.com/general/latest/gr/sts.html
+    """
+    if (
+        not region
+        or not partition
+        or not isinstance(region, str)
+        or not isinstance(partition, str)
+    ):
+        return None
+
+    if partition == "aws":
+        # For the 'aws' partition, STS endpoints are generally regional
+        # except for the global endpoint (sts.amazonaws.com) which is
+        # generally resolved to us-east-1 under the hood by the SDKs
+        # when a region is not explicitly specified.
+        # However, for explicit regional endpoints, the format is sts.<region>.amazonaws.com
+        return f"sts.{region}.amazonaws.com"
+    elif partition == "aws-cn":
+        # China regions have a different domain suffix
+        return f"sts.{region}.amazonaws.com.cn"
+    elif partition == "aws-us-gov":
+        return (
+            f"sts.{region}.amazonaws.com"  # GovCloud uses .com, but dedicated regions
+        )
+    else:
+        logger.warning("Invalid AWS partition: %s", partition)
+        return None
+
+
 def create_aws_attestation() -> WorkloadIdentityAttestation | None:
     """Tries to create a workload identity attestation for AWS.
 
@@ -125,8 +202,12 @@ def create_aws_attestation() -> WorkloadIdentityAttestation | None:
     if not arn:
         logger.debug("No AWS caller identity was found.")
         return None
+    partition = get_aws_partition(arn)
+    if not partition:
+        logger.debug("No AWS partition was found.")
+        return None
 
-    sts_hostname = f"sts.{region}.amazonaws.com"
+    sts_hostname = get_aws_sts_hostname(region, partition)
     request = AWSRequest(
         method="POST",
         url=f"https://{sts_hostname}/?Action=GetCallerIdentity&Version=2011-06-15",
@@ -234,9 +315,8 @@ def create_azure_attestation(
     issuer, subject = extract_iss_and_sub_without_signature_verification(jwt_str)
     if not issuer or not subject:
         return None
-    if not (
-        issuer.startswith("https://sts.windows.net/")
-        or issuer.startswith("https://login.microsoftonline.com/")
+    if not any(
+        issuer.startswith(issuer_prefix) for issuer_prefix in AZURE_ISSUER_PREFIXES
     ):
         # This might happen if we're running on a different platform that responds to the same metadata request signature as Azure.
         logger.debug("Unexpected Azure token issuer '%s'", issuer)

test/unit/test_auth_workload_identity.py

@@ -14,7 +14,12 @@
     HTTPError,
     Timeout,
 )
-from snowflake.connector.wif_util import AttestationProvider
+from snowflake.connector.wif_util import (
+    AZURE_ISSUER_PREFIXES,
+    AttestationProvider,
+    get_aws_partition,
+    get_aws_sts_hostname,
+)
 
 from ..csp_helpers import FakeAwsEnvironment, FakeGceMetadataService, gen_dummy_id_token
 
@@ -154,6 +159,73 @@ def test_explicit_aws_generates_unique_assertion_content(
     )
 
 
+@pytest.mark.parametrize(
+    "arn, expected_partition",
+    [
+        ("arn:aws:iam::123456789012:role/MyTestRole", "aws"),
+        (
+            "arn:aws-cn:ec2:cn-north-1:987654321098:instance/i-1234567890abcdef0",
+            "aws-cn",
+        ),
+        ("arn:aws-us-gov:s3:::my-gov-bucket", "aws-us-gov"),
+        ("arn:aws:s3:::my-bucket/my/key", "aws"),
+        ("arn:aws:lambda:us-east-1:123456789012:function:my-function", "aws"),
+        ("arn:aws:sns:eu-west-1:111122223333:my-topic", "aws"),
+        # Edge cases / Invalid inputs
+        ("invalid-arn", None),
+        ("arn::service:region:account:resource", None),  # Missing partition
+        ("arn:aws:iam:", "aws"),  # Incomplete ARN, but partition is present
+        ("", None),  # Empty string
+        (None, None),  # None input
+        (123, None),  # Non-string input
+    ],
+)
+def test_get_aws_partition_valid_and_invalid_arns(arn, expected_partition):
+    assert get_aws_partition(arn) == expected_partition
+
+
+@pytest.mark.parametrize(
+    "region, partition, expected_hostname",
+    [
+        # AWS partition
+        ("us-east-1", "aws", "sts.us-east-1.amazonaws.com"),
+        ("eu-west-2", "aws", "sts.eu-west-2.amazonaws.com"),
+        ("ap-southeast-1", "aws", "sts.ap-southeast-1.amazonaws.com"),
+        (
+            "us-east-1",
+            "aws",
+            "sts.us-east-1.amazonaws.com",
+        ),  # Redundant but good for coverage
+        # AWS China partition
+        ("cn-north-1", "aws-cn", "sts.cn-north-1.amazonaws.com.cn"),
+        ("cn-northwest-1", "aws-cn", "sts.cn-northwest-1.amazonaws.com.cn"),
+        ("", "aws-cn", None),  # No global endpoint for 'aws-cn' without region
+        # AWS GovCloud partition
+        ("us-gov-west-1", "aws-us-gov", "sts.us-gov-west-1.amazonaws.com"),
+        ("us-gov-east-1", "aws-us-gov", "sts.us-gov-east-1.amazonaws.com"),
+        ("", "aws-us-gov", None),  # No global endpoint for 'aws-us-gov' without region
+        # Invalid/Edge cases
+        ("us-east-1", "unknown-partition", None),  # Unknown partition
+        ("some-region", "invalid-partition", None),  # Invalid partition
+        (None, "aws", None),  # None region
+        ("us-east-1", None, None),  # None partition
+        (123, "aws", None),  # Non-string region
+        ("us-east-1", 456, None),  # Non-string partition
+        ("", "", None),  # Empty region and partition
+        ("us-east-1", "", None),  # Empty partition
+        (
+            "invalid-region",
+            "aws",
+            "sts.invalid-region.amazonaws.com",
+        ),  # Valid format, invalid region name
+    ],
+)
+def test_get_aws_sts_hostname_valid_and_invalid_inputs(
+    region, partition, expected_hostname
+):
+    assert get_aws_sts_hostname(region, partition) == expected_hostname
+
+
 # -- GCP Tests --
 
 
@@ -312,6 +384,22 @@ def test_explicit_azure_uses_explicit_entra_resource(fake_azure_metadata_service
     assert parsed["aud"] == "api://non-standard"
 
 
+@pytest.mark.parametrize(
+    "issuer",
+    [
+        "https://sts.windows.net/067802cd-8f92-4c7c-bceb-ea8f15d31cc5",
+        "https://sts.chinacloudapi.cn/067802cd-8f92-4c7c-bceb-ea8f15d31cc5",
+        "https://login.microsoftonline.com/067802cd-8f92-4c7c-bceb-ea8f15d31cc5/v2.0",
+        "https://login.microsoftonline.us/067802cd-8f92-4c7c-bceb-ea8f15d31cc5/v2.0",
+        "https://login.partner.microsoftonline.cn/067802cd-8f92-4c7c-bceb-ea8f15d31cc5/v2.0",
+    ],
+)
+def test_azure_issuer_prefixes(issuer):
+    assert any(
+        issuer.startswith(issuer_prefix) for issuer_prefix in AZURE_ISSUER_PREFIXES
+    )
+
+
 # -- Auto-detect Tests --