Add ChallengeAuthenticationPolicy

sdk/core/azure-core/CHANGELOG.md

@@ -5,6 +5,8 @@
 ### New Features
 
 - Added `azure.core.credentials.AzureNamedKeyCredential` credential #17548.
+- Added `azure.core.pipeline.policies.ChallengeAuthenticationPolicy` and
+  `.AsyncChallengeAuthenticationPolicy`
 
 ## 1.13.0 (2021-04-02)
 

sdk/core/azure-core/azure/core/pipeline/policies/__init__.py

@@ -25,7 +25,12 @@
 # --------------------------------------------------------------------------
 
 from ._base import HTTPPolicy, SansIOHTTPPolicy, RequestHistory
-from ._authentication import BearerTokenCredentialPolicy, AzureKeyCredentialPolicy, AzureSasCredentialPolicy
+from ._authentication import (
+    BearerTokenCredentialPolicy,
+    AzureKeyCredentialPolicy,
+    AzureSasCredentialPolicy,
+    ChallengeAuthenticationPolicy,
+)
 from ._custom_hook import CustomHookPolicy
 from ._redirect import RedirectPolicy
 from ._retry import RetryPolicy, RetryMode
@@ -44,6 +49,7 @@
     'HTTPPolicy',
     'SansIOHTTPPolicy',
     'BearerTokenCredentialPolicy',
+    'ChallengeAuthenticationPolicy',
     'AzureKeyCredentialPolicy',
     'AzureSasCredentialPolicy',
     'HeadersPolicy',
@@ -65,12 +71,13 @@
 
 try:
     from ._base_async import AsyncHTTPPolicy
-    from ._authentication_async import AsyncBearerTokenCredentialPolicy
+    from ._authentication_async import AsyncBearerTokenCredentialPolicy, AsyncChallengeAuthenticationPolicy
     from ._redirect_async import AsyncRedirectPolicy
     from ._retry_async import AsyncRetryPolicy
     __all__.extend([
         'AsyncHTTPPolicy',
         'AsyncBearerTokenCredentialPolicy',
+        'AsyncChallengeAuthenticationPolicy',
         'AsyncRedirectPolicy',
         'AsyncRetryPolicy'
     ])

sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py

@@ -6,7 +6,7 @@
 import time
 import six
 
-from . import SansIOHTTPPolicy
+from . import HTTPPolicy, SansIOHTTPPolicy
 from ...exceptions import ServiceRequestError
 
 try:
@@ -18,7 +18,27 @@
     # pylint:disable=unused-import
     from typing import Any, Dict, Optional
     from azure.core.credentials import AccessToken, TokenCredential, AzureKeyCredential, AzureSasCredential
-    from azure.core.pipeline import PipelineRequest
+    from azure.core.pipeline import PipelineRequest, PipelineResponse
+
+
+def _enforce_https(request):
+    # type: (PipelineRequest) -> None
+    """Raise ServiceRequestError if the request URL is non-HTTPS and the sender did not specify "enforce_https=False"
+    """
+
+    # move 'enforce_https' from options to context so it persists
+    # across retries but isn't passed to a transport implementation
+    option = request.context.options.pop("enforce_https", None)
+
+    # True is the default setting; we needn't preserve an explicit opt in to the default behavior
+    if option is False:
+        request.context["enforce_https"] = option
+
+    enforce_https = request.context.get("enforce_https", True)
+    if enforce_https and not request.http_request.url.lower().startswith("https"):
+        raise ServiceRequestError(
+            "Bearer token authentication is not permitted for non-TLS protected (non-https) URLs."
+        )
 
 
 # pylint:disable=too-few-public-methods
@@ -40,20 +60,7 @@ def __init__(self, credential, *scopes, **kwargs):  # pylint:disable=unused-argu
     @staticmethod
     def _enforce_https(request):
         # type: (PipelineRequest) -> None
