TokenCache.search() also wipes stale access tokens

Author: rayluo

msal/token_cache.py

@@ -130,7 +130,7 @@ def _is_matching(entry: dict, query: dict, target_set: set = None) -> bool:
             target_set <= set(entry.get("target", "").split())
             if target_set else True)
 
-    def search(self, credential_type, target=None, query=None):  # O(n) generator
+    def search(self, credential_type, target=None, query=None, *, now=None):  # O(n) generator
         """Returns a generator of matching entries.
 
         It is O(1) for AT hits, and O(n) for other types.
@@ -157,18 +157,32 @@ def search(self, credential_type, target=None, query=None):  # O(n) generator
         target_set = set(target)
         with self._lock:
             # O(n) search. The key is NOT used in search.
+            now = int(time.time() if now is None else now)
+            expired_access_tokens = [
+                # Especially when/if we key ATs by ephemeral fields such as key_id,
+                # stale ATs keyed by an old key_id would stay forever.
+                # Here we collect them for their removal.
+            ]
             for entry in self._cache.get(credential_type, {}).values():
+                if (  # Automatically delete expired access tokens
+                    credential_type == self.CredentialType.ACCESS_TOKEN
+                    and int(entry["expires_on"]) < now
+                ):
+                    expired_access_tokens.append(entry)  # Can't delete them within current for-loop
+                    continue
                 if (entry != preferred_result  # Avoid yielding the same entry twice
                     and self._is_matching(entry, query, target_set=target_set)
                 ):
                     yield entry
+            for at in expired_access_tokens:
+                self.remove_at(at)
 
-    def find(self, credential_type, target=None, query=None):
+    def find(self, credential_type, target=None, query=None, *, now=None):
         """Equivalent to list(search(...))."""
         warnings.warn(
             "Use list(search(...)) instead to explicitly get a list.",
             DeprecationWarning)
-        return list(self.search(credential_type, target=target, query=query))
+        return list(self.search(credential_type, target=target, query=query, now=now))
 
     def add(self, event, now=None):
         """Handle a token obtaining event, and add tokens into cache."""

tests/test_application.py

@@ -340,6 +340,7 @@ class TestApplicationForRefreshInBehaviors(unittest.TestCase):
     account = {"home_account_id": "{}.{}".format(uid, utid)}
     rt = "this is a rt"
     client_id = "my_app"
+    soon = 60  # application.py considers tokens within 5 minutes as expired
 
     @classmethod
     def setUpClass(cls):  # Initialization at runtime, not interpret-time
@@ -414,7 +415,8 @@ def mock_post(url, headers=None, *args, **kwargs):
 
     def test_expired_token_and_unavailable_aad_should_return_error(self):
         # a.k.a. Attempt refresh expired token when AAD unavailable
-        self.populate_cache(access_token="expired at", expires_in=-1, refresh_in=-900)
+        self.populate_cache(
+            access_token="expired at", expires_in=self.soon, refresh_in=-900)
         error = "something went wrong"
         def mock_post(url, headers=None, *args, **kwargs):
             self.assertEqual("4|84,3|", (headers or {}).get(CLIENT_CURRENT_TELEMETRY))
@@ -425,7 +427,8 @@ def mock_post(url, headers=None, *args, **kwargs):
 
     def test_expired_token_and_available_aad_should_return_new_token(self):
         # a.k.a. Attempt refresh expired token when AAD available
-        self.populate_cache(access_token="expired at", expires_in=-1, refresh_in=-900)
+        self.populate_cache(
+            access_token="expired at", expires_in=self.soon, refresh_in=-900)
         new_access_token = "new AT"
         new_refresh_in = 123
         def mock_post(url, headers=None, *args, **kwargs):

tests/test_token_cache.py

@@ -58,6 +58,7 @@ def testAddByAad(self):
         client_id = "my_client_id"
         id_token = build_id_token(
             oid="object1234", preferred_username="John Doe", aud=client_id)
+        now = 1000
         self.cache.add({
             "client_id": client_id,
             "scope": ["s2", "s1", "s3"],  # Not in particular order
@@ -66,7 +67,7 @@ def testAddByAad(self):
                 uid="uid", utid="utid",  # client_info
                 expires_in=3600, access_token="an access token",
                 id_token=id_token, refresh_token="a refresh token"),
-            }, now=1000)
+            }, now=now)
         access_token_entry = {
                 'cached_at': "1000",
                 'client_id': 'my_client_id',
@@ -84,7 +85,7 @@ def testAddByAad(self):
             self.at_key_maker(**access_token_entry)))
         self.assertIn(
             access_token_entry,
-            self.cache.find(self.cache.CredentialType.ACCESS_TOKEN),
+            self.cache.find(self.cache.CredentialType.ACCESS_TOKEN, now=now),
             "find(..., query=None) should not crash, even though MSAL does not use it")
         self.assertEqual(
             {
@@ -203,17 +204,20 @@ def testAddByAdfs(self):
                 "appmetadata-fs.msidlab8.com-my_client_id")
             )
 
-    def assertFoundAccessToken(self, *, scopes, query, data=None):
+    def assertFoundAccessToken(self, *, scopes, query, data=None, now=None):
         cached_at = None
         for cached_at in self.cache.search(
-                TokenCache.CredentialType.ACCESS_TOKEN, target=scopes, query=query):
+                TokenCache.CredentialType.ACCESS_TOKEN,
+                target=scopes, query=query, now=now,
+        ):
             for k, v in (data or {}).items():  # The extra data, if any
                 self.assertEqual(cached_at.get(k), v, f"AT should contain {k}={v}")
         self.assertTrue(cached_at, "AT should be cached and searchable")
         return cached_at
 
     def _test_data_should_be_saved_and_searchable_in_access_token(self, data):
         scopes = ["s2", "s1", "s3"]  # Not in particular order
+        now = 1000
         self.cache.add({
             "data": data,
             "client_id": "my_client_id",
@@ -223,8 +227,8 @@ def _test_data_should_be_saved_and_searchable_in_access_token(self, data):
                 uid="uid", utid="utid",  # client_info
                 expires_in=3600, access_token="an access token",
                 refresh_token="a refresh token"),
-            }, now=1000)
-        self.assertFoundAccessToken(scopes=scopes, data=data, query=dict(
+            }, now=now)
+        self.assertFoundAccessToken(scopes=scopes, data=data, now=now, query=dict(
             data,  # Also use the extra data as a query criteria
             client_id="my_client_id",
             environment="login.example.com",