Fix Key Vault challenge-based authentication tenant ID parsing for DSTSv2 authorities

.vscode/cspell.json

@@ -1111,6 +1111,8 @@
     {
       "filename": "sdk/keyvault/**",
       "words": [
+        "DSTS",
+        "dstsv",
         "eddsa",
         "Thawte"
       ]

sdk/keyvault/azure-keyvault-administration/CHANGELOG.md

@@ -8,6 +8,7 @@
 
 ### Bugs Fixed
 
+- Fixed challenge-based authentication to correctly parse the tenant ID from DSTSv2 authority URIs ([#45326](https://github.com/Azure/azure-sdk-for-python/issues/45326)).
 - Fixed a bug in the challenge authentication policy where the authentication challenge was cached before the challenge resource was verified. The challenge is now cached only after resource verification succeeds [#48710](https://github.com/Azure/azure-sdk-for-python/pull/48710).
 
 ### Other Changes

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/http_challenge.py

@@ -6,6 +6,8 @@
 from typing import Dict, MutableMapping, Optional
 from urllib import parse
 
+_DSTS_V2_PATH_SEGMENT = "dstsv2"
+
 
 class HttpChallenge(object):
     """An object representing the content of a Key Vault authentication challenge.
@@ -66,18 +68,35 @@ def __init__(
         if "authorization" not in self._parameters and "authorization_uri" not in self._parameters:
             raise ValueError("Invalid challenge parameters")
 
-        authorization_uri = self.get_authorization_server()
-        # the authorization server URI should look something like https://login.windows.net/tenant-id
-        raw_uri_path = str(parse.urlparse(authorization_uri).path)
-        uri_path = raw_uri_path.lstrip("/")
-        self.tenant_id = uri_path.split("/", maxsplit=1)[0] or None
+        self.tenant_id = self._parse_tenant_id(self.get_authorization_server())
 
         # if the response headers were supplied
         if response_headers:
             # get the message signing key and message key encryption key from the headers
             self.server_signature_key = response_headers.get("x-ms-message-signing-key", None)
             self.server_encryption_key = response_headers.get("x-ms-message-encryption-key", None)
 
+    @staticmethod
+    def _parse_tenant_id(authorization_uri: str) -> "Optional[str]":
+        """Extracts the tenant ID from the authorization server URI of a challenge.
+
+        For Microsoft Entra ID authorities the tenant ID is the first path segment, for example
+        https://login.microsoftonline.com/<tenant-id>. For DSTSv2 authorities the first path segment is the literal
+        "dstsv2" and the tenant ID is the second path segment, for example
+        https://uswest2-passive-dsts.dsts.core.windows.net/dstsv2/<tenant-id>.
+
+        :param str authorization_uri: The authorization server URI from the challenge.
+
+        :returns: The tenant ID, or None if the URI does not contain one.
+        :rtype: str or None
+        """
+        raw_uri_path = str(parse.urlparse(authorization_uri).path)
+        path_segments = raw_uri_path.lstrip("/").split("/")
+        tenant_id = path_segments[0]
+        if tenant_id.lower() == _DSTS_V2_PATH_SEGMENT and len(path_segments) > 1 and path_segments[1]:
+            tenant_id = path_segments[1]
+        return tenant_id or None
+
     def is_bearer_challenge(self) -> bool:
         """Tests whether the HttpChallenge is a Bearer challenge.
 

sdk/keyvault/azure-keyvault-administration/tests/test_challenge_auth.py

@@ -15,7 +15,7 @@
 from azure.core.credentials import AccessToken, AccessTokenInfo
 from azure.core.pipeline import AsyncPipeline, Pipeline
 from azure.core.rest import HttpRequest
-from azure.keyvault.administration._internal import ChallengeAuthPolicy, HttpChallengeCache
+from azure.keyvault.administration._internal import ChallengeAuthPolicy, HttpChallenge, HttpChallengeCache
 from azure.keyvault.administration._internal.async_challenge_auth_policy import AsyncChallengeAuthPolicy
 
 TOKEN_TYPES = [AccessToken, AccessTokenInfo]
@@ -253,3 +253,134 @@ async def get_token(*_, **__):
     await pipeline.run(first_request)
 
     await pipeline.run(HttpRequest("GET", second_url))
+
+
+ENTRA_TENANT_ID = "72f988bf-86f1-41af-91ab-2d7cd022db57"
+DSTS_TENANT_ID = "de763a21-49f7-4b08-a8e1-52c8fbc103b4"
+DSTS_AUTHORITY = "https://uswest2-passive-dsts.dsts.core.windows.net"
+
+
+@pytest.mark.parametrize(
+    "authority,expected_tenant",
+    [
+        (f"https://login.microsoftonline.com/{ENTRA_TENANT_ID}", ENTRA_TENANT_ID),
+        (f"https://login.microsoftonline.com/{ENTRA_TENANT_ID}/oauth2/authorize", ENTRA_TENANT_ID),
+        (f"{DSTS_AUTHORITY}/dstsv2/{DSTS_TENANT_ID}", DSTS_TENANT_ID),
+        (f"{DSTS_AUTHORITY}/DSTSv2/{DSTS_TENANT_ID}/", DSTS_TENANT_ID),
+        # a DSTSv2 authority without a tenant segment keeps the previous behavior
+        (f"{DSTS_AUTHORITY}/dstsv2", "dstsv2"),
+        (f"{DSTS_AUTHORITY}/dstsv2/", "dstsv2"),
+        # an empty segment after "dstsv2" is not used as the tenant ID
+        (f"{DSTS_AUTHORITY}/dstsv2//{DSTS_TENANT_ID}", "dstsv2"),
+        # path segments after the DSTSv2 tenant ID are ignored
+        (f"{DSTS_AUTHORITY}/dstsv2/{DSTS_TENANT_ID}/oauth2/token", DSTS_TENANT_ID),
+        # only an exact "dstsv2" first segment denotes a DSTSv2 authority
+        (f"{DSTS_AUTHORITY}/dstsv2x/{DSTS_TENANT_ID}", "dstsv2x"),
+        (f"https://login.microsoftonline.com/{ENTRA_TENANT_ID}/dstsv2/{DSTS_TENANT_ID}", ENTRA_TENANT_ID),
+        ("https://login.microsoftonline.com/", None),
+    ],
+)
+def test_challenge_parsing_tenant_id(authority, expected_tenant):
+    """The tenant ID should be parsed from both Microsoft Entra ID and DSTSv2 authorization URIs"""
+
+    challenge = HttpChallenge(
+        "https://request.uri", challenge=f'Bearer authorization="{authority}", resource=https://vault.azure.net'
+    )
+
+    assert challenge.get_authorization_server() == authority
+    assert challenge.tenant_id == expected_tenant
+
+
+@empty_challenge_cache
+@pytest.mark.parametrize("token_type", TOKEN_TYPES)
+def test_tenant_dstsv2(token_type):
+    """The policy's token requests should pass the tenant ID that follows the "dstsv2" segment of the authority"""
+
+    expected_token = "expected_token"
+    authority = f"{DSTS_AUTHORITY}/dstsv2/{DSTS_TENANT_ID}"
+    challenge = Mock(
+        status_code=401,
+        headers={"WWW-Authenticate": f'Bearer authorization="{authority}", resource=https://vault.azure.net'},
+    )
+
+    class Requests:
+        count = 0
+
+    def send(request):
+        Requests.count += 1
+        if Requests.count == 1:
+            # first request should be unauthorized
+            assert "Authorization" not in request.headers
+            return challenge
+        elif Requests.count == 2:
+            # second request should be authorized according to the challenge
+            assert expected_token in request.headers["Authorization"]
+            return Mock(status_code=200)
+        raise ValueError("unexpected request")
+
+    def get_token(*_, options=None, **kwargs):
+        options_bag = options if options else kwargs
+        assert options_bag.get("tenant_id") == DSTS_TENANT_ID
+        return token_type(expected_token, 0)
+
+    if token_type == AccessToken:
+        credential = Mock(spec_set=["get_token"], get_token=Mock(wraps=get_token))
+    else:
+        credential = Mock(spec_set=["get_token_info"], get_token_info=Mock(wraps=get_token))
+
+    pipeline = Pipeline(policies=[ChallengeAuthPolicy(credential=credential)], transport=Mock(send=send))
+    pipeline.run(HttpRequest("GET", get_random_url()))
+
+    assert Requests.count == 2
+    if hasattr(credential, "get_token"):
+        assert credential.get_token.call_count == 1
+    else:
+        assert credential.get_token_info.call_count == 1
+
+
+@pytest.mark.asyncio
+@async_empty_challenge_cache
+@pytest.mark.parametrize("token_type", TOKEN_TYPES)
+async def test_tenant_dstsv2_async(token_type):
+    """The policy's token requests should pass the tenant ID that follows the "dstsv2" segment of the authority"""
+
+    expected_token = "expected_token"
+    authority = f"{DSTS_AUTHORITY}/dstsv2/{DSTS_TENANT_ID}"
+    challenge = Mock(
+        status_code=401,
+        headers={"WWW-Authenticate": f'Bearer authorization="{authority}", resource=https://vault.azure.net'},
+    )
+
+    class Requests:
+        count = 0
+
+    async def send(request):
+        Requests.count += 1
+        if Requests.count == 1:
+            # first request should be unauthorized
+            assert "Authorization" not in request.headers
+            return challenge
+        elif Requests.count == 2:
+            # second request should be authorized according to the challenge
+            assert expected_token in request.headers["Authorization"]
+            return Mock(status_code=200)
+        raise ValueError("unexpected request")
+
+    async def get_token(*_, options=None, **kwargs):
+        options_bag = options if options else kwargs
+        assert options_bag.get("tenant_id") == DSTS_TENANT_ID
+        return token_type(expected_token, 0)
+
+    if token_type == AccessToken:
+        credential = Mock(spec_set=["get_token"], get_token=Mock(wraps=get_token))
+    else:
+        credential = Mock(spec_set=["get_token_info"], get_token_info=Mock(wraps=get_token))
+
+    pipeline = AsyncPipeline(policies=[AsyncChallengeAuthPolicy(credential=credential)], transport=Mock(send=send))
+    await pipeline.run(HttpRequest("GET", get_random_url()))
+
+    assert Requests.count == 2
+    if hasattr(credential, "get_token"):
+        assert credential.get_token.call_count == 1
+    else:
+        assert credential.get_token_info.call_count == 1

sdk/keyvault/azure-keyvault-certificates/CHANGELOG.md

@@ -8,6 +8,7 @@
 
 ### Bugs Fixed
 
+- Fixed challenge-based authentication to correctly parse the tenant ID from DSTSv2 authority URIs ([#45326](https://github.com/Azure/azure-sdk-for-python/issues/45326)).
 - Fixed a bug in the challenge authentication policy where the authentication challenge was cached before the challenge resource was verified. The challenge is now cached only after resource verification succeeds [#48710](https://github.com/Azure/azure-sdk-for-python/pull/48710).
 
 ### Other Changes

sdk/keyvault/azure-keyvault-certificates/azure/keyvault/certificates/_shared/http_challenge.py

@@ -6,6 +6,8 @@
 from typing import Dict, MutableMapping, Optional
 from urllib import parse
 
+_DSTS_V2_PATH_SEGMENT = "dstsv2"
+
 
 class HttpChallenge(object):
     """An object representing the content of a Key Vault authentication challenge.
@@ -66,18 +68,35 @@ def __init__(
         if "authorization" not in self._parameters and "authorization_uri" not in self._parameters:
             raise ValueError("Invalid challenge parameters")
 
-        authorization_uri = self.get_authorization_server()
-        # the authorization server URI should look something like https://login.windows.net/tenant-id
-        raw_uri_path = str(parse.urlparse(authorization_uri).path)
-        uri_path = raw_uri_path.lstrip("/")
-        self.tenant_id = uri_path.split("/", maxsplit=1)[0] or None
+        self.tenant_id = self._parse_tenant_id(self.get_authorization_server())
 
         # if the response headers were supplied
         if response_headers:
             # get the message signing key and message key encryption key from the headers
             self.server_signature_key = response_headers.get("x-ms-message-signing-key", None)
             self.server_encryption_key = response_headers.get("x-ms-message-encryption-key", None)
 
+    @staticmethod
+    def _parse_tenant_id(authorization_uri: str) -> "Optional[str]":
+        """Extracts the tenant ID from the authorization server URI of a challenge.
+
+        For Microsoft Entra ID authorities the tenant ID is the first path segment, for example
+        https://login.microsoftonline.com/<tenant-id>. For DSTSv2 authorities the first path segment is the literal
+        "dstsv2" and the tenant ID is the second path segment, for example
+        https://uswest2-passive-dsts.dsts.core.windows.net/dstsv2/<tenant-id>.
+
+        :param str authorization_uri: The authorization server URI from the challenge.
+
+        :returns: The tenant ID, or None if the URI does not contain one.
+        :rtype: str or None
+        """
+        raw_uri_path = str(parse.urlparse(authorization_uri).path)
+        path_segments = raw_uri_path.lstrip("/").split("/")
+        tenant_id = path_segments[0]
+        if tenant_id.lower() == _DSTS_V2_PATH_SEGMENT and len(path_segments) > 1 and path_segments[1]:
+            tenant_id = path_segments[1]
+        return tenant_id or None
+
     def is_bearer_challenge(self) -> bool:
         """Tests whether the HttpChallenge is a Bearer challenge.
 

sdk/keyvault/azure-keyvault-certificates/tests/test_challenge_auth.py

@@ -16,7 +16,7 @@
 from azure.core.credentials import AccessToken, AccessTokenInfo
 from azure.core.pipeline import Pipeline
 from azure.core.rest import HttpRequest
-from azure.keyvault.certificates._shared import ChallengeAuthPolicy, HttpChallengeCache
+from azure.keyvault.certificates._shared import ChallengeAuthPolicy, HttpChallenge, HttpChallengeCache
 
 TOKEN_TYPES = [AccessToken, AccessTokenInfo]
 
@@ -136,3 +136,86 @@ def get_token(*_, **__):
     pipeline.run(first_request)
 
     pipeline.run(HttpRequest("GET", second_url))
+
+
+ENTRA_TENANT_ID = "72f988bf-86f1-41af-91ab-2d7cd022db57"
+DSTS_TENANT_ID = "de763a21-49f7-4b08-a8e1-52c8fbc103b4"
+DSTS_AUTHORITY = "https://uswest2-passive-dsts.dsts.core.windows.net"
+
+
+@pytest.mark.parametrize(
+    "authority,expected_tenant",
+    [
+        (f"https://login.microsoftonline.com/{ENTRA_TENANT_ID}", ENTRA_TENANT_ID),
+        (f"https://login.microsoftonline.com/{ENTRA_TENANT_ID}/oauth2/authorize", ENTRA_TENANT_ID),
+        (f"{DSTS_AUTHORITY}/dstsv2/{DSTS_TENANT_ID}", DSTS_TENANT_ID),
+        (f"{DSTS_AUTHORITY}/DSTSv2/{DSTS_TENANT_ID}/", DSTS_TENANT_ID),
+        # a DSTSv2 authority without a tenant segment keeps the previous behavior
+        (f"{DSTS_AUTHORITY}/dstsv2", "dstsv2"),
+        (f"{DSTS_AUTHORITY}/dstsv2/", "dstsv2"),
+        # an empty segment after "dstsv2" is not used as the tenant ID
+        (f"{DSTS_AUTHORITY}/dstsv2//{DSTS_TENANT_ID}", "dstsv2"),
+        # path segments after the DSTSv2 tenant ID are ignored
+        (f"{DSTS_AUTHORITY}/dstsv2/{DSTS_TENANT_ID}/oauth2/token", DSTS_TENANT_ID),
+        # only an exact "dstsv2" first segment denotes a DSTSv2 authority
+        (f"{DSTS_AUTHORITY}/dstsv2x/{DSTS_TENANT_ID}", "dstsv2x"),
+        (f"https://login.microsoftonline.com/{ENTRA_TENANT_ID}/dstsv2/{DSTS_TENANT_ID}", ENTRA_TENANT_ID),
+        ("https://login.microsoftonline.com/", None),
+    ],
+)
+def test_challenge_parsing_tenant_id(authority, expected_tenant):
+    """The tenant ID should be parsed from both Microsoft Entra ID and DSTSv2 authorization URIs"""
+
+    challenge = HttpChallenge(
+        "https://request.uri", challenge=f'Bearer authorization="{authority}", resource=https://vault.azure.net'
+    )
+
+    assert challenge.get_authorization_server() == authority
+    assert challenge.tenant_id == expected_tenant
+
+
+@empty_challenge_cache
+@pytest.mark.parametrize("token_type", TOKEN_TYPES)
+def test_tenant_dstsv2(token_type):
+    """The policy's token requests should pass the tenant ID that follows the "dstsv2" segment of the authority"""
+
+    expected_token = "expected_token"
+    authority = f"{DSTS_AUTHORITY}/dstsv2/{DSTS_TENANT_ID}"
+    challenge = Mock(
+        status_code=401,
+        headers={"WWW-Authenticate": f'Bearer authorization="{authority}", resource=https://vault.azure.net'},
+    )
+
+    class Requests:
+        count = 0
+
+    def send(request):
+        Requests.count += 1
+        if Requests.count == 1:
+            # first request should be unauthorized
+            assert "Authorization" not in request.headers
+            return challenge
+        elif Requests.count == 2:
+            # second request should be authorized according to the challenge
+            assert expected_token in request.headers["Authorization"]
+            return Mock(status_code=200)
+        raise ValueError("unexpected request")
+
+    def get_token(*_, options=None, **kwargs):
+        options_bag = options if options else kwargs
+        assert options_bag.get("tenant_id") == DSTS_TENANT_ID
+        return token_type(expected_token, 0)
+
+    if token_type == AccessToken:
+        credential = Mock(spec_set=["get_token"], get_token=Mock(wraps=get_token))
+    else:
+        credential = Mock(spec_set=["get_token_info"], get_token_info=Mock(wraps=get_token))
+
+    pipeline = Pipeline(policies=[ChallengeAuthPolicy(credential=credential)], transport=Mock(send=send))
+    pipeline.run(HttpRequest("GET", get_random_url()))
+
+    assert Requests.count == 2
+    if hasattr(credential, "get_token"):
+        assert credential.get_token.call_count == 1
+    else:
+        assert credential.get_token_info.call_count == 1