acquire_token_for_client() can use regional endpoint

Author: rayluo

msal/application.py

@@ -9,6 +9,7 @@
 import sys
 import warnings
 from threading import Lock
+import os
 
 import requests
 
@@ -108,14 +109,21 @@ class ClientApplication(object):
     GET_ACCOUNTS_ID = "902"
     REMOVE_ACCOUNT_ID = "903"
 
+    ATTEMPT_REGION_DISCOVERY = "TryAutoDetect"
+
     def __init__(
             self, client_id,
             client_credential=None, authority=None, validate_authority=True,
             token_cache=None,
             http_client=None,
             verify=True, proxies=None, timeout=None,
             client_claims=None, app_name=None, app_version=None,
-            client_capabilities=None):
+            client_capabilities=None,
+            region=None,  # Note: We choose to add this param in this base class,
+                # despite it is currently only needed by ConfidentialClientApplication.
+                # This way, it holds the same positional param place for PCA,
+                # when we would eventually want to add this feature to PCA in future.
+            ):
         """Create an instance of application.
 
         :param str client_id: Your app has a client_id after you register it on AAD.
@@ -220,6 +228,25 @@ def __init__(
             MSAL will combine them into
             `claims parameter <https://openid.net/specs/openid-connect-core-1_0-final.html#ClaimsParameter`_
             which you will later provide via one of the acquire-token request.
+
+        :param str region:
+            Added since MSAL Python 1.12.0.
+
+            If enabled, MSAL token requests would remain inside that region.
+            Currently, regional endpoint only supports using
+            ``acquire_token_for_client()`` for some scopes.
+
+            The default value is None, which means region support remains turned off.
+
+            App developer can opt in to regional endpoint,
+            by provide a region name, such as "westus", "eastus2".
+
+            An app running inside Azure VM can use a special keyword
+            ``ClientApplication.ATTEMPT_REGION_DISCOVERY`` to auto-detect region.
+            (Attempting this on a non-VM could hang indefinitely.
+            Make sure you configure a short timeout,
+            or provide a custom http_client which has a short timeout.
+            That way, the latency would be under your control.)
         """
         self.client_id = client_id
         self.client_credential = client_credential
@@ -249,7 +276,10 @@ def __init__(
                 self.http_client, validate_authority=validate_authority)
             # Here the self.authority is not the same type as authority in input
         self.token_cache = token_cache or TokenCache()
-        self.client = self._build_client(client_credential, self.authority)
+        self._region_configured = region
+        self._region_detected = None
+        self.client, self._regional_client = self._build_client(
+            client_credential, self.authority)
         self.authority_groups = None
         self._telemetry_buffer = {}
         self._telemetry_lock = Lock()
@@ -260,6 +290,26 @@ def _build_telemetry_context(
             self._telemetry_buffer, self._telemetry_lock, api_id,
             correlation_id=correlation_id, refresh_reason=refresh_reason)
 
+    def _detect_region(self):
+        return os.environ.get("REGION_NAME")  # TODO: or Call IMDS
+
+    def _get_regional_authority(self, central_authority):
+        self._region_detected = self._region_detected or self._detect_region()
+        if self._region_configured and self._region_detected != self._region_configured:
+            logger.warning('Region configured ({}) != region detected ({})'.format(
+                repr(self._region_configured), repr(self._region_detected)))
+        region_to_use = self._region_configured or self._region_detected
+        if region_to_use:
+            logger.info('Region to be used: {}'.format(repr(region_to_use)))
+            regional_host = ("{}.login.microsoft.com".format(region_to_use)
+                if central_authority.instance == "login.microsoftonline.com"
+                else "{}.{}".format(region_to_use, central_authority.instance))
+            return Authority(
+                "https://{}/{}".format(regional_host, central_authority.tenant),
+                self.http_client,
+                validate_authority=False)  # The central_authority has already been validated
+        return None
+
     def _build_client(self, client_credential, authority):
         client_assertion = None
         client_assertion_type = None
@@ -298,15 +348,15 @@ def _build_client(self, client_credential, authority):
             client_assertion_type = Client.CLIENT_ASSERTION_TYPE_JWT
         else:
             default_body['client_secret'] = client_credential
