Consolidate proactive token refreshing code

sdk/identity/azure-identity/azure/identity/_authn_client.py

@@ -75,18 +75,6 @@ def __init__(self, endpoint=None, authority=None, tenant=None, **kwargs):  # pyl
     def auth_url(self):
         return self._auth_url
 
-    def should_refresh(self, token):
-        # type: (AccessToken) -> bool
-        """ check if the token needs refresh or not
-        """
-        expires_on = int(token.expires_on)
-        now = int(time.time())
-        if expires_on - now > self._token_refresh_offset:
-            return False
-        if now - self._last_refresh_time < self._token_refresh_retry_delay:
-            return False
-        return True
-
     def get_cached_token(self, scopes):
         # type: (Iterable[str]) -> Optional[AccessToken]
         tokens = self._cache.find(TokenCache.CredentialType.ACCESS_TOKEN, target=list(scopes))

sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py

@@ -6,15 +6,15 @@
 
 from azure.core.exceptions import ClientAuthenticationError
 from .._internal.aad_client import AadClient
-from .._internal.decorators import log_get_token
+from .._internal.get_token_mixin import GetTokenMixin
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
     from typing import Any, Iterable, Optional
     from azure.core.credentials import AccessToken
 
 
-class AuthorizationCodeCredential(object):
+class AuthorizationCodeCredential(GetTokenMixin):
     """Authenticates by redeeming an authorization code previously obtained from Azure Active Directory.
 
     See https://docs.microsoft.com/azure/active-directory/develop/v2-oauth2-auth-code-flow for more information
@@ -38,8 +38,8 @@ def __init__(self, tenant_id, client_id, authorization_code, redirect_uri, **kwa
         self._client_secret = kwargs.pop("client_secret", None)
         self._client = kwargs.pop("client", None) or AadClient(tenant_id, client_id, **kwargs)
         self._redirect_uri = redirect_uri
+        super(AuthorizationCodeCredential, self).__init__()
 
-    @log_get_token("AuthorizationCodeCredential")
     def get_token(self, *scopes, **kwargs):
         # type: (*str, **Any) -> AccessToken
         """Request an access token for `scopes`.
@@ -56,38 +56,31 @@ def get_token(self, *scopes, **kwargs):
           attribute gives a reason. Any error response from Azure Active Directory is available as the error's
           ``response`` attribute.
         """
-        if not scopes:
-            raise ValueError("'get_token' requires at least one scope")
+        return super(AuthorizationCodeCredential, self).get_token(*scopes)
 
+    def _acquire_token_silently(self, *scopes):
+        # type: (*str) -> Optional[AccessToken]
+        return self._client.get_cached_access_token(scopes)
+
+    def _request_token(self, *scopes, **kwargs):
+        # type: (*str, **Any) -> AccessToken
         if self._authorization_code:
             token = self._client.obtain_token_by_authorization_code(
                 scopes=scopes, code=self._authorization_code, redirect_uri=self._redirect_uri, **kwargs
             )
             self._authorization_code = None  # auth codes are single-use
             return token
 
-        token = self._client.get_cached_access_token(scopes)
-        if not token:
-            token = self._redeem_refresh_token(scopes, **kwargs)
-        elif self._client.should_refresh(token):
-            try:
-                self._redeem_refresh_token(scopes, **kwargs)
-            except Exception:  # pylint: disable=broad-except
-                pass
+        token = None
+        for refresh_token in self._client.get_cached_refresh_tokens(scopes):
+            if "secret" in refresh_token:
+                token = self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs)
+                if token:
+                    break
 
         if not token:
             raise ClientAuthenticationError(
                 message="No authorization code, cached access token, or refresh token available."
             )
 
         return token
-
-    def _redeem_refresh_token(self, scopes, **kwargs):
-        # type: (Iterable[str], **Any) -> Optional[AccessToken]
-        for refresh_token in self._client.get_cached_refresh_tokens(scopes):
-            if "secret" not in refresh_token:
-                continue
-            token = self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs)
-            if token:
-                return token
-        return None

sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py

@@ -23,6 +23,7 @@
 from .._authn_client import AuthnClient
 from .._constants import Endpoints, EnvironmentVariables
 from .._internal.decorators import log_get_token
+from .._internal.get_token_mixin import GetTokenMixin
 from .._internal.user_agent import USER_AGENT
 
 try:
@@ -156,7 +157,7 @@ def _create_config(**kwargs):
     }
 
 
-class ImdsCredential(_ManagedIdentityBase):
+class ImdsCredential(_ManagedIdentityBase, GetTokenMixin):
     """Authenticates with a managed identity via the IMDS endpoint.
 
     :keyword str client_id: ID of a user-assigned identity. Leave unspecified to use a system-assigned identity.
