Add more test behavior to validate id token behaviors

Author: rayluo

msal/token_cache.py

@@ -2,20 +2,16 @@
 import threading
 import time
 import logging
-import base64
 
 from .authority import canonicalize
+from .oauth2cli.oidc import base64decode, decode_id_token
 
 
 logger = logging.getLogger(__name__)
 
 def is_subdict_of(small, big):
     return dict(big, **small) == big
 
-def base64decode(raw):  # This can handle a padding-less raw input
-    raw += '=' * (-len(raw) % 4)  # https://stackoverflow.com/a/32517907/728675
-    return base64.b64decode(raw).decode("utf-8")
-
 
 class TokenCache(object):
     """This is considered as a base class containing minimal cache behavior.
@@ -112,8 +108,8 @@ def add(self, event, now=None):
                     }
 
             if client_info:
-                decoded_id_token = json.loads(
-                    base64decode(id_token.split('.')[1])) if id_token else {}
+                decoded_id_token = decode_id_token(
+                    id_token, client_id=event["client_id"]) if id_token else {}
                 key = self._build_account_key(home_account_id, environment, realm)
                 self._cache.setdefault(self.CredentialType.ACCOUNT, {})[key] = {
                     "home_account_id": home_account_id,

tests/test_application.py

@@ -181,7 +181,7 @@ def setUp(self):
             "token_endpoint": "{}/oauth2/v2.0/token".format(self.authority_url),
             "response": TokenCacheTestCase.build_response(
                 access_token="Siblings won't share AT. test_remove_account() will.",
-                id_token=TokenCacheTestCase.build_id_token(),
+                id_token=TokenCacheTestCase.build_id_token(aud=self.preexisting_family_app_id),
                 uid=self.uid, utid=self.utid, refresh_token=self.frt, foci="1"),
             })  # The add(...) helper populates correct home_account_id for future searching
 

tests/test_token_cache.py

@@ -1,6 +1,7 @@
 import logging
 import base64
 import json
+import time
 
 from msal.token_cache import *
 from tests import unittest
@@ -13,12 +14,17 @@
 class TokenCacheTestCase(unittest.TestCase):
 
     @staticmethod
-    def build_id_token(sub="sub", oid="oid", preferred_username="me", **kwargs):
+    def build_id_token(
+            iss="issuer", sub="subject", aud="my_client_id", exp=None, iat=None,
+            preferred_username="me", **claims):
         return "header.%s.signature" % base64.b64encode(json.dumps(dict({
+            "iss": iss,
             "sub": sub,
-            "oid": oid,
+            "aud": aud,
+            "exp": exp or (time.time() + 100),
+            "iat": iat or time.time(),
             "preferred_username": preferred_username,
-            }, **kwargs)).encode()).decode('utf-8')
+            }, **claims)).encode()).decode('utf-8')
 
     @staticmethod
     def build_response(  # simulate a response from AAD
@@ -54,9 +60,11 @@ def setUp(self):
         self.cache = TokenCache()
 
     def testAdd(self):
-        id_token = self.build_id_token(oid="object1234", preferred_username="John Doe")
+        client_id = "my_client_id"
+        id_token = self.build_id_token(
+            oid="object1234", preferred_username="John Doe", aud=client_id)
         self.cache.add({
-            "client_id": "my_client_id",
+            "client_id": client_id,
             "scope": ["s2", "s1", "s3"],  # Not in particular order
             "token_endpoint": "https://login.example.com/contoso/v2/token",
             "response": self.build_response(