token refresh offset

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

@@ -64,7 +64,12 @@ def get_token(self, *scopes, **kwargs):
             self._authorization_code = None  # auth codes are single-use
             return token
 
-        token = self._client.get_cached_access_token(scopes) or self._redeem_refresh_token(scopes, **kwargs)
+        token = self._client.get_cached_access_token(scopes)
+        if not token:
+            token = self._redeem_refresh_token(scopes, **kwargs)
+        elif self._client.is_refresh(token):
+            self._redeem_refresh_token(scopes, **kwargs)
+
         if not token:
             raise ClientAuthenticationError(
                 message="No authorization code, cached access token, or refresh token available."

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

@@ -48,6 +48,8 @@ def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
         token = self._client.get_cached_access_token(scopes, query={"client_id": self._client_id})
         if not token:
             token = self._client.obtain_token_by_client_certificate(scopes, self._certificate, **kwargs)
+        elif self._client.is_refresh(token):
+            self._client.obtain_token_by_client_certificate(scopes, self._certificate, **kwargs)
         return token
 
     def _get_auth_client(self, tenant_id, client_id, **kwargs):

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

@@ -49,6 +49,8 @@ def get_token(self, *scopes, **kwargs):
         token = self._client.get_cached_access_token(scopes, query={"client_id": self._client_id})
         if not token:
             token = self._client.obtain_token_by_client_secret(scopes, self._secret, **kwargs)
+        elif self._client.is_refresh(token):
+            self._client.obtain_token_by_client_secret(scopes, self._secret, **kwargs)
         return token
 
     def _get_auth_client(self, tenant_id, client_id, **kwargs):

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

@@ -48,9 +48,17 @@ def get_token(self, *scopes, **kwargs):
 
         token = self._client.get_cached_access_token(scopes)
 
-        if token:
-            return token
+        if not token:
+            token = self._redeem_refresh_token(scopes, **kwargs)
+        elif self._client.is_refresh(token):
+            try:
+                self._redeem_refresh_token(scopes, **kwargs)
+            except Exception:  # pylint: disable=broad-except
+                pass
+        return token
 
+    def _redeem_refresh_token(self, scopes, **kwargs):
+        # type: (Sequence[str], **Any) -> Optional[AccessToken]
         if not self._refresh_token:
             self._refresh_token = get_credentials()
             if not self._refresh_token:

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

@@ -48,13 +48,16 @@ def __init__(self, tenant_id, client_id, authority=None, cache=None, **kwargs):
         self._cache = cache or TokenCache()
         self._client_id = client_id
         self._pipeline = self._build_pipeline(**kwargs)
+        self._token_refresh_timeout = 30    # default 30s
+        self._token_refresh_offset = 120    # default 2 min
+        self._last_refresh_time = int(time.time())
 
     def get_cached_access_token(self, scopes, query=None):
         # type: (Sequence[str], Optional[dict]) -> Optional[AccessToken]
         tokens = self._cache.find(TokenCache.CredentialType.ACCESS_TOKEN, target=list(scopes), query=query)
         for token in tokens:
             expires_on = int(token["expires_on"])
-            if expires_on - 300 > int(time.time()):
+            if expires_on - 30 > int(time.time()):
                 return AccessToken(token["secret"], expires_on)
         return None
 
@@ -63,6 +66,19 @@ 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 is_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_offset:
+            return False
+        return True
+
+
     @abc.abstractmethod
     def obtain_token_by_authorization_code(self, scopes, code, redirect_uri, client_secret=None, **kwargs):
         pass
@@ -85,6 +101,7 @@ def _build_pipeline(self, config=None, policies=None, transport=None, **kwargs):
 
     def _process_response(self, response, request_time):
         # type: (PipelineResponse, int) -> AccessToken
+        self._last_refresh_time = time.time()   # no matter succeed or not, update the last refresh time
 
         content = ContentDecodePolicy.deserialize_from_http_generics(response.http_response)
 

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

@@ -80,6 +80,8 @@ async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
         token = self._client.get_cached_access_token(scopes)
         if not token:
             token = await self._redeem_refresh_token(scopes, **kwargs)
+        elif self._client.is_refresh(token):
+            await self._redeem_refresh_token(scopes, **kwargs)
 
         if not token:
             raise ClientAuthenticationError(

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

@@ -54,6 +54,8 @@ async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # py
         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.is_refresh(token):
+            await self._client.obtain_token_by_client_certificate(scopes, self._certificate, **kwargs)
         return token
 
     def _get_auth_client(self, tenant_id, client_id, **kwargs):

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

@@ -55,6 +55,8 @@ async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
         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_secret(scopes, self._secret, **kwargs)
+        elif self._client.is_refresh(token):
+            await self._client.obtain_token_by_client_secret(scopes, self._secret, **kwargs)
         return token
 
     def _get_auth_client(self, tenant_id, client_id, **kwargs):

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

@@ -52,9 +52,16 @@ async def get_token(self, *scopes, **kwargs):
             raise ValueError("'get_token' requires at least one scope")
 
         token = self._client.get_cached_access_token(scopes)
-        if token:
-            return token
+        if not token:
+            token = await self._redeem_refresh_token(scopes, **kwargs)
+        elif self._client.is_refresh(token):
+            try:
+                await self._redeem_refresh_token(scopes, **kwargs)
+            except Exception:  # pylint: disable=broad-except
+                pass
+        return token
 
+    async def _redeem_refresh_token(self, scopes: "Sequence[str]", **kwargs: "Any") -> "Optional[AccessToken]":
         if not self._refresh_token:
             self._refresh_token = get_credentials()
             if not self._refresh_token:

sdk/identity/azure-identity/tests/test_token_refresh_offset.py

@@ -0,0 +1,35 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import time
+from azure.identity._internal.aad_client import AadClient
+from azure.core.credentials import AccessToken
+import pytest
+
+try:
+    from unittest import mock
+except ImportError:  # python < 3.3
+    import mock
+
+
+def test_if_refresh():
+    client = AadClient("test", "test")
+    now = int(time.time())
+
+    # do not need refresh
+    token = AccessToken("token", now + 500)
+    is_refresh = client.is_refresh(token)
+    assert not is_refresh
+
+    # need refresh
+    token = AccessToken("token", now + 100)
+    client._last_refresh_time = now - 500
+    is_refresh = client.is_refresh(token)
+    assert is_refresh
+
+    # not exceed cool down time, do not refresh
+    token = AccessToken("token", now + 100)
+    client._last_refresh_time = now - 5
+    is_refresh = client.is_refresh(token)
+    assert not is_refresh

sdk/identity/azure-identity/tests/test_vscode_credential.py

@@ -92,9 +92,13 @@ def test_cache_refresh_token():
 
 
 def test_no_obtain_token_if_cached():
+    def mock_is_refresh(token):
+        return False
+
     expected_token = AccessToken("token", 42)
 
     mock_client = mock.Mock(spec=object)
+    mock_client.is_refresh = mock_is_refresh
     mock_client.obtain_token_by_refresh_token = mock.Mock(return_value=expected_token)
     mock_client.get_cached_access_token = mock.Mock(return_value="VALUE")
 

sdk/identity/azure-identity/tests/test_vscode_credential_async.py

@@ -95,9 +95,13 @@ async def test_cache_refresh_token():
 
 @pytest.mark.asyncio
 async def test_no_obtain_token_if_cached():
+    def mock_is_refresh(token):
+        return False
+
     expected_token = AccessToken("token", 42)
 
     mock_client = mock.Mock(spec=object)
+    mock_client.is_refresh = mock_is_refresh
     token_by_refresh_token = mock.Mock(return_value=expected_token)
     mock_client.obtain_token_by_refresh_token = wrap_in_future(token_by_refresh_token)
     mock_client.get_cached_access_token = mock.Mock(return_value="VALUE")