Don't cache tokens rejected by the auth code nonce or max_age check

msal/oauth2cli/oidc.py

@@ -241,13 +241,18 @@ def obtain_token_by_authorization_code(self, code, nonce=None, **kwargs):
         """
         warnings.warn(
             "Use obtain_token_by_auth_code_flow() instead", DeprecationWarning)
+        on_obtaining_tokens = (
+            kwargs.pop("on_obtaining_tokens", None) or self.on_obtaining_tokens)
+        events = []  # Do not store any token until the nonce is validated
         result = super(Client, self).obtain_token_by_authorization_code(
-            code, **kwargs)
+            code, on_obtaining_tokens=events.append, **kwargs)
         nonce_in_id_token = result.get("id_token_claims", {}).get("nonce")
         if "id_token_claims" in result and nonce and nonce != nonce_in_id_token:
             raise ValueError(
                 'The nonce in id token ("%s") should match your nonce ("%s")' %
                 (nonce_in_id_token, nonce))
+        for event in events:
+            on_obtaining_tokens(event)
         return result
 
     def initiate_auth_code_flow(
@@ -297,8 +302,12 @@ def obtain_token_by_auth_code_flow(self, auth_code_flow, auth_response, **kwargs
         See :func:`oauth2.Client.obtain_token_by_auth_code_flow` in parent class
         for descriptions on other parameters and return value.
         """
+        on_obtaining_tokens = (
+            kwargs.pop("on_obtaining_tokens", None) or self.on_obtaining_tokens)
+        events = []  # Do not store any token until the ID token passes the checks below
         result = super(Client, self).obtain_token_by_auth_code_flow(
-            auth_code_flow, auth_response, **kwargs)
+            auth_code_flow, auth_response,
+            on_obtaining_tokens=events.append, **kwargs)
         if "id_token_claims" in result:
             nonce_in_id_token = result.get("id_token_claims", {}).get("nonce")
             expected_hash = _nonce_hash(auth_code_flow["nonce"])
@@ -326,6 +335,8 @@ def obtain_token_by_auth_code_flow(self, auth_code_flow, auth_response, **kwargs
                         now=now,
                         id_token=json.dumps(result["id_token_claims"], indent=2),
                         ))
+        for event in events:
+            on_obtaining_tokens(event)
         return result
 
     def obtain_token_by_browser(

tests/test_application.py

@@ -15,6 +15,7 @@
     ClientApplication, PublicClientApplication, ConfidentialClientApplication,
     _str2bytes, _merge_claims_challenge_and_capabilities,
 )
+from msal.oauth2cli.oidc import _nonce_hash
 from tests import unittest
 from tests.test_token_cache import build_id_token, build_response
 from tests.http_client import MinimalHttpClient, MinimalResponse
@@ -1461,6 +1462,73 @@ def test_cached_token_is_isolated_by_client_claims(self):
             "A plain silent call must not read a client_claims token")
 
 
+@patch(_OIDC_DISCOVERY, new=_OIDC_DISCOVERY_MOCK)
+class TestAuthCodeTokensRejectedByIdTokenChecksAreNotCached(unittest.TestCase):
+    """Tokens from a response that fails the nonce or max_age check
+    must not be stored, otherwise acquire_token_silent() would serve them."""
+
+    def _build_app(self):
+        return ConfidentialClientApplication(
+            "client_id", client_credential="secret",
+            authority="https://login.microsoftonline.com/my_tenant")
+
+    def _post(self, **id_token_claims):
+        return lambda *a, **k: MinimalResponse(status_code=200, text=json.dumps(
+            build_response(
+                uid="user_oid", utid="my_tenant", access_token="an AT",
+                refresh_token="an RT", id_token=build_id_token(
+                    aud="client_id", oid="user_oid", tid="my_tenant",
+                    preferred_username="user@contoso.com", **id_token_claims))))
+
+    def _assert_nothing_cached(self, app):
+        self.assertEqual([], app.get_accounts())
+        for credential_type in (
+                msal.TokenCache.CredentialType.ACCESS_TOKEN,
+                msal.TokenCache.CredentialType.REFRESH_TOKEN):
+            self.assertEqual([], list(app.token_cache.search(credential_type)))
+
+    def test_tokens_with_mismatching_nonce_should_not_be_cached(self):
+        app = self._build_app()
+        flow = app.initiate_auth_code_flow(["s"])
+        with self.assertRaises(RuntimeError):
+            app.acquire_token_by_auth_code_flow(
+                flow, {"state": flow["state"], "code": "code"},
+                post=self._post(nonce="a different nonce"))
+        self._assert_nothing_cached(app)
+
+    def test_tokens_with_stale_auth_time_should_not_be_cached(self):
+        app = self._build_app()
+        flow = app.initiate_auth_code_flow(["s"], max_age=60)
+        with self.assertRaises(RuntimeError):
+            app.acquire_token_by_auth_code_flow(
+                flow, {"state": flow["state"], "code": "code"},
+                post=self._post(
+                    nonce=_nonce_hash(flow["nonce"]),
+                    auth_time=int(time.time()) - 3600))
+        self._assert_nothing_cached(app)
+
+    def test_tokens_with_mismatching_nonce_via_authorization_code_should_not_be_cached(self):
+        app = self._build_app()
+        with self.assertRaises(ValueError):
+            app.acquire_token_by_authorization_code(
+                "code", ["s"], nonce="expected nonce",
+                post=self._post(nonce="a different nonce"))
+        self._assert_nothing_cached(app)
+
+    def test_tokens_passing_id_token_checks_should_be_cached(self):
+        app = self._build_app()
+        flow = app.initiate_auth_code_flow(["s"], max_age=60)
+        app.acquire_token_by_auth_code_flow(
+            flow, {"state": flow["state"], "code": "code"},
+            post=self._post(
+                nonce=_nonce_hash(flow["nonce"]), auth_time=int(time.time())))
+        accounts = app.get_accounts()
+        self.assertEqual(1, len(accounts))
+        result = app.acquire_token_silent(["s"], accounts[0])
+        self.assertEqual("an AT", result["access_token"])
+        self.assertEqual(result[app._TOKEN_SOURCE], app._TOKEN_SOURCE_CACHE)
+
+
 @patch(_OIDC_DISCOVERY, new=_OIDC_DISCOVERY_MOCK)
 class TestUserFicWithClientClaims(unittest.TestCase):
     """acquire_token_by_user_federated_identity_credential(forwarded_client_claims=...)