WIP: ManagedIdentityClient sends claims and token_sha256_to_refresh to SF

TODO: Need to change it to accept client_capabilities and the relay it
to SF via xms_cc

Author: rayluo

msal/managed_identity.py

@@ -2,6 +2,7 @@
 # All rights reserved.
 #
 # This code is licensed under the MIT License.
+import hashlib
 import json
 import logging
 import os
@@ -266,8 +267,7 @@ def acquire_token_for_client(
             and then a *claims challenge* will be returned by the target resource,
             as a `claims_challenge` directive in the `www-authenticate` header,
             even if the app developer did not opt in for the "CP1" client capability.
-            Upon receiving a `claims_challenge`, MSAL will skip a token cache read,
-            and will attempt to acquire a new token.
+            Upon receiving a `claims_challenge`, MSAL will attempt to acquire a new token.
 
         .. note::
 
@@ -278,11 +278,13 @@ def acquire_token_for_client(
             This is a service-side behavior that cannot be changed by this library.
             `Azure VM docs <https://learn.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token#get-a-token-using-http>`_
         """
+        access_token_to_refresh = None  # This could become a public parameter in the future
         access_token_from_cache = None
         client_id_in_cache = self._managed_identity.get(
             ManagedIdentity.ID, "SYSTEM_ASSIGNED_MANAGED_IDENTITY")
         now = time.time()
-        if not claims_challenge:  # Then attempt token cache search
+        if True:  # Attempt cache search even if receiving claims_challenge,
+                  # because we want to locate the existing token (if any) and refresh it
             matches = self._token_cache.find(
                 self._token_cache.CredentialType.ACCESS_TOKEN,
                 target=[resource],
@@ -297,6 +299,11 @@ def acquire_token_for_client(
                 expires_in = int(entry["expires_on"]) - now
                 if expires_in < 5*60:  # Then consider it expired
                     continue  # Removal is not necessary, it will be overwritten
+                if claims_challenge and not access_token_to_refresh:
+                    # Since caller did not pinpoint the token causing claims challenge,
+                    # we have to assume it is the first token we found in cache.
+                    access_token_to_refresh = entry["secret"]
+                    break
                 logger.debug("Cache hit an AT")
                 access_token_from_cache = {  # Mimic a real response
                     "access_token": entry["secret"],
@@ -310,7 +317,13 @@ def acquire_token_for_client(
                         break  # With a fallback in hand, we break here to go refresh
                 return access_token_from_cache  # It is still good as new
         try:
-            result = _obtain_token(self._http_client, self._managed_identity, resource)
+            result = _obtain_token(
+                self._http_client, self._managed_identity, resource,
+                claims_challenge=claims_challenge,
+                access_token_sha256_to_refresh=hashlib.sha256(
+                    access_token_to_refresh.encode("utf-8")).hexdigest()
+                    if access_token_to_refresh else None,
+            )
             if "access_token" in result:
                 expires_in = result.get("expires_in", 3600)
                 if "refresh_in" not in result and expires_in >= 7200:
@@ -385,8 +398,14 @@ def get_managed_identity_source():
     return DEFAULT_TO_VM
 
 
-def _obtain_token(http_client, managed_identity, resource):
-    # A unified low-level API that talks to different Managed Identity
+def _obtain_token(
+    http_client, managed_identity, resource,
+    *,
+    claims_challenge: Optional[str] = None,
+    access_token_sha256_to_refresh: Optional[str] = None,
+):
+    if claims_challenge and len(claims_challenge) > 1024:
+        claims_challenge = None  # MSIv1 does not support long url
     if ("IDENTITY_ENDPOINT" in os.environ and "IDENTITY_HEADER" in os.environ
             and "IDENTITY_SERVER_THUMBPRINT" in os.environ
     ):
@@ -402,6 +421,8 @@ def _obtain_token(http_client, managed_identity, resource):
             os.environ["IDENTITY_HEADER"],
             os.environ["IDENTITY_SERVER_THUMBPRINT"],
             resource,
+            claims_challenge=claims_challenge,
+            access_token_sha256_to_refresh=access_token_sha256_to_refresh,
         )
     if "IDENTITY_ENDPOINT" in os.environ and "IDENTITY_HEADER" in os.environ:
         return _obtain_token_on_app_service(
@@ -553,6 +574,9 @@ def _obtain_token_on_machine_learning(
 
 def _obtain_token_on_service_fabric(
     http_client, endpoint, identity_header, server_thumbprint, resource,
+    *,
+    claims_challenge: str = None,
+    access_token_sha256_to_refresh: str = None,
 ):
     """Obtains token for
     `Service Fabric <https://learn.microsoft.com/en-us/azure/service-fabric/>`_
@@ -563,7 +587,12 @@ def _obtain_token_on_service_fabric(
     logger.debug("Obtaining token via managed identity on Azure Service Fabric")
     resp = http_client.get(
         endpoint,
-        params={"api-version": "2019-07-01-preview", "resource": resource},
+        params={k: v for k, v in {
+            "api-version": "2019-07-01-preview",
+            "resource": resource,
+            "claims": claims_challenge,
+            "token_sha256_to_refresh": access_token_sha256_to_refresh,
+            }.items() if v is not None},
         headers={"Secret": identity_header},
         )
     try:

tests/test_mi.py

@@ -1,3 +1,4 @@
+import hashlib
 import json
 import os
 import sys
@@ -79,7 +80,13 @@ def assertCacheStatus(self, app):
             "Should have expected client_id")
         self.assertEqual("managed_identity", at["realm"], "Should have expected realm")
 
-    def _test_happy_path(self, app, mocked_http, expires_in, resource="R"):
+    def _test_happy_path(
+        self, app, mocked_http, expires_in, *, resource="R", claims_challenge=None,
+    ):
+        """It tests a normal token request that is expected to hit IdP,
+        a subsequent same token request that is expected to hit cache,
+        and then a request with claims_challenge that shall hit IdP again.
+        """
         result = app.acquire_token_for_client(resource=resource)
         mocked_http.assert_called()
         call_count = mocked_http.call_count
@@ -115,7 +122,8 @@ def _test_happy_path(self, app, mocked_http, expires_in, resource="R"):
                 expected_refresh_on - 5 < result["refresh_on"] <= expected_refresh_on,
                 "Should have a refresh_on time around the middle of the token's life")
 
-        result = app.acquire_token_for_client(resource=resource, claims_challenge="foo")
+        result = app.acquire_token_for_client(
+            resource=resource, claims_challenge=claims_challenge or "placeholder")
         self.assertEqual("identity_provider", result["token_source"], "Should miss cache")
 
 
@@ -132,6 +140,9 @@ def _test_happy_path(self) -> callable:
 
     def test_happy_path_of_vm(self):
         self._test_happy_path().assert_called_with(
+            # The last call contained claims_challenge
+            # but since IMDS doesn't support token_sha256_to_refresh,
+            # the request shall remain the same as before
             'http://169.254.169.254/metadata/identity/oauth2/token',
             params={'api-version': '2018-02-01', 'resource': 'R'},
             headers={'Metadata': 'true'},
@@ -245,18 +256,52 @@ def test_machine_learning_error_should_be_normalized(self):
 })
 class ServiceFabricTestCase(ClientTestCase):
 
