refresh_in defaults to half of expires_in if expires_in >= 7200

Author: rayluo

msal/imds.py

@@ -16,6 +16,7 @@
     from collections import UserDict
 except:
     UserDict = dict  # The real UserDict is an old-style class which fails super()
+from .token_cache import TokenCache
 
 
 logger = logging.getLogger(__name__)
@@ -104,6 +105,7 @@ def _scope_to_resource(scope):  # This is an experimental reasonable-effort appr
 
 
 def _obtain_token(http_client, managed_identity, resource):
+    # A unified low-level API that talks to different Managed Identity
     if ("IDENTITY_ENDPOINT" in os.environ and "IDENTITY_HEADER" in os.environ
             and "IDENTITY_SERVER_THUMBPRINT" in os.environ
     ):
@@ -303,6 +305,12 @@ def _obtain_token_on_arc(http_client, endpoint, resource):
 
 
 class ManagedIdentityClient(object):
+    """A low level API that encapulate multiple managed identity backends:
+    VM, App Service, Azure Automation (Runbooks), Azure Function, Service Fabric,
+    and Azure Arc.
+
+    It also provides token cache support.
+    """
     _instance, _tenant = socket.getfqdn(), "managed_identity"  # Placeholders
 
     def __init__(self, http_client, managed_identity, token_cache=None):
@@ -319,16 +327,17 @@ def __init__(self, http_client, managed_identity, token_cache=None):
 
         :param token_cache:
             Optional. It accepts a :class:`msal.TokenCache` instance to store tokens.
+            It will use an in-memory token cache by default.
 
-        Example: Hard code a managed identity for your app::
+        Recipe 1: Hard code a managed identity for your app::
 
             import msal, requests
             client = msal.ManagedIdentityClient(
                 requests.Session(),
                 msal.UserAssignedManagedIdentity(client_id="foo"),
                 )
 
-        Recipe: Write once, run everywhere.
+        Recipe 2: Write once, run everywhere.
         If you use different managed identity on different deployment,
         you may use an environment variable (such as AZURE_MANAGED_IDENTITY)
         to store a json blob like
@@ -346,7 +355,7 @@ def __init__(self, http_client, managed_identity, token_cache=None):
         """
         self._http_client = http_client
         self._managed_identity = managed_identity
-        self._token_cache = token_cache
+        self._token_cache = token_cache or TokenCache()
 
     def acquire_token(self, resource=None):
         """Acquire token for the managed identity.
@@ -361,7 +370,8 @@ def acquire_token(self, resource=None):
         access_token_from_cache = None
         client_id_in_cache = self._managed_identity.get(
             ManagedIdentity.ID, "SYSTEM_ASSIGNED_MANAGED_IDENTITY")
-        if self._token_cache:
+        if True:  # Does not offer an "if not force_refresh" option, because
+                  # there would be built-in token cache in the service side anyway
             matches = self._token_cache.find(
                 self._token_cache.CredentialType.ACCESS_TOKEN,
                 target=[resource],
@@ -386,17 +396,26 @@ def acquire_token(self, resource=None):
                 if "refresh_on" in entry and int(entry["refresh_on"]) < now:  # aging
                     break  # With a fallback in hand, we break here to go refresh
                 return access_token_from_cache  # It is still good as new
-        result = _obtain_token(self._http_client, self._managed_identity, resource)
-        if self._token_cache and "access_token" in result:
-            self._token_cache.add(dict(
-                client_id=client_id_in_cache,
-                scope=[resource],
-                token_endpoint="https://{}/{}".format(self._instance, self._tenant),
-                response=result,
-                params={},
-                data={},
-                #grant_type="placeholder",
-            ))
-            return result
-        return access_token_from_cache or result
+        try:
+            result = _obtain_token(self._http_client, self._managed_identity, resource)
+            if "access_token" in result:
+                expires_in = result.get("expires_in", 3600)
+                if "refresh_in" not in result and expires_in >= 7200:
+                    result["refresh_in"] = int(expires_in / 2)
+                self._token_cache.add(dict(
+                    client_id=client_id_in_cache,
+                    scope=[resource],
+                    token_endpoint="https://{}/{}".format(self._instance, self._tenant),
+                    response=result,
+                    params={},
+                    data={},
+                    #grant_type="placeholder",
+                ))
+            if (result and "error" not in result) or (not access_token_from_cache):
+                return result
+        except:  # The exact HTTP exception is transportation-layer dependent
+            # Typically network error. Potential AAD outage?
+            if not access_token_from_cache:  # It means there is no fall back option
+                raise  # We choose to bubble up the exception
+        return access_token_from_cache
 

tests/test_mi.py

@@ -11,7 +11,6 @@
 
 from tests.http_client import MinimalResponse
 from msal import (
-    TokenCache,
     SystemAssignedManagedIdentity, UserAssignedManagedIdentity,
     ManagedIdentityClient)
 
@@ -46,7 +45,7 @@ def setUp(self):
                 # the client has no hard dependency on ManagedIdentity object
                 "ManagedIdentityIdType": "SystemAssigned", "Id": None,
             },
-            token_cache=TokenCache())
+            )
 
     def _test_token_cache(self, app):
         cache = app._token_cache._cache
@@ -140,7 +139,7 @@ def test_unified_api_service_should_ignore_unnecessary_client_id(self):
         self._test_happy_path(ManagedIdentityClient(
             requests.Session(),
             {"ManagedIdentityIdType": "ClientId", "Id": "foo"},
-            token_cache=TokenCache()))
+            ))
 
     def test_sf_error_should_be_normalized(self):
         raw_error = '''