Merge pull request #27 from AzureAD/foci

Family of Client Id (FoCI) support

Author: rayluo

msal/application.py

@@ -305,26 +305,71 @@ def acquire_token_silent(
                     "token_type": "Bearer",
                     "expires_in": int(expires_in),  # OAuth2 specs defines it as int
                     }
+        return self._acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
+                the_authority, decorate_scope(scopes, self.client_id), account,
+                **kwargs)
 
+    def _acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
+            self, authority, scopes, account, **kwargs):
+        query = {
+            "environment": authority.instance,
+            "home_account_id": (account or {}).get("home_account_id"),
+            # "realm": authority.tenant,  # AAD RTs are tenant-independent
+            }
+        apps = self.token_cache.find(  # Use find(), rather than token_cache.get(...)
+            TokenCache.CredentialType.APP_METADATA, query={
+                "environment": authority.instance, "client_id": self.client_id})
+        app_metadata = apps[0] if apps else {}
+        if not app_metadata:  # Meaning this app is now used for the first time.
+            # When/if we have a way to directly detect current app's family,
+            # we'll rewrite this block, to support multiple families.
+            # For now, we try existing RTs (*). If it works, we are in that family.
+            # (*) RTs of a different app/family are not supposed to be
+            # shared with or accessible by us in the first place.
+            at = self._acquire_token_silent_by_finding_specific_refresh_token(
+                authority, scopes,
+                dict(query, family_id="1"),  # A hack, we have only 1 family for now
+                rt_remover=lambda rt_item: None,  # NO-OP b/c RTs are likely not mine
+                break_condition=lambda response:  # Break loop when app not in family
+                    # Based on an AAD-only behavior mentioned in internal doc here
+                    # https://msazure.visualstudio.com/One/_git/ESTS-Docs/pullrequest/1138595
+                    "client_mismatch" in response.get("error_additional_info", []),
+                **kwargs)
+            if at:
+                return at
+        if app_metadata.get("family_id"):  # Meaning this app belongs to this family
+            at = self._acquire_token_silent_by_finding_specific_refresh_token(
+                authority, scopes, dict(query, family_id=app_metadata["family_id"]),
+                **kwargs)
+            if at:
+                return at
+        # Either this app is an orphan, so we will naturally use its own RT;
+        # or all attempts above have failed, so we fall back to non-foci behavior.
+        return self._acquire_token_silent_by_finding_specific_refresh_token(
+            authority, scopes, dict(query, client_id=self.client_id), **kwargs)
+
+    def _acquire_token_silent_by_finding_specific_refresh_token(
+            self, authority, scopes, query,
+            rt_remover=None, break_condition=lambda response: False, **kwargs):
         matches = self.token_cache.find(
             self.token_cache.CredentialType.REFRESH_TOKEN,
             # target=scopes,  # AAD RTs are scope-independent
-            query={
-                "client_id": self.client_id,
-                "environment": the_authority.instance,
-                "home_account_id": (account or {}).get("home_account_id"),
-                # "realm": the_authority.tenant,  # AAD RTs are tenant-independent
-                })
-        client = self._build_client(self.client_credential, the_authority)
+            query=query)
+        logger.debug("Found %d RTs matching %s", len(matches), query)
+        client = self._build_client(self.client_credential, authority)
         for entry in matches:
-            logger.debug("Cache hit an RT")
+            logger.debug("Cache attempts an RT")
             response = client.obtain_token_by_refresh_token(
                 entry, rt_getter=lambda token_item: token_item["secret"],
-                scope=decorate_scope(scopes, self.client_id))
+                on_removing_rt=rt_remover or self.token_cache.remove_rt,
+                scope=scopes,
+                **kwargs)
             if "error" not in response:
                 return response
             logger.debug(
                 "Refresh failed. {error}: {error_description}".format(**response))
+            if break_condition(response):
+                break
 
 
 class PublicClientApplication(ClientApplication):  # browser app or mobile app

msal/token_cache.py

@@ -30,6 +30,7 @@ class CredentialType:
         REFRESH_TOKEN = "RefreshToken"
         ACCOUNT = "Account"  # Not exactly a credential type, but we put it here
         ID_TOKEN = "IdToken"
+        APP_METADATA = "AppMetadata"
 
     class AuthorityType:
         ADFS = "ADFS"