@@ -167,16 +168,12 @@ def __init__(self, **kwargs):
         super(ImdsCredential, self).__init__(endpoint=Endpoints.IMDS, client_cls=AuthnClient, **kwargs)
         self._endpoint_available = None  # type: Optional[bool]
 
-    def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
-        # type: (*str, **Any) -> AccessToken
-        """Request an access token for `scopes`.
-
-        This method is called automatically by Azure SDK clients.
+    def _acquire_token_silently(self, *scopes):
+        # type: (*str) -> Optional[AccessToken]
+        return self._client.get_cached_token(scopes)
 
-        :param str scopes: desired scope for the access token. This credential allows only one scope per request.
-        :rtype: :class:`azure.core.credentials.AccessToken`
-        :raises ~azure.identity.CredentialUnavailableError: the IMDS endpoint is unreachable
-        """
+    def _request_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
+        # type: (*str, **Any) -> AccessToken
         if self._endpoint_available is None:
             # Lacking another way to determine whether the IMDS endpoint is listening,
             # we send a request it would immediately reject (missing a required header),
@@ -199,18 +196,6 @@ def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
         if len(scopes) != 1:
             raise ValueError("This credential requires exactly one scope per token request.")
 
-        token = self._client.get_cached_token(scopes)
-        if not token:
-            token = self._refresh_token(*scopes)
-        elif self._client.should_refresh(token):
-            try:
-                token = self._refresh_token(*scopes)
-            except Exception:  # pylint: disable=broad-except
-                pass
-
-        return token
-
-    def _refresh_token(self, *scopes):
         resource = scopes[0]
         if resource.endswith("/.default"):
             resource = resource[: -len("/.default")]

sdk/identity/azure-identity/azure/identity/_credentials/vscode.py

@@ -9,7 +9,7 @@
 from .._constants import AZURE_VSCODE_CLIENT_ID
 from .._internal import validate_tenant_id
 from .._internal.aad_client import AadClient
-from .._internal.decorators import log_get_token
+from .._internal.get_token_mixin import GetTokenMixin
 
 if sys.platform.startswith("win"):
     from .._internal.win_vscode_adapter import get_credentials
@@ -20,11 +20,11 @@
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
-    from typing import Any, Iterable, Optional
+    from typing import Any, Optional
     from azure.core.credentials import AccessToken
 
 
-class VisualStudioCodeCredential(object):
+class VisualStudioCodeCredential(GetTokenMixin):
     """Authenticates as the Azure user signed in to Visual Studio Code.
 
     :keyword str authority: Authority of an Azure Active Directory endpoint, for example 'login.microsoftonline.com',
@@ -36,14 +36,14 @@ class VisualStudioCodeCredential(object):
 
     def __init__(self, **kwargs):
         # type: (**Any) -> None
+        super(VisualStudioCodeCredential, self).__init__()
         self._refresh_token = None
         self._client = kwargs.pop("_client", None)
         self._tenant_id = kwargs.pop("tenant_id", None) or "organizations"
         validate_tenant_id(self._tenant_id)
         if not self._client:
             self._client = AadClient(self._tenant_id, AZURE_VSCODE_CLIENT_ID, **kwargs)
 
-    @log_get_token("VisualStudioCodeCredential")
     def get_token(self, *scopes, **kwargs):
         # type: (*str, **Any) -> AccessToken
         """Request an access token for `scopes` as the user currently signed in to Visual Studio Code.
