Add the missing token query check

Author: rayluo

msal/token_cache.py

@@ -117,6 +117,12 @@ def _get(self, credential_type, key, default=None):  # O(1)
         with self._lock:
             return self._cache.get(credential_type, {}).get(key, default)
 
+    @staticmethod
+    def _is_matching(entry: dict, query: dict, target_set: set):
+        return is_subdict_of(query or {}, entry) and (
+            target_set <= set(entry.get("target", "").split())
+            if target_set else True)
+
     def _find(self, credential_type, target=None, query=None):  # O(n) generator
         """Returns a generator of matching entries.
 
@@ -125,6 +131,7 @@ def _find(self, credential_type, target=None, query=None):  # O(n) generator
         """
         target = sorted(target or [])  # Match the order sorted by add()
         assert isinstance(target, list), "Invalid parameter type"
+        target_set = set(target)
 
         preferred_result = None
         if (credential_type == self.CredentialType.ACCESS_TOKEN
@@ -135,20 +142,20 @@ def _find(self, credential_type, target=None, query=None):  # O(n) generator
             preferred_result = self._get_access_token(
                 query["home_account_id"], query["environment"],
                 query["client_id"], query["realm"], target)
-            if preferred_result:
+            if preferred_result and self._is_matching(
+                preferred_result, query, target_set,
+            ):
                 yield preferred_result
 
-        target_set = set(target)
         with self._lock:
             # Since the target inside token cache key is (per schema) unsorted,
             # there is no point to attempt an O(1) key-value search here.
             # So we always do an O(n) in-memory search.
             for entry in self._cache.get(credential_type, {}).values():
-                if is_subdict_of(query or {}, entry) and (
-                        target_set <= set(entry.get("target", "").split())
-                        if target else True):
-                    if entry != preferred_result:  # Avoid yielding the same entry twice
-                        yield entry
+                if (entry != preferred_result  # Avoid yielding the same entry twice
+                    and self._is_matching(entry, query, target_set)
+                ):
+                    yield entry
 
     def find(self, credential_type, target=None, query=None):  # Obsolete. Use _find() instead.
         return list(self._find(credential_type, target=target, query=query))

tests/test_e2e.py

@@ -679,11 +679,28 @@ def _test_acquire_token_by_client_secret(
 
 class PopWithExternalKeyTestCase(LabBasedTestCase):
     def _test_service_principal(self):
-        # Any SP can obtain an ssh-cert. Here we use the lab app.
-        result = get_lab_app().acquire_token_for_client(self.SCOPE, data=self.DATA1)
+        app = get_lab_app()  # Any SP can obtain an ssh-cert. Here we use the lab app.
+        result = app.acquire_token_for_client(self.SCOPE, data=self.DATA1)
         self.assertIsNotNone(result.get("access_token"), "Encountered {}: {}".format(
             result.get("error"), result.get("error_description")))
         self.assertEqual(self.EXPECTED_TOKEN_TYPE, result["token_type"])
+        self.assertEqual(result["token_source"], "identity_provider")
+
+        # Test cache hit
+        cached_result = app.acquire_token_for_client(self.SCOPE, data=self.DATA1)
+        self.assertIsNotNone(
+            cached_result.get("access_token"), "Encountered {}: {}".format(
+            cached_result.get("error"), cached_result.get("error_description")))
+        self.assertEqual(self.EXPECTED_TOKEN_TYPE, cached_result["token_type"])
+        self.assertEqual(cached_result["token_source"], "cache")
+
+        # refresh_token grant can fetch an ssh-cert bound to a different key
+        refreshed_result = app.acquire_token_for_client(self.SCOPE, data=self.DATA2)
+        self.assertIsNotNone(
+            refreshed_result.get("access_token"), "Encountered {}: {}".format(
+            refreshed_result.get("error"), refreshed_result.get("error_description")))
+        self.assertEqual(self.EXPECTED_TOKEN_TYPE, refreshed_result["token_type"])
+        self.assertEqual(refreshed_result["token_source"], "identity_provider")
 
     def _test_user_account(self):
         lab_user = self.get_lab_user(usertype="cloud")
@@ -701,16 +718,30 @@ def _test_user_account(self):
         self.assertIsNotNone(result.get("access_token"), "Encountered {}: {}".format(
             result.get("error"), result.get("error_description")))
         self.assertEqual(self.EXPECTED_TOKEN_TYPE, result["token_type"])
+        self.assertEqual(result["token_source"], "identity_provider")
         logger.debug("%s.cache = %s",
             self.id(), json.dumps(self.app.token_cache._cache, indent=4))
 
+        # refresh_token grant can hit an ssh-cert bound to the same key
+        account = self.app.get_accounts()[0]
+        cached_result = self.app.acquire_token_silent(
+            self.SCOPE, account=account, data=self.DATA1)
+        self.assertIsNotNone(cached_result)
+        self.assertEqual(self.EXPECTED_TOKEN_TYPE, cached_result["token_type"])
+        ## Actually, the self._test_acquire_token_interactive() already contained
+        ## a built-in refresh test, so the token in cache has been refreshed already.
+        ## Therefore, the following line won't pass, which is expected.
+        #self.assertEqual(result["access_token"], cached_result['access_token'])
+        self.assertEqual(cached_result["token_source"], "cache")
+
         # refresh_token grant can fetch an ssh-cert bound to a different key
         account = self.app.get_accounts()[0]
-        refreshed_ssh_cert = self.app.acquire_token_silent(
+        refreshed_result = self.app.acquire_token_silent(
             self.SCOPE, account=account, data=self.DATA2)
-        self.assertIsNotNone(refreshed_ssh_cert)
-        self.assertEqual(self.EXPECTED_TOKEN_TYPE, refreshed_ssh_cert["token_type"])
-        self.assertNotEqual(result["access_token"], refreshed_ssh_cert['access_token'])
+        self.assertIsNotNone(refreshed_result)
+        self.assertEqual(self.EXPECTED_TOKEN_TYPE, refreshed_result["token_type"])
+        self.assertNotEqual(result["access_token"], refreshed_result['access_token'])
+        self.assertEqual(refreshed_result["token_source"], "identity_provider")
 
 
 class SshCertTestCase(PopWithExternalKeyTestCase):