-        server_configuration = {
+        central_configuration = {
             "authorization_endpoint": authority.authorization_endpoint,
             "token_endpoint": authority.token_endpoint,
             "device_authorization_endpoint":
                 authority.device_authorization_endpoint or
                 urljoin(authority.token_endpoint, "devicecode"),
             }
-        return Client(
-            server_configuration,
+        central_client = Client(
+            central_configuration,
             self.client_id,
             http_client=self.http_client,
             default_headers=default_headers,
@@ -318,6 +368,30 @@ def _build_client(self, client_credential, authority):
             on_removing_rt=self.token_cache.remove_rt,
             on_updating_rt=self.token_cache.update_rt)
 
+        regional_client = None
+        regional_authority = self._get_regional_authority(authority)
+        if regional_authority:
+            regional_configuration = {
+                "authorization_endpoint": regional_authority.authorization_endpoint,
+                "token_endpoint": regional_authority.token_endpoint,
+                "device_authorization_endpoint":
+                    regional_authority.device_authorization_endpoint or
+                    urljoin(regional_authority.token_endpoint, "devicecode"),
+                }
+            regional_client = Client(
+                regional_configuration,
+                self.client_id,
+                http_client=self.http_client,
+                default_headers=default_headers,
+                default_body=default_body,
+                client_assertion=client_assertion,
+                client_assertion_type=client_assertion_type,
+                on_obtaining_tokens=lambda event: self.token_cache.add(dict(
+                    event, environment=authority.instance)),
+                on_removing_rt=self.token_cache.remove_rt,
+                on_updating_rt=self.token_cache.update_rt)
+        return central_client, regional_client
+
     def initiate_auth_code_flow(
             self,
             scopes,  # type: list[str]
@@ -953,7 +1027,7 @@ def _acquire_token_silent_by_finding_specific_refresh_token(
             # target=scopes,  # AAD RTs are scope-independent
             query=query)
         logger.debug("Found %d RTs matching %s", len(matches), query)
-        client = self._build_client(self.client_credential, authority)
+        client, _ = self._build_client(self.client_credential, authority)
 
         response = None  # A distinguishable value to mean cache is empty
         telemetry_context = self._build_telemetry_context(
@@ -1304,7 +1378,8 @@ def acquire_token_for_client(self, scopes, claims_challenge=None, **kwargs):
         self._validate_ssh_cert_input_data(kwargs.get("data", {}))
         telemetry_context = self._build_telemetry_context(
             self.ACQUIRE_TOKEN_FOR_CLIENT_ID)
-        response = _clean_up(self.client.obtain_token_for_client(
+        client = self._regional_client or self.client
+        response = _clean_up(client.obtain_token_for_client(
             scope=scopes,  # This grant flow requires no scope decoration
             headers=telemetry_context.generate_headers(),
             data=dict(

tests/test_e2e.py

@@ -99,6 +99,12 @@ def assertCacheWorksForUser(
             "We should get an AT from acquire_token_silent(...) call")
 
     def assertCacheWorksForApp(self, result_from_wire, scope):
+        logger.debug(
+            "%s: cache = %s, id_token_claims = %s",
+            self.id(),
+            json.dumps(self.app.token_cache._cache, indent=4),
+            json.dumps(result_from_wire.get("id_token_claims"), indent=4),
+            )
         # Going to test acquire_token_silent(...) to locate an AT from cache
         result_from_cache = self.app.acquire_token_silent(scope, account=None)
         self.assertIsNotNone(result_from_cache)
@@ -345,7 +351,9 @@ def test_device_flow(self):
 def get_lab_app(
         env_client_id="LAB_APP_CLIENT_ID",
         env_client_secret="LAB_APP_CLIENT_SECRET",
-        ):
+        authority="https://login.microsoftonline.com/"
+            "72f988bf-86f1-41af-91ab-2d7cd011db47",  # Microsoft tenant ID
+        **kwargs):
     """Returns the lab app as an MSAL confidential client.
 
     Get it from environment variables if defined, otherwise fall back to use MSI.
@@ -367,10 +375,12 @@ def get_lab_app(
                 env_client_id, env_client_secret)
         # See also https://microsoft.sharepoint-df.com/teams/MSIDLABSExtended/SitePages/Programmatically-accessing-LAB-API's.aspx
         raise unittest.SkipTest("MSI-based mechanism has not been implemented yet")
-    return msal.ConfidentialClientApplication(client_id, client_secret,
-            authority="https://login.microsoftonline.com/"
-                "72f988bf-86f1-41af-91ab-2d7cd011db47",  # Microsoft tenant ID
-            http_client=MinimalHttpClient())
+    return msal.ConfidentialClientApplication(
+            client_id,
+            client_credential=client_secret,
+            authority=authority,
+            http_client=MinimalHttpClient(),
+            **kwargs)
 
 def get_session(lab_app, scopes):  # BTW, this infrastructure tests the confidential client flow
     logger.info("Creating session")
@@ -726,6 +736,85 @@ def test_b2c_acquire_token_by_ropc(self):
             )
 
 
+class WorldWideRegionalEndpointTestCase(LabBasedTestCase):
+    region = "westus"
+
+    def test_acquire_token_for_client_should_hit_regional_endpoint(self):
+        """This is the only grant supported by regional endpoint, for now"""
+        self.app = get_lab_app(  # Regional endpoint only supports confidential client
+            ## Would fail the OIDC Discovery
+            #authority="https://westus2.login.microsoftonline.com/"
+            #    "72f988bf-86f1-41af-91ab-2d7cd011db47",  # Microsoft tenant ID
+
+            #authority="https://westus.login.microsoft.com/microsoft.onmicrosoft.com",
+            #validate_authority=False,
+
+            authority="https://login.microsoftonline.com/microsoft.onmicrosoft.com",
+            region=self.region,  # Explicitly use this region, regardless of detection
+            )
+        scopes = ["https://graph.microsoft.com/.default"]
+        result = self.app.acquire_token_for_client(
+            scopes,
+            params={"AllowEstsRNonMsi": "true"},  # For testing regional endpoint
+            )
+        self.assertIn('access_token', result)
+        self.assertCacheWorksForApp(result, scopes)
+        # TODO: Test the request hit the regional endpoint self.region?
+
+
+class RegionalEndpointViaEnvVarTestCase(WorldWideRegionalEndpointTestCase):
+
+    def setUp(self):
+        os.environ["REGION_NAME"] = "eastus"
+
+    def tearDown(self):
+        del os.environ["REGION_NAME"]
+
+    @unittest.skipUnless(
+        os.getenv("LAB_OBO_CLIENT_SECRET"),
+        "Need LAB_OBO_CLIENT_SECRET from https://aka.ms/GetLabSecret?Secret=TodoListServiceV2-OBO")
+    @unittest.skipUnless(
+        os.getenv("LAB_OBO_CONFIDENTIAL_CLIENT_ID"),
+        "Need LAB_OBO_CONFIDENTIAL_CLIENT_ID from https://docs.msidlab.com/flows/onbehalfofflow.html")
+    @unittest.skipUnless(
+        os.getenv("LAB_OBO_PUBLIC_CLIENT_ID"),
+        "Need LAB_OBO_PUBLIC_CLIENT_ID from https://docs.msidlab.com/flows/onbehalfofflow.html")
+    def test_cca_obo_should_bypass_regional_endpoint_therefore_still_work(self):
+        """We test OBO because it is implemented in sub class ConfidentialClientApplication"""
+        config = self.get_lab_user(usertype="cloud")
+
+        config_cca = {}
+        config_cca.update(config)
+        config_cca["client_id"] = os.getenv("LAB_OBO_CONFIDENTIAL_CLIENT_ID")
+        config_cca["scope"] = ["https://graph.microsoft.com/.default"]
+        config_cca["client_secret"] = os.getenv("LAB_OBO_CLIENT_SECRET")
+
+        config_pca = {}
+        config_pca.update(config)
+        config_pca["client_id"] = os.getenv("LAB_OBO_PUBLIC_CLIENT_ID")
+        config_pca["password"] = self.get_lab_user_secret(config_pca["lab_name"])
+        config_pca["scope"] = ["api://%s/read" % config_cca["client_id"]]
+
+        self._test_acquire_token_obo(config_pca, config_cca)
+
+    @unittest.skipUnless(
+        os.getenv("LAB_OBO_CLIENT_SECRET"),
+        "Need LAB_OBO_CLIENT_SECRET from https://aka.ms/GetLabSecret?Secret=TodoListServiceV2-OBO")
+    @unittest.skipUnless(
+        os.getenv("LAB_OBO_CONFIDENTIAL_CLIENT_ID"),
+        "Need LAB_OBO_CONFIDENTIAL_CLIENT_ID from https://docs.msidlab.com/flows/onbehalfofflow.html")
+    def test_cca_ropc_should_bypass_regional_endpoint_therefore_still_work(self):
+        """We test ROPC because it is implemented in base class ClientApplication"""
+        config = self.get_lab_user(usertype="cloud")
+        config["password"] = self.get_lab_user_secret(config["lab_name"])
+        # We repurpose the obo confidential app to test ROPC
+        # Swap in the OBO confidential app
+        config["client_id"] = os.getenv("LAB_OBO_CONFIDENTIAL_CLIENT_ID")
+        config["scope"] = ["https://graph.microsoft.com/.default"]
+        config["client_secret"] = os.getenv("LAB_OBO_CLIENT_SECRET")
+        self._test_username_password(**config)
+
+
 class ArlingtonCloudTestCase(LabBasedTestCase):
     environment = "azureusgovernment"