-    def _test_happy_path(self, app):
+    def _test_happy_path(self, app, *, claims_challenge=None) -> callable:
         expires_in = 1234
         with patch.object(app._http_client, "get", return_value=MinimalResponse(
             status_code=200,
             text='{"access_token": "AT", "expires_on": %s, "resource": "R", "token_type": "Bearer"}' % (
                 int(time.time()) + expires_in),
         )) as mocked_method:
             super(ServiceFabricTestCase, self)._test_happy_path(
-                app, mocked_method, expires_in)
+                app, mocked_method, expires_in, claims_challenge=claims_challenge)
+            return mocked_method
 
-    def test_happy_path(self):
-        self._test_happy_path(self.app)
+    def test_happy_path_with_small_claim_challenge(self):
+        claims_challenge='{"access_token": {"nbf": {"essential": true, "value": "1563308371"}}}'
+        last_call = self._test_happy_path(self.app, claims_challenge=claims_challenge)
+        last_call.assert_called_with(
+            # The last call contained claims_challenge
+            # and the claim_challenge has a size less than 1KB,
+            # so it should relay both claims and hash to SF
+            'http://localhost',
+            params={
+                'api-version': '2019-07-01-preview',
+                'resource': 'R',
+                'claims': claims_challenge,
+                'token_sha256_to_refresh': hashlib.sha256(b"AT").hexdigest(),
+            },
+            headers={'Secret': 'foo'}
+        )
+
+    def test_happy_path_with_large_claim_challenge(self):
+        claims_challenge=json.dumps({
+            "filler": "x" * 1024,  # 1KB payload
+            "access_token": {"nbf": {"essential": True, "value": "1563308371"}},
+            })
+        last_call = self._test_happy_path(self.app, claims_challenge=claims_challenge)
+        last_call.assert_called_with(
+            # The last call contained claims_challenge
+            # and the claim_challenge has a size more than 1KB,
+            # so it should relay token_sha256_to_refresh only to SF
+            'http://localhost',
+            params={
+                'api-version': '2019-07-01-preview',
+                'resource': 'R',
+                'token_sha256_to_refresh': hashlib.sha256(b"AT").hexdigest(),
+            },
+            headers={'Secret': 'foo'}
+        )
 
     def test_unified_api_service_should_ignore_unnecessary_client_id(self):
         self._test_happy_path(ManagedIdentityClient(