Implementing CCS Routing info

Author: rayluo

msal/application.py

@@ -590,6 +590,8 @@ def initiate_auth_code_flow(
                 self._client_capabilities, claims_challenge),
             )
         flow["claims_challenge"] = claims_challenge
+        if login_hint:
+            flow["login_hint"] = login_hint  # To be relayed to token endpoint
         return flow
 
     def get_authorization_request_url(
@@ -726,11 +728,15 @@ def authorize():  # A controller in a web app
         self._validate_ssh_cert_input_data(kwargs.get("data", {}))
         telemetry_context = self._build_telemetry_context(
             self.ACQUIRE_TOKEN_BY_AUTHORIZATION_CODE_ID)
+        headers = telemetry_context.generate_headers()
+        if "login_hint" in auth_code_flow:  # Then use it as the CCS Routing info
+            headers["X-AnchorMailbox"] = "UPN:{}".format(
+                auth_code_flow.pop("login_hint"))
         response =_clean_up(self.client.obtain_token_by_auth_code_flow(
             auth_code_flow,
             auth_response,
             scope=self._decorate_scope(scopes) if scopes else None,
-            headers=telemetry_context.generate_headers(),
+            headers=headers,
             data=dict(
                 kwargs.pop("data", {}),
                 claims=_merge_claims_challenge_and_capabilities(
@@ -1178,6 +1184,10 @@ def _acquire_token_silent_by_finding_specific_refresh_token(
                 key=lambda e: int(e.get("last_modification_time", "0")),
                 reverse=True):
             logger.debug("Cache attempts an RT")
+            headers = telemetry_context.generate_headers()
+            if "home_account_id" in query:  # Then use it as CCS Routing info
+                headers["X-AnchorMailbox"] = "Oid:{}".format(
+                    query["home_account_id"].replace(".", "@"))
             response = client.obtain_token_by_refresh_token(
                 entry, rt_getter=lambda token_item: token_item["secret"],
                 on_removing_rt=lambda rt_item: None,  # Disable RT removal,
@@ -1189,7 +1199,7 @@ def _acquire_token_silent_by_finding_specific_refresh_token(
                     skip_account_creation=True,  # To honor a concurrent remove_account()
                     )),
                 scope=scopes,
-                headers=telemetry_context.generate_headers(),
+                headers=headers,
                 data=dict(
                     kwargs.pop("data", {}),
                     claims=_merge_claims_challenge_and_capabilities(
@@ -1284,6 +1294,8 @@ def acquire_token_by_username_password(
         telemetry_context = self._build_telemetry_context(
             self.ACQUIRE_TOKEN_BY_USERNAME_PASSWORD_ID)
         headers = telemetry_context.generate_headers()
+            # No need to add CCS Routing info,
+            # because username param will be recognized as CCS Routing info.
         data = dict(
             kwargs.pop("data", {}),
             claims=_merge_claims_challenge_and_capabilities(
@@ -1425,6 +1437,9 @@ def acquire_token_interactive(
             self._client_capabilities, claims_challenge)
         telemetry_context = self._build_telemetry_context(
             self.ACQUIRE_TOKEN_INTERACTIVE)
+        headers = telemetry_context.generate_headers()
+        if login_hint:  # Then use it as the CCS Routing info
+            headers["X-AnchorMailbox"] = "UPN:{}".format(login_hint)
         response = _clean_up(self.client.obtain_token_by_browser(
             scope=self._decorate_scope(scopes) if scopes else None,
             extra_scope_to_consent=extra_scopes_to_consent,
@@ -1439,7 +1454,7 @@ def acquire_token_interactive(
                 "domain_hint": domain_hint,
                 },
             data=dict(kwargs.pop("data", {}), claims=claims),
-            headers=telemetry_context.generate_headers(),
+            headers=headers,
             browser_name=_preferred_browser(),
             **kwargs))
         telemetry_context.update_telemetry(response)

tests/test_ccs.py

@@ -0,0 +1,53 @@
+import unittest
+try:
+    from unittest.mock import patch, ANY
+except:
+    from mock import patch, ANY
+
+from tests.http_client import MinimalResponse
+from tests.test_token_cache import build_response
+
+import msal
+
+
+class TestCcsRoutingInfoTestCase(unittest.TestCase):
+
+    def test_acquire_token_by_auth_code_flow(self):
+        app = msal.ClientApplication("client_id")
+        state = "foo"
+        flow = app.initiate_auth_code_flow(
+            ["some", "scope"], login_hint="[email protected]", state=state)
+        with patch.object(app.http_client, "post", return_value=MinimalResponse(
+                status_code=400, text='{"error": "mock"}')) as mocked_method:
+            app.acquire_token_by_auth_code_flow(flow, {"state": state, "code": "bar"})
+            self.assertEqual(
+                "UPN:[email protected]",
+                mocked_method.call_args[1].get("headers", {}).get('X-AnchorMailbox'),
+                "CSS routing info should be derived from login_hint")
+
+    def test_acquire_token_silent(self):
+        uid = "foo"
+        utid = "bar"
+        client_id = "my_client_id"
+        scopes = ["some", "scope"]
+        authority_url = "https://login.microsoftonline.com/common"
+        token_cache = msal.TokenCache()
+        token_cache.add({  # Pre-populate the cache
+            "client_id": client_id,
+            "scope": scopes,
+            "token_endpoint": "{}/oauth2/v2.0/token".format(authority_url),
+            "response": build_response(
+                access_token="an expired AT to trigger refresh", expires_in=-99,
+                uid=uid, utid=utid, refresh_token="this is a RT"),
+            })  # The add(...) helper populates correct home_account_id for future searching
+        app = msal.ClientApplication(
+            client_id, authority=authority_url, token_cache=token_cache)
+        with patch.object(app.http_client, "post", return_value=MinimalResponse(
+                status_code=400, text='{"error": "mock"}')) as mocked_method:
+            account = {"home_account_id": "{}.{}".format(uid, utid)}
+            app.acquire_token_silent(["scope"], account)
+            self.assertEqual(
+                "Oid:{}@{}".format(uid, utid),  # It would look like "Oid:foo@bar"
+                mocked_method.call_args[1].get("headers", {}).get('X-AnchorMailbox'),
+                "CSS routing info should be derived from home_account_id")
+

tests/test_e2e.py

@@ -516,8 +516,8 @@ def _test_acquire_token_by_auth_code_flow(
             client_id, authority=authority, http_client=MinimalHttpClient())
         with AuthCodeReceiver(port=port) as receiver:
             flow = self.app.initiate_auth_code_flow(
+                scope,
                 redirect_uri="http://localhost:%d" % receiver.get_port(),
-                scopes=scope,
                 )
             auth_response = receiver.get_auth_response(
                 auth_uri=flow["auth_uri"], state=flow["state"], timeout=60,