-
-        # move 'enforce_https' from options to context so it persists
-        # across retries but isn't passed to a transport implementation
-        option = request.context.options.pop("enforce_https", None)
-
-        # True is the default setting; we needn't preserve an explicit opt in to the default behavior
-        if option is False:
-            request.context["enforce_https"] = option
-
-        enforce_https = request.context.get("enforce_https", True)
-        if enforce_https and not request.http_request.url.lower().startswith("https"):
-            raise ServiceRequestError(
-                "Bearer token authentication is not permitted for non-TLS protected (non-https) URLs."
-            )
+        return _enforce_https(request)
 
     @staticmethod
     def _update_headers(headers, token):
@@ -94,6 +101,84 @@ def on_request(self, request):
         self._update_headers(request.http_request.headers, self._token.token)
 
 
+class ChallengeAuthenticationPolicy(HTTPPolicy):
+    """Base class for policies that authorize requests with bearer tokens and expect authentication challenges
+
+    :param ~azure.core.credentials.TokenCredential credential: an object which can provide access tokens, such as a
+        credential from :mod:`azure.identity`
+    :param str scopes: required authentication scopes
+    """
+
+    def __init__(self, credential, *scopes, **kwargs):  # pylint:disable=unused-argument
+        # type: (TokenCredential, *str, **Any) -> None
+        super(ChallengeAuthenticationPolicy, self).__init__()
+        self._scopes = scopes
+        self._credential = credential
+        self._token = None  # type: Optional[AccessToken]
+
+    def _need_new_token(self):
+        # type: () -> bool
+        return not self._token or self._token.expires_on - time.time() < 300
+
+    def authorize_request(self, request, *scopes, **kwargs):
+        # type: (PipelineRequest, *str, **Any) -> None
+        """Acquire a token from the credential and authorize the request with it.
+
+        Keyword arguments are passed to the credential's get_token method. The token will be cached and used to
+        authorize future requests.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request
+        :param str scopes: required scopes of authentication
+        """
+        self._token = self._credential.get_token(*scopes, **kwargs)
+        request.http_request.headers["Authorization"] = "Bearer " + self._token.token
+
+    def on_request(self, request):
+        # type: (PipelineRequest) -> None
+        """Called before the policy sends a request.
+
+        The base implementation authorizes the request with a bearer token.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request
+        """
+
+        if self._token is None or self._need_new_token():
+            self._token = self._credential.get_token(*self._scopes)
+        request.http_request.headers["Authorization"] = "Bearer " + self._token.token
+
+    def send(self, request):
+        # type: (PipelineRequest) -> PipelineResponse
+        """Authorizes a request with a bearer token, possibly handling an authentication challenge
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request
+        """
+        _enforce_https(request)
+
+        self.on_request(request)
+
+        response = self.next.send(request)
+
+        if response.http_response.status_code == 401:
+            self._token = None  # any cached token is invalid
+            if "WWW-Authenticate" in response.http_response.headers and self.on_challenge(request, response):
+                response = self.next.send(request)
+
+        return response
+
+    def on_challenge(self, request, response):
+        # type: (PipelineRequest, PipelineResponse) -> bool
+        """Authorize request according to an authentication challenge
+
+        This method is called when the resource provider responds 401 with a WWW-Authenticate header.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request which elicited an authentication challenge
+        :param ~azure.core.pipeline.PipelineResponse response: the resource provider's response
+        :returns: a bool indicating whether the policy should send the request
+        """
+        # pylint:disable=unused-argument,no-self-use
+        return False
+
+
 class AzureKeyCredentialPolicy(SansIOHTTPPolicy):
     """Adds a key header for the provided credential.
 

sdk/core/azure-core/azure/core/pipeline/policies/_authentication_async.py

@@ -4,10 +4,17 @@
 # license information.
 # -------------------------------------------------------------------------
 import asyncio
+import time
+from typing import TYPE_CHECKING
 
-from azure.core.pipeline import PipelineRequest
-from azure.core.pipeline.policies import SansIOHTTPPolicy
-from azure.core.pipeline.policies._authentication import _BearerTokenCredentialPolicyBase
+from azure.core.pipeline.policies import AsyncHTTPPolicy, SansIOHTTPPolicy
+from azure.core.pipeline.policies._authentication import _BearerTokenCredentialPolicyBase, _enforce_https
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+    from azure.core.credentials import AccessToken
+    from azure.core.credentials_async import AsyncTokenCredential
+    from azure.core.pipeline import PipelineRequest, PipelineResponse
 
 
 class AsyncBearerTokenCredentialPolicy(_BearerTokenCredentialPolicyBase, SansIOHTTPPolicy):
@@ -23,7 +30,7 @@ def __init__(self, credential, *scopes, **kwargs):
         super().__init__(credential, *scopes, **kwargs)
         self._lock = asyncio.Lock()
 
-    async def on_request(self, request: PipelineRequest):  # pylint:disable=invalid-overridden-method
+    async def on_request(self, request: "PipelineRequest"):  # pylint:disable=invalid-overridden-method
         """Adds a bearer token Authorization header to request and sends request to next policy.
 
         :param request: The pipeline request object to be modified.
