Add access_token_sha256_to_refresh

Author: rayluo

msal/application.py

@@ -1,11 +1,12 @@
 import functools
+import hashlib
 import json
 import time
 import logging
 import sys
 import warnings
 from threading import Lock
-from typing import Optional  # Needed in Python 3.7 & 3.8
+from typing import List, Optional  # Needed in Python 3.7 & 3.8
 from urllib.parse import urlparse
 import os
 
@@ -1520,12 +1521,17 @@ def _acquire_token_silent_from_cache_and_possibly_refresh_it(
             correlation_id=None,
             http_exceptions=None,
             auth_scheme=None,
+            *,
+            access_token_sha256_to_refresh: Optional[str] = None,
             **kwargs):
         # This internal method has two calling patterns:
         # it accepts a non-empty account to find token for a user,
         # and accepts account=None to find a token for the current app.
         access_token_from_cache = None
-        if not (force_refresh or claims_challenge or auth_scheme):  # Then attempt AT cache
+        if access_token_sha256_to_refresh and force_refresh:
+            raise ValueError(
+                "access_token_sha256_to_refresh and force_refresh are mutually exclusive")
+        if access_token_sha256_to_refresh or not (force_refresh or auth_scheme):  # Then attempt AT cache
             query={
                     "client_id": self.client_id,
                     "environment": authority.instance,
@@ -1549,6 +1555,11 @@ def _acquire_token_silent_from_cache_and_possibly_refresh_it(
                 if expires_in < 5*60:  # Then consider it expired
                     refresh_reason = msal.telemetry.AT_EXPIRED
                     continue  # Removal is not necessary, it will be overwritten
+                if access_token_sha256_to_refresh and hashlib.sha256(
+                        entry["secret"].encode()
+                        ).hexdigest() == access_token_sha256_to_refresh:
+                    refresh_reason = msal.telemetry.AT_REJECTED
+                    continue  # Might have another useful AT in next loop
                 logger.debug("Cache hit an AT")
                 access_token_from_cache = {  # Mimic a real response
                     "access_token": entry["secret"],
@@ -2347,7 +2358,14 @@ class ConfidentialClientApplication(ClientApplication):  # server-side web app
     except that ``allow_broker`` parameter shall remain ``None``.
     """
 
-    def acquire_token_for_client(self, scopes, claims_challenge=None, **kwargs):
+    def acquire_token_for_client(
+        self,
+        scopes: List[str],
+        claims_challenge: Optional[str] = None,
+        *,
+        access_token_sha256_to_refresh: Optional[str] = None,
+        **kwargs
+    ):
         """Acquires token for the current confidential client, not for an end user.
 
         Since MSAL Python 1.23, it will automatically look for token from cache,
@@ -2360,6 +2378,11 @@ def acquire_token_for_client(self, scopes, claims_challenge=None, **kwargs):
             in the form of a claims_challenge directive in the www-authenticate header to be
             returned from the UserInfo Endpoint and/or in the ID Token and/or Access Token.
             It is a string of a JSON object which contains lists of claims being requested from these locations.
+        :param access_token_sha256_to_refresh:
+            If you have an access token that is known to be rejected by the resource,
+            you can provide its sha256 here so that MSAL will refresh that specific token.
+
+            New in version 1.32.0.
 
         :return: A dict representing the json response from Microsoft Entra:
 
@@ -2371,7 +2394,9 @@ def acquire_token_for_client(self, scopes, claims_challenge=None, **kwargs):
                 "Historically, this method does not support force_refresh behavior. "
             )
         return _clean_up(self._acquire_token_silent_with_error(
-            scopes, None, claims_challenge=claims_challenge, **kwargs))
+            scopes, account=None, claims_challenge=claims_challenge,
+            access_token_sha256_to_refresh=access_token_sha256_to_refresh,
+            **kwargs))
 
     def _acquire_token_for_client(
         self,

msal/telemetry.py

@@ -13,6 +13,7 @@
 AT_EXPIRED = 3
 AT_AGING = 4
 RESERVED = 5
+AT_REJECTED = 6
 
 
 def _get_new_correlation_id():

tests/test_application.py

@@ -1,5 +1,6 @@
 # Note: Since Aug 2019 we move all e2e tests into test_e2e.py,
 # so this test_application file contains only unit tests without dependency.
+import hashlib
 import json
 import logging
 import sys
@@ -62,6 +63,35 @@ def test_bytes_to_bytes(self):
         self.assertEqual(type(_str2bytes(b"some bytes")), type(b"bytes"))
 
 
+def fake_token_getter(
+    *,
+    access_token: str = "an access token",
+    status_code: int = 200,
+    expires_in: int = 3600,
+    token_type: str = "Bearer",
+    payload: dict = None,
+    headers: dict = None,
+):
+    """A helper to create a fake token getter,
+    which will be consumed by ClientApplication's acquire methods' post parameter.
+
+    Generic mock.patch() is inconvenient because:
+    1. If you patch it at or above oauth2.py _obtain_token(), token cache is not populated.
+    2. If you patch it at request.post(), your test cases become fragile because
+       more http round-trips may be added for future flows,
+       then your existing test case would break until you mock new round-trips.
+    """
+    return lambda url, *args, **kwargs: MinimalResponse(
+        status_code=status_code,
+        text=json.dumps(payload or {
+            "access_token": access_token,
+            "expires_in": expires_in,
+            "token_type": token_type,
+            }),
+        headers=headers,
+    )
+
+
 class TestClientApplicationAcquireTokenSilentErrorBehaviors(unittest.TestCase):
 
     @patch(_OIDC_DISCOVERY, new=_OIDC_DISCOVERY_MOCK)
@@ -874,3 +904,61 @@ def test_app_did_not_register_redirect_uri_should_error_out(self):
                 parent_window_handle=app.CONSOLE_WINDOW_HANDLE,
                 )
             self.assertEqual(result.get("error"), "broker_error")