@@ -162,6 +163,17 @@ def add(self, event, now=None):
                     rt["family_id"] = response["foci"]
                 self._cache.setdefault(self.CredentialType.REFRESH_TOKEN, {})[key] = rt
 
+            key = self._build_appmetadata_key(environment, event.get("client_id"))
+            self._cache.setdefault(self.CredentialType.APP_METADATA, {})[key] = {
+                "client_id": event.get("client_id"),
+                "environment": environment,
+                "family_id": response.get("foci"),  # None is also valid
+                }
+
+    @staticmethod
+    def _build_appmetadata_key(environment, client_id):
+        return "appmetadata-{}-{}".format(environment or "", client_id or "")
+
     @classmethod
     def _build_rt_key(
             cls,

requirements.txt

@@ -1 +1,2 @@
 .
+mock; python_version < '3.3'

tests/test_application.py

@@ -2,8 +2,15 @@
 import json
 import logging
 
+try:
+    from unittest.mock import *  # Python 3
+except:
+    from mock import *  # Need an external mock package
+
 from msal.application import *
+import msal
 from tests import unittest
+from tests.test_token_cache import TokenCacheTestCase
 
 
 THIS_FOLDER = os.path.dirname(__file__)
@@ -155,3 +162,80 @@ def test_auth_code(self):
                 error_description=result.get("error_description")))
         self.assertCacheWorks(result)
 
+
+class TestClientApplicationAcquireTokenSilentFociBehaviors(unittest.TestCase):
+
+    def setUp(self):
+        self.authority_url = "https://login.microsoftonline.com/common"
+        self.authority = msal.authority.Authority(self.authority_url)
+        self.scopes = ["s1", "s2"]
+        self.uid = "my_uid"
+        self.utid = "my_utid"
+        self.account = {"home_account_id": "{}.{}".format(self.uid, self.utid)}
+        self.frt = "what the frt"
+        self.cache = msal.SerializableTokenCache()
+        self.cache.add({  # Pre-populate a FRT
+            "client_id": "preexisting_family_app",
+            "scope": self.scopes,
+            "token_endpoint": "{}/oauth2/v2.0/token".format(self.authority_url),
+            "response": TokenCacheTestCase.build_response(
+                uid=self.uid, utid=self.utid, refresh_token=self.frt, foci="1"),
+            })  # The add(...) helper populates correct home_account_id for future searching
+
+    def test_unknown_orphan_app_will_attempt_frt_and_not_remove_it(self):
+        app = ClientApplication(
+            "unknown_orphan", authority=self.authority_url, token_cache=self.cache)
+        logger.debug("%s.cache = %s", self.id(), self.cache.serialize())
+        def tester(url, data=None, **kwargs):
+            self.assertEqual(self.frt, data.get("refresh_token"), "Should attempt the FRT")
+            return Mock(status_code=200, json=Mock(return_value={
+                "error": "invalid_grant",
+                "error_description": "Was issued to another client"}))
+        app._acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
+            self.authority, self.scopes, self.account, post=tester)
+        self.assertNotEqual([], app.token_cache.find(
+            msal.TokenCache.CredentialType.REFRESH_TOKEN, query={"secret": self.frt}),
+            "The FRT should not be removed from the cache")
+
+    def test_known_orphan_app_will_skip_frt_and_only_use_its_own_rt(self):
+        app = ClientApplication(
+            "known_orphan", authority=self.authority_url, token_cache=self.cache)
+        rt = "RT for this orphan app. We will check it being used by this test case."
+        self.cache.add({  # Populate its RT and AppMetadata, so it becomes a known orphan app
+            "client_id": app.client_id,
+            "scope": self.scopes,
+            "token_endpoint": "{}/oauth2/v2.0/token".format(self.authority_url),
+            "response": TokenCacheTestCase.build_response(
+                uid=self.uid, utid=self.utid, refresh_token=rt),
+            })
+        logger.debug("%s.cache = %s", self.id(), self.cache.serialize())
+        def tester(url, data=None, **kwargs):
+            self.assertEqual(rt, data.get("refresh_token"), "Should attempt the RT")
+            return Mock(status_code=200, json=Mock(return_value={}))
+        app._acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
+            self.authority, self.scopes, self.account, post=tester)
+
+    def test_unknown_family_app_will_attempt_frt_and_join_family(self):
+        def tester(url, data=None, **kwargs):
+            self.assertEqual(
+                self.frt, data.get("refresh_token"), "Should attempt the FRT")
+            return Mock(
+                status_code=200,
+                json=Mock(return_value=TokenCacheTestCase.build_response(
+                    uid=self.uid, utid=self.utid, foci="1", access_token="at")))
+        app = ClientApplication(
+            "unknown_family_app", authority=self.authority_url, token_cache=self.cache)
+        at = app._acquire_token_silent_by_finding_rt_belongs_to_me_or_my_family(
+            self.authority, self.scopes, self.account, post=tester)
+        logger.debug("%s.cache = %s", self.id(), self.cache.serialize())
+        self.assertEqual("at", at.get("access_token"), "New app should get a new AT")
+        app_metadata = app.token_cache.find(
+            msal.TokenCache.CredentialType.APP_METADATA,
+            query={"client_id": app.client_id})
+        self.assertNotEqual([], app_metadata, "Should record new app's metadata")
+        self.assertEqual("1", app_metadata[0].get("family_id"),
+            "The new family app should be recorded as in the same family")
+    # Known family app will simply use FRT, which is largely the same as this one
+
+    # Will not test scenario of app leaving family. Per specs, it won't happen.
+

