[Key Vault] Correctly handle token refreshes in AD FS (#33634)

Author: mccoyp

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

@@ -26,7 +26,6 @@
 from . import http_challenge_cache as ChallengeCache
 from .challenge_auth_policy import _enforce_tls, _update_challenge
 
-
 class AsyncChallengeAuthPolicy(AsyncBearerTokenCredentialPolicy):
     """Policy for handling HTTP authentication challenges.
 
@@ -49,7 +48,11 @@ async def on_request(self, request: PipelineRequest) -> None:
             if self._need_new_token():
                 # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
                 scope = challenge.get_scope() or challenge.get_resource() + "/.default"
-                self._token = await self._credential.get_token(scope, tenant_id=challenge.tenant_id)
+                # Exclude tenant for AD FS authentication
+                if challenge.tenant_id and challenge.tenant_id.lower().endswith("adfs"):
+                    self._token = await self._credential.get_token(scope)
+                else:
+                    self._token = await self._credential.get_token(scope, tenant_id=challenge.tenant_id)
 
             # ignore mypy's warning -- although self._token is Optional, get_token raises when it fails to get a token
             request.http_request.headers["Authorization"] = f"Bearer {self._token.token}"  # type: ignore

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

@@ -66,7 +66,7 @@ class ChallengeAuthPolicy(BearerTokenCredentialPolicy):
     def __init__(self, credential: TokenCredential, *scopes: str, **kwargs: Any) -> None:
         super(ChallengeAuthPolicy, self).__init__(credential, *scopes, **kwargs)
         self._credential = credential
-        self._token: "Optional[AccessToken]" = None
+        self._token: Optional[AccessToken] = None
         self._verify_challenge_resource = kwargs.pop("verify_challenge_resource", True)
 
     def on_request(self, request: PipelineRequest) -> None:
@@ -77,7 +77,11 @@ def on_request(self, request: PipelineRequest) -> None:
             if self._need_new_token:
                 # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
                 scope = challenge.get_scope() or challenge.get_resource() + "/.default"
-                self._token = self._credential.get_token(scope, tenant_id=challenge.tenant_id)
+                # Exclude tenant for AD FS authentication
+                if challenge.tenant_id and challenge.tenant_id.lower().endswith("adfs"):
+                    self._token = self._credential.get_token(scope)
+                else:
+                    self._token = self._credential.get_token(scope, tenant_id=challenge.tenant_id)
 
             # ignore mypy's warning -- although self._token is Optional, get_token raises when it fails to get a token
             request.http_request.headers["Authorization"] = f"Bearer {self._token.token}"  # type: ignore

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

@@ -48,7 +48,11 @@ async def on_request(self, request: PipelineRequest) -> None:
             if self._need_new_token():
                 # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
                 scope = challenge.get_scope() or challenge.get_resource() + "/.default"
-                self._token = await self._credential.get_token(scope, tenant_id=challenge.tenant_id)
+                # Exclude tenant for AD FS authentication
+                if challenge.tenant_id and challenge.tenant_id.lower().endswith("adfs"):
+                    self._token = await self._credential.get_token(scope)
+                else:
+                    self._token = await self._credential.get_token(scope, tenant_id=challenge.tenant_id)
 
             # ignore mypy's warning -- although self._token is Optional, get_token raises when it fails to get a token
             request.http_request.headers["Authorization"] = f"Bearer {self._token.token}"  # type: ignore

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

@@ -77,7 +77,11 @@ def on_request(self, request: PipelineRequest) -> None:
             if self._need_new_token:
                 # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
                 scope = challenge.get_scope() or challenge.get_resource() + "/.default"
-                self._token = self._credential.get_token(scope, tenant_id=challenge.tenant_id)
+                # Exclude tenant for AD FS authentication
+                if challenge.tenant_id and challenge.tenant_id.lower().endswith("adfs"):
+                    self._token = self._credential.get_token(scope)
+                else:
+                    self._token = self._credential.get_token(scope, tenant_id=challenge.tenant_id)
 
             # ignore mypy's warning -- although self._token is Optional, get_token raises when it fails to get a token
             request.http_request.headers["Authorization"] = f"Bearer {self._token.token}"  # type: ignore

sdk/keyvault/azure-keyvault-keys/azure/keyvault/keys/_shared/async_challenge_auth_policy.py

@@ -48,7 +48,11 @@ async def on_request(self, request: PipelineRequest) -> None:
             if self._need_new_token():
                 # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
                 scope = challenge.get_scope() or challenge.get_resource() + "/.default"
-                self._token = await self._credential.get_token(scope, tenant_id=challenge.tenant_id)
+                # Exclude tenant for AD FS authentication
+                if challenge.tenant_id and challenge.tenant_id.lower().endswith("adfs"):
+                    self._token = await self._credential.get_token(scope)
+                else:
+                    self._token = await self._credential.get_token(scope, tenant_id=challenge.tenant_id)
 
             # ignore mypy's warning -- although self._token is Optional, get_token raises when it fails to get a token
             request.http_request.headers["Authorization"] = f"Bearer {self._token.token}"  # type: ignore

sdk/keyvault/azure-keyvault-keys/azure/keyvault/keys/_shared/challenge_auth_policy.py

@@ -77,7 +77,11 @@ def on_request(self, request: PipelineRequest) -> None:
             if self._need_new_token:
                 # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
                 scope = challenge.get_scope() or challenge.get_resource() + "/.default"
-                self._token = self._credential.get_token(scope, tenant_id=challenge.tenant_id)
+                # Exclude tenant for AD FS authentication
+                if challenge.tenant_id and challenge.tenant_id.lower().endswith("adfs"):
+                    self._token = self._credential.get_token(scope)
+                else:
+                    self._token = self._credential.get_token(scope, tenant_id=challenge.tenant_id)
 
             # ignore mypy's warning -- although self._token is Optional, get_token raises when it fails to get a token
             request.http_request.headers["Authorization"] = f"Bearer {self._token.token}"  # type: ignore

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

@@ -262,7 +262,7 @@ def send(request):
                 assert not request.body
                 assert request.headers["Content-Length"] == "0"
                 return challenge
-            elif Requests.count == 2:
+            elif Requests.count in (2, 3):
                 # second request should be authorized according to challenge and have the expected content
                 assert request.headers["Content-Length"]
                 assert request.body == expected_content
@@ -276,13 +276,17 @@ def get_token(*_, **kwargs):
             return AccessToken(expected_token, 0)
 
         credential = Mock(get_token=Mock(wraps=get_token))
-        pipeline = Pipeline(policies=[ChallengeAuthPolicy(credential=credential)], transport=Mock(send=send))
+        policy = ChallengeAuthPolicy(credential=credential)
+        pipeline = Pipeline(policies=[policy], transport=Mock(send=send))
         request = HttpRequest("POST", get_random_url())
         request.set_bytes_body(expected_content)
         pipeline.run(request)
-
         assert credential.get_token.call_count == 1
 
+        # Regression test: https://github.com/Azure/azure-sdk-for-python/issues/33621
+        policy._token = None
+        pipeline.run(request)
+
     tenant = "tenant-id"
     # AD FS challenges have an unusual authority format; see https://github.com/Azure/azure-sdk-for-python/issues/28648
     endpoint = f"https://adfs.redmond.azurestack.corp.microsoft.com/adfs/{tenant}"

sdk/keyvault/azure-keyvault-keys/tests/test_challenge_auth_async.py

@@ -218,7 +218,7 @@ async def send(request):
                 assert not request.body
                 assert request.headers["Content-Length"] == "0"
                 return challenge
-            elif Requests.count == 2:
+            elif Requests.count in (2, 3):
                 # second request should be authorized according to challenge and have the expected content
                 assert request.headers["Content-Length"]
                 assert request.body == expected_content
@@ -232,15 +232,17 @@ async def get_token(*_, **kwargs):
             return AccessToken(expected_token, 0)
 
         credential = Mock(get_token=Mock(wraps=get_token))
-        pipeline = AsyncPipeline(
-            policies=[AsyncChallengeAuthPolicy(credential=credential)], transport=Mock(send=send)
-        )
+        policy = AsyncChallengeAuthPolicy(credential=credential)
+        pipeline = AsyncPipeline(policies=[policy], transport=Mock(send=send))
         request = HttpRequest("POST", get_random_url())
         request.set_bytes_body(expected_content)
         await pipeline.run(request)
-
         assert credential.get_token.call_count == 1
 
+        # Regression test: https://github.com/Azure/azure-sdk-for-python/issues/33621
+        policy._token = None
+        await pipeline.run(request)
+
     tenant = "tenant-id"
     # AD FS challenges have an unusual authority format; see https://github.com/Azure/azure-sdk-for-python/issues/28648
     endpoint = f"https://adfs.redmond.azurestack.corp.microsoft.com/adfs/{tenant}"

sdk/keyvault/azure-keyvault-secrets/azure/keyvault/secrets/_shared/async_challenge_auth_policy.py

@@ -48,7 +48,11 @@ async def on_request(self, request: PipelineRequest) -> None:
             if self._need_new_token():
                 # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
                 scope = challenge.get_scope() or challenge.get_resource() + "/.default"
-                self._token = await self._credential.get_token(scope, tenant_id=challenge.tenant_id)
+                # Exclude tenant for AD FS authentication
+                if challenge.tenant_id and challenge.tenant_id.lower().endswith("adfs"):
+                    self._token = await self._credential.get_token(scope)
+                else:
+                    self._token = await self._credential.get_token(scope, tenant_id=challenge.tenant_id)
 
             # ignore mypy's warning -- although self._token is Optional, get_token raises when it fails to get a token
             request.http_request.headers["Authorization"] = f"Bearer {self._token.token}"  # type: ignore

sdk/keyvault/azure-keyvault-secrets/azure/keyvault/secrets/_shared/challenge_auth_policy.py

@@ -77,7 +77,11 @@ def on_request(self, request: PipelineRequest) -> None:
             if self._need_new_token:
                 # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
                 scope = challenge.get_scope() or challenge.get_resource() + "/.default"
-                self._token = self._credential.get_token(scope, tenant_id=challenge.tenant_id)
+                # Exclude tenant for AD FS authentication
+                if challenge.tenant_id and challenge.tenant_id.lower().endswith("adfs"):
+                    self._token = self._credential.get_token(scope)
+                else:
+                    self._token = self._credential.get_token(scope, tenant_id=challenge.tenant_id)
 
             # ignore mypy's warning -- although self._token is Optional, get_token raises when it fails to get a token
             request.http_request.headers["Authorization"] = f"Bearer {self._token.token}"  # type: ignore