@@ -55,31 +55,21 @@ def get_token(self, *scopes, **kwargs):
         :raises ~azure.identity.CredentialUnavailableError: the credential cannot retrieve user details from Visual
           Studio Code
         """
-        if not scopes:
-            raise ValueError("'get_token' requires at least one scope")
-
         if self._tenant_id.lower() == "adfs":
             raise CredentialUnavailableError(
                 message="VisualStudioCodeCredential authentication unavailable. ADFS is not supported."
             )
+        return super(VisualStudioCodeCredential, self).get_token(*scopes, **kwargs)
 
-        token = self._client.get_cached_access_token(scopes)
-
-        if not token:
-            token = self._redeem_refresh_token(scopes, **kwargs)
-        elif self._client.should_refresh(token):
-            try:
-                self._redeem_refresh_token(scopes, **kwargs)
-            except Exception:  # pylint: disable=broad-except
-                pass
-        return token
+    def _acquire_token_silently(self, *scopes):
+        # type: (*str) -> Optional[AccessToken]
+        return self._client.get_cached_access_token(scopes)
 
-    def _redeem_refresh_token(self, scopes, **kwargs):
-        # type: (Iterable[str], **Any) -> Optional[AccessToken]
+    def _request_token(self, *scopes, **kwargs):
+        # type: (*str, **Any) -> AccessToken
         if not self._refresh_token:
             self._refresh_token = get_credentials()
             if not self._refresh_token:
                 raise CredentialUnavailableError(message="Failed to get Azure user details from Visual Studio Code.")
 
-        token = self._client.obtain_token_by_refresh_token(scopes, self._refresh_token, **kwargs)
-        return token
+        return self._client.obtain_token_by_refresh_token(scopes, self._refresh_token, **kwargs)

sdk/identity/azure-identity/azure/identity/_internal/aad_client_base.py

@@ -69,19 +69,6 @@ def get_cached_refresh_tokens(self, scopes):
         """Assumes all cached refresh tokens belong to the same user"""
         return self._cache.find(TokenCache.CredentialType.REFRESH_TOKEN, target=list(scopes))
 
-    def should_refresh(self, token):
-        # type: (AccessToken) -> bool
-        """ check if the token needs refresh or not
-        """
-        expires_on = int(token.expires_on)
-        now = int(time.time())
-        if expires_on - now > self._token_refresh_offset:
-            return False
-        if now - self._last_refresh_time < self._token_refresh_retry_delay:
-            return False
-        return True
-
-
     @abc.abstractmethod
     def obtain_token_by_authorization_code(self, scopes, code, redirect_uri, client_secret=None, **kwargs):
         pass

sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py

@@ -6,15 +6,15 @@
 
 from azure.core.exceptions import ClientAuthenticationError
 from .._internal import AadClient, AsyncContextManager
-from .._internal.decorators import log_get_token_async
+from .._internal.get_token_mixin import GetTokenMixin
 
 if TYPE_CHECKING:
     # pylint:disable=unused-import,ungrouped-imports
     from typing import Any, Iterable, Optional
     from azure.core.credentials import AccessToken
 
 
-class AuthorizationCodeCredential(AsyncContextManager):
+class AuthorizationCodeCredential(AsyncContextManager, GetTokenMixin):
     """Authenticates by redeeming an authorization code previously obtained from Azure Active Directory.
 
     See https://docs.microsoft.com/azure/active-directory/develop/v2-oauth2-auth-code-flow for more information