@@ -36,3 +43,85 @@ async def on_request(self, request: PipelineRequest):  # pylint:disable=invalid-
             if self._need_new_token:
                 self._token = await self._credential.get_token(*self._scopes)  # type: ignore
         self._update_headers(request.http_request.headers, self._token.token)
+
+
+class AsyncChallengeAuthenticationPolicy(AsyncHTTPPolicy):
+    """Base class for policies that authorize requests with bearer tokens and expect authentication challenges
+
+    :param ~azure.core.credentials.AsyncTokenCredential credential: an object which can asynchronously provide access
+        tokens, such as a credential from :mod:`azure.identity.aio`
+    :param str scopes: required authentication scopes
+    """
+
+    def __init__(self, credential: "AsyncTokenCredential", *scopes: str, **kwargs: "Any") -> None:
+        # pylint:disable=unused-argument
+        super().__init__()
+        self._credential = credential
+        self._lock = asyncio.Lock()
+        self._scopes = scopes
+        self._token = None  # type: Optional[AccessToken]
+
+    def _need_new_token(self) -> bool:
+        return not self._token or self._token.expires_on - time.time() < 300
+
+    async def on_request(self, request: "PipelineRequest") -> None:
+        """Called before the policy sends a request.
+
+        The base implementation authorizes the request with a bearer token.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request
+        """
+
+        if self._token is None or self._need_new_token():
+            async with self._lock:
+                # double check because another coroutine may have acquired a token while we waited to acquire the lock
+                if not self._token or self._need_new_token():
+                    self._token = await self._credential.get_token(*self._scopes)
+
+        request.http_request.headers["Authorization"] = "Bearer " + self._token.token
+
+    async def authorize_request(self, request: "PipelineRequest", *scopes: str, **kwargs: "Any") -> None:
+        """Acquire a token from the credential and authorize the request with it.
+
+        Keyword arguments are passed to the credential's get_token method. The token will be cached and used to
+        authorize future requests.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request
+        :param str scopes: required scopes of authentication
+        """
+
+        async with self._lock:
+            self._token = await self._credential.get_token(*scopes, **kwargs)
+        request.http_request.headers["Authorization"] = "Bearer " + self._token.token
+
+    async def send(self, request: "PipelineRequest") -> "PipelineResponse":
+        """Authorizes a request with a bearer token, possibly handling an authentication challenge
+
+        :param ~azure.core.pipeline.PipelineRequest request: The request
+        """
+        _enforce_https(request)
+
+        await self.on_request(request)
+
+        response = await self.next.send(request)
+
+        if response.http_response.status_code == 401:
+            self._token = None  # any cached token is invalid
+            if "WWW-Authenticate" in response.http_response.headers:
+                request_authorized = await self.on_challenge(request, response)
+                if request_authorized:
+                    response = await self.next.send(request)
+
+        return response
+
+    async def on_challenge(self, request: "PipelineRequest", response: "PipelineResponse") -> bool:
+        """Authorize request according to an authentication challenge
+
+        This method is called when the resource provider responds 401 with a WWW-Authenticate header.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request which elicited an authentication challenge
+        :param ~azure.core.pipeline.PipelineResponse response: the resource provider's response
+        :returns: a bool indicating whether the policy should send the request
+        """
+        # pylint:disable=unused-argument,no-self-use
+        return False