+
+
+@patch("msal.authority.tenant_discovery", new=Mock(return_value={
+    "authorization_endpoint": "https://contoso.com/placeholder",
+    "token_endpoint": "https://contoso.com/placeholder",
+    }))
+class AccessTokenToRefreshTestCase(unittest.TestCase):
+    scopes = ["scope"]
+    token1 = "AT one"
+    token1_hash = hashlib.sha256(token1.encode()).hexdigest()
+    token2 = "AT two"
+
+    def setUp(self):
+        self.app = msal.ConfidentialClientApplication("id", client_credential="*")
+        # Prepopulate cache
+        self.app.acquire_token_for_client(
+            self.scopes, post=fake_token_getter(access_token=self.token1))
+        self.assertNotEqual(
+            self.app.token_cache._cache, {}, "Cache should have been populated")
+
+    def test_mismatching_hash_should_not_trigger_refresh(self):
+        result = self.app.acquire_token_for_client(
+            self.scopes,
+            access_token_sha256_to_refresh="mismatching hash",
+            post=fake_token_getter(access_token=self.token2))
+        self.assertEqual(result.get("access_token"), self.token1, "Should hit old token")
+        self.assertEqual(result.get("token_source"), self.app._TOKEN_SOURCE_CACHE)
+
+    def test_matching_hash_should_trigger_refresh(self):
+        result = self.app.acquire_token_for_client(
+            self.scopes,
+            access_token_sha256_to_refresh=self.token1_hash,
+            post=fake_token_getter(access_token=self.token2))
+        self.assertEqual(result.get("access_token"), self.token2, "Should obtain new token")
+        self.assertEqual(result.get("token_source"), self.app._TOKEN_SOURCE_IDP)
+
+        # A client using old token1 and valid old hash, even with claims challenge,
+        # should not trigger refresh, because we can serve it with token2 in cache.
+        result = self.app.acquire_token_for_client(
+            self.scopes,
+            access_token_sha256_to_refresh=self.token1_hash,
+            claims_challenge='''{"access_token": {
+                "access_token": {"nbf": {"essential": true, "value": "1563308371"}
+                }}''',
+            post=fake_token_getter(access_token="AT three"))
+        self.assertEqual(result.get("access_token"), self.token2, "Token 2 should be returned")
+        self.assertEqual(result.get("token_source"), self.app._TOKEN_SOURCE_CACHE)
+
+    def test_force_refresh_alone_should_trigger_refresh(self):
+        # Note: MSAL Python's acquire_token_for_client() never support force_refresh,
+        # but let's ensure _acquire_token_silent_with_error(..., force_refresh=True)
+        # bypasses cache, so that other account-based flows can still use it.
+        result = self.app._acquire_token_silent_with_error(
+            self.scopes, account=None, force_refresh=True,
+            post=fake_token_getter(access_token=self.token2))
+        self.assertEqual(result.get("access_token"), self.token2, "Should hit new token")
+        self.assertEqual(result.get("token_source"), self.app._TOKEN_SOURCE_IDP)
+