@@ -50,8 +50,8 @@ def __init__(
         self._client_secret = kwargs.pop("client_secret", None)
         self._client = kwargs.pop("client", None) or AadClient(tenant_id, client_id, **kwargs)
         self._redirect_uri = redirect_uri
+        super().__init__()
 
-    @log_get_token_async
     async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
         """Request an access token for `scopes`.
 
@@ -67,37 +67,29 @@ async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
           attribute gives a reason. Any error response from Azure Active Directory is available as the error's
           ``response`` attribute.
         """
-        if not scopes:
-            raise ValueError("'get_token' requires at least one scope")
+        return await super().get_token(*scopes, **kwargs)
 
+    async def _acquire_token_silently(self, *scopes: str) -> "Optional[AccessToken]":
+        return self._client.get_cached_access_token(scopes)
+
+    async def _request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
         if self._authorization_code:
             token = await self._client.obtain_token_by_authorization_code(
                 scopes=scopes, code=self._authorization_code, redirect_uri=self._redirect_uri, **kwargs
             )
-
             self._authorization_code = None  # auth codes are single-use
             return token
 
-        token = self._client.get_cached_access_token(scopes)
-        if not token:
-            token = await self._redeem_refresh_token(scopes, **kwargs)
-        elif self._client.should_refresh(token):
-            try:
-                await self._redeem_refresh_token(scopes, **kwargs)
-            except Exception:  # pylint: disable=broad-except
-                pass
+        token = None
+        for refresh_token in self._client.get_cached_refresh_tokens(scopes):
+            if "secret" in refresh_token:
+                token = await self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs)
+                if token:
+                    break
+
         if not token:
             raise ClientAuthenticationError(
                 message="No authorization code, cached access token, or refresh token available."
             )
 
         return token
-
-    async def _redeem_refresh_token(self, scopes: "Iterable[str]", **kwargs: "Any") -> "Optional[AccessToken]":
-        for refresh_token in self._client.get_cached_refresh_tokens(scopes):
-            if "secret" not in refresh_token:
-                continue
-            token = await self._client.obtain_token_by_refresh_token(scopes, refresh_token["secret"], **kwargs)
-            if token:
-                return token
-        return None

sdk/identity/azure-identity/azure/identity/aio/_credentials/certificate.py

@@ -7,17 +7,17 @@
 import msal
 
 from .._internal import AadClient, AsyncContextManager
-from .._internal.decorators import log_get_token_async
+from .._internal.get_token_mixin import GetTokenMixin
 from ..._credentials.certificate import get_client_credential
 from ..._internal import AadClientCertificate, validate_tenant_id
-from ..._persistent_cache import _load_persistent_cache, TokenCachePersistenceOptions
+from ..._persistent_cache import _load_persistent_cache
 
 if TYPE_CHECKING:
     from typing import Any, Optional
     from azure.core.credentials import AccessToken
 
 
-class CertificateCredential(AsyncContextManager):
+class CertificateCredential(AsyncContextManager, GetTokenMixin):
     """Authenticates as a service principal using a certificate.
 
     The certificate must have an RSA private key, because this credential signs assertions using RS256.
@@ -57,6 +57,7 @@ def __init__(self, tenant_id, client_id, certificate_path=None, **kwargs):
 
         self._client = AadClient(tenant_id, client_id, cache=cache, **kwargs)
         self._client_id = client_id
+        super().__init__()
 
     async def __aenter__(self):
         await self._client.__aenter__()
@@ -67,27 +68,8 @@ async def close(self):
 
         await self._client.__aexit__()
 
-    @log_get_token_async
-    async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # pylint:disable=unused-argument
-        """Asynchronously request an access token for `scopes`.
-
-        This method is called automatically by Azure SDK clients.
-
-        :param str scopes: desired scopes for the access token. This method requires at least one scope.
-        :rtype: :class:`azure.core.credentials.AccessToken`
-        :raises ~azure.core.exceptions.ClientAuthenticationError: authentication failed. The error's ``message``
-          attribute gives a reason. Any error response from Azure Active Directory is available as the error's
-          ``response`` attribute.
-        """
-        if not scopes:
-            raise ValueError("'get_token' requires at least one scope")
-
-        token = self._client.get_cached_access_token(scopes, query={"client_id": self._client_id})
-        if not token:
-            token = await self._client.obtain_token_by_client_certificate(scopes, self._certificate, **kwargs)
-        elif self._client.should_refresh(token):
-            try:
-                await self._client.obtain_token_by_client_certificate(scopes, self._certificate, **kwargs)
-            except Exception:  # pylint: disable=broad-except
-                pass
-        return token
+    async def _acquire_token_silently(self, *scopes: str) -> "Optional[AccessToken]":
+        return self._client.get_cached_access_token(scopes, query={"client_id": self._client_id})
+
+    async def _request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
+        return await self._client.obtain_token_by_client_certificate(scopes, self._certificate, **kwargs)