tests/test_token_cache.py

@@ -12,30 +12,57 @@
 
 class TokenCacheTestCase(unittest.TestCase):
 
+    @staticmethod
+    def build_id_token(sub="sub", oid="oid", preferred_username="me", **kwargs):
+        return "header.%s.signature" % base64.b64encode(json.dumps(dict({
+            "sub": sub,
+            "oid": oid,
+            "preferred_username": preferred_username,
+            }, **kwargs)).encode()).decode('utf-8')
+
+    @staticmethod
+    def build_response(  # simulate a response from AAD
+            uid="uid", utid="utid",  # They will form client_info
+            access_token=None, expires_in=3600, token_type="some type",
+            refresh_token=None,
+            foci=None,
+            id_token=None,  # or something generated by build_id_token()
+            error=None,
+            ):
+        response = {
+            "client_info": base64.b64encode(json.dumps({
+                "uid": uid, "utid": utid,
+                }).encode()).decode('utf-8'),
+            }
+        if error:
+            response["error"] = error
+        if access_token:
+            response.update({
+                "access_token": access_token,
+                "expires_in": expires_in,
+                "token_type": token_type,
+                })
+        if refresh_token:
+            response["refresh_token"] = refresh_token
+        if id_token:
+            response["id_token"] = id_token
+        if foci:
+            response["foci"] = foci
+        return response
+
     def setUp(self):
         self.cache = TokenCache()
 
     def testAdd(self):
-        client_info = base64.b64encode(b'''
-            {"uid": "uid", "utid": "utid"}
-            ''').decode('utf-8')
-        id_token = "header.%s.signature" % base64.b64encode(b'''{
-            "sub": "subject",
-            "oid": "object1234",
-            "preferred_username": "John Doe"
-            }''').decode('utf-8')
+        id_token = self.build_id_token(oid="object1234", preferred_username="John Doe")
         self.cache.add({
             "client_id": "my_client_id",
             "scope": ["s2", "s1", "s3"],  # Not in particular order
             "token_endpoint": "https://login.example.com/contoso/v2/token",
-            "response": {
-		"access_token": "an access token",
-		"token_type": "some type",
-		"expires_in": 3600,
-		"refresh_token": "a refresh token",
-                "client_info": client_info,
-                "id_token": id_token,
-		},
+            "response": self.build_response(
+                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)
         self.assertEqual(
             {
@@ -88,6 +115,15 @@ def testAdd(self):
             self.cache._cache["IdToken"].get(
                 'uid.utid-login.example.com-idtoken-my_client_id-contoso-')
             )
+        self.assertEqual(
+            {
+                "client_id": "my_client_id",
+                'environment': 'login.example.com',
+                "family_id": None,
+            },
+            self.cache._cache.get("AppMetadata", {}).get(
+                "appmetadata-login.example.com-my_client_id")
+            )
 
 
 class SerializableTokenCacheTestCase(TokenCacheTestCase):