More changes and address PR comments

Author: bgavrilMS

msal/application.py

@@ -11,7 +11,12 @@
 
 from .oauth2cli import Client, JwtAssertionCreator
 from .oauth2cli.oidc import decode_part
-from .authority import Authority, WORLD_WIDE
+from .authority import (
+    Authority,
+    WORLD_WIDE,
+    _get_instance_discovery_endpoint,
+    _get_instance_discovery_host,
+)
 from .mex import send_request as mex_send_request
 from .wstrust_request import send_request as wst_send_request
 from .wstrust_response import *
@@ -671,7 +676,7 @@ def __init__(
         self._region_detected = None
         self.client, self._regional_client = self._build_client(
             client_credential, self.authority)
-        self.authority_groups = None
+        self.authority_groups = {}
         self._telemetry_buffer = {}
         self._telemetry_lock = Lock()
         _msal_extension_check()
@@ -1304,9 +1309,16 @@ def _find_msal_accounts(self, environment):
             }
         return list(grouped_accounts.values())
 
-    def _get_instance_metadata(self):  # This exists so it can be mocked in unit test
+    def _get_instance_metadata(self, instance):  # This exists so it can be mocked in unit test
+        instance_discovery_host = _get_instance_discovery_host(instance)
         resp = self.http_client.get(
-            "https://login.microsoftonline.com/common/discovery/instance?api-version=1.1&authorization_endpoint=https://login.microsoftonline.com/common/oauth2/authorize",  # TBD: We may extend this to use self._instance_discovery endpoint
+            _get_instance_discovery_endpoint(instance),
+            params={
+                'api-version': '1.1',
+                'authorization_endpoint': (
+                    "https://{}/common/oauth2/authorize".format(instance_discovery_host)
+                    ),
+            },
             headers={'Accept': 'application/json'})
         resp.raise_for_status()
         return json.loads(resp.text)['metadata']
@@ -1318,10 +1330,10 @@ def _get_authority_aliases(self, instance):
             # Then it is an ADFS/B2C/known_authority_hosts situation
             # which may not reach the central endpoint, so we skip it.
             return []
-        if not self.authority_groups:
-            self.authority_groups = [
-                set(group['aliases']) for group in self._get_instance_metadata()]
-        for group in self.authority_groups:
+        if instance not in self.authority_groups:
+            self.authority_groups[instance] = [
+                set(group['aliases']) for group in self._get_instance_metadata(instance)]
+        for group in self.authority_groups[instance]:
             if instance in group:
                 return [alias for alias in group if alias != instance]
         return []

msal/authority.py

@@ -44,6 +44,15 @@
 _CIAM_DOMAIN_SUFFIX = ".ciamlogin.com"
 
 
+def _get_instance_discovery_host(instance):
+    return instance if instance in WELL_KNOWN_AUTHORITY_HOSTS else WORLD_WIDE
+
+
+def _get_instance_discovery_endpoint(instance):
+    return 'https://{}/common/discovery/instance'.format(
+        _get_instance_discovery_host(instance))
+
+
 class AuthorityBuilder(object):
     def __init__(self, instance, tenant):
         """A helper to save caller from doing string concatenation.
@@ -152,10 +161,8 @@ def _initialize_entra_authority(
             ) or (len(parts) == 3 and parts[2].lower().startswith("b2c_"))
         self._is_known_to_developer = self.is_adfs or self._is_b2c or not validate_authority
         is_known_to_microsoft = self.instance in WELL_KNOWN_AUTHORITY_HOSTS
-        instance_discovery_host = (
-            self.instance if self.instance in WELL_KNOWN_AUTHORITY_HOSTS else WORLD_WIDE)
-        instance_discovery_endpoint = 'https://{}/common/discovery/instance'.format(  # Note: This URL seemingly returns V1 endpoint only
-            instance_discovery_host
+        instance_discovery_endpoint = _get_instance_discovery_endpoint(  # Note: This URL seemingly returns V1 endpoint only
+            self.instance
             ) if instance_discovery in (None, True) else instance_discovery
         if instance_discovery_endpoint and not (
                 is_known_to_microsoft or self._is_known_to_developer):

tests/http_client.py

@@ -40,3 +40,44 @@ def raise_for_status(self):
         if self._raw_resp is not None:  # Turns out `if requests.response` won't work
                                         # cause it would be True when 200<=status<400
             self._raw_resp.raise_for_status()
+
+
+class RecordingHttpClient(object):
+    def __init__(self):
+        self.get_calls = []
+        self.post_calls = []
+        self._get_routes = []
+        self._post_routes = []
+
+    def add_get_route(self, matcher, responder):
+        self._get_routes.append((matcher, responder))
+
+    def add_post_route(self, matcher, responder):
+        self._post_routes.append((matcher, responder))
+
+    def get(self, url, params=None, headers=None, **kwargs):
+        call = {
+            "url": url,
+            "params": params,
+            "headers": headers,
+            "kwargs": kwargs,
+        }
+        self.get_calls.append(call)
+        for matcher, responder in self._get_routes:
+            if matcher(call):
+                return responder(call)
+        return MinimalResponse(status_code=404, text="")
+
+    def post(self, url, params=None, data=None, headers=None, **kwargs):
+        call = {
+            "url": url,
+            "params": params,
+            "data": data,
+            "headers": headers,
+            "kwargs": kwargs,
+        }
+        self.post_calls.append(call)
+        for matcher, responder in self._post_routes:
+            if matcher(call):
+                return responder(call)
+        return MinimalResponse(status_code=404, text="")

tests/test_authority.py

@@ -56,16 +56,11 @@ def test_wellknown_host_and_tenant(self):
                 continue
             self._test_given_host_and_tenant(host, "common")
 
-    def test_new_sovereign_hosts_should_be_known_authorities(self):
-        self.assertIn(AZURE_GOV_FR, WELL_KNOWN_AUTHORITY_HOSTS)
-        self.assertIn(AZURE_GOV_DE, WELL_KNOWN_AUTHORITY_HOSTS)
-        self.assertIn(AZURE_GOV_SG, WELL_KNOWN_AUTHORITY_HOSTS)
-
     @patch("msal.authority._instance_discovery")
     @patch("msal.authority.tenant_discovery")
     def test_new_sovereign_hosts_should_build_authority_endpoints(
             self, tenant_discovery_mock, instance_discovery_mock):
-        for host in (AZURE_GOV_FR, AZURE_GOV_DE, AZURE_GOV_SG):
+        for host in WELL_KNOWN_AUTHORITY_HOSTS:
             tenant_discovery_mock.return_value = {
                 "authorization_endpoint": "https://{}/common/oauth2/v2.0/authorize".format(host),
                 "token_endpoint": "https://{}/common/oauth2/v2.0/token".format(host),
@@ -90,21 +85,21 @@ def test_new_sovereign_hosts_should_build_authority_endpoints(
     @patch("msal.authority.tenant_discovery")
     def test_known_authority_should_use_same_host_and_skip_instance_discovery(
             self, tenant_discovery_mock, instance_discovery_mock):
-        host = AZURE_US_GOVERNMENT
-        tenant_discovery_mock.return_value = {
-            "authorization_endpoint": "https://{}/common/oauth2/v2.0/authorize".format(host),
-            "token_endpoint": "https://{}/common/oauth2/v2.0/token".format(host),
-            "issuer": "https://{}/common/v2.0".format(host),
-        }
-        c = MinimalHttpClient()
-        Authority("https://{}/common".format(host), c)
-        c.close()
+        for host in WELL_KNOWN_AUTHORITY_HOSTS:
+            tenant_discovery_mock.return_value = {
+                "authorization_endpoint": "https://{}/common/oauth2/v2.0/authorize".format(host),
+                "token_endpoint": "https://{}/common/oauth2/v2.0/token".format(host),
+                "issuer": "https://{}/common/v2.0".format(host),
+            }
+            c = MinimalHttpClient()
+            Authority("https://{}/common".format(host), c)
+            c.close()
 
-        instance_discovery_mock.assert_not_called()
-        tenant_discovery_endpoint = tenant_discovery_mock.call_args[0][0]
-        self.assertTrue(
-            tenant_discovery_endpoint.startswith(
-                "https://{}/common/v2.0/.well-known/openid-configuration".format(host)))
+            instance_discovery_mock.assert_not_called()
+            tenant_discovery_endpoint = tenant_discovery_mock.call_args[0][0]
+            self.assertTrue(
+                tenant_discovery_endpoint.startswith(
+                    "https://{}/common/v2.0/.well-known/openid-configuration".format(host)))
 
     @patch("msal.authority._instance_discovery")
     @patch("msal.authority.tenant_discovery")
@@ -361,7 +356,24 @@ def test_by_default_a_known_to_microsoft_authority_should_skip_validation_but_st
         app = msal.ClientApplication("id", authority="https://login.microsoftonline.com/common")
         known_to_microsoft_validation.assert_not_called()
         app.get_accounts()  # This could make an instance metadata call for authority aliases
-        instance_metadata.assert_called_once_with()
+        instance_metadata.assert_called_once_with("login.microsoftonline.com")
+
+    def test_by_default_a_sovereign_known_authority_should_use_cloud_local_instance_metadata(
+            self, instance_metadata, known_to_microsoft_validation, _):
+        app = msal.ClientApplication("id", authority="https://login.microsoftonline.us/common")
+        known_to_microsoft_validation.assert_not_called()
+        app.get_accounts()  # This could make an instance metadata call for authority aliases
+        instance_metadata.assert_called_once_with("login.microsoftonline.us")
+
+    def test_fr_known_authority_should_still_work_when_instance_metadata_has_no_alias_entry(
+            self, instance_metadata, known_to_microsoft_validation, _):
+        app = msal.ClientApplication("id", authority="https://{}/common".format(AZURE_GOV_FR))
+        known_to_microsoft_validation.assert_not_called()
+
+        accounts = app.get_accounts()
+
+        self.assertEqual([], accounts)
+        instance_metadata.assert_called_once_with(AZURE_GOV_FR)
 
     def test_validate_authority_boolean_should_skip_validation_and_instance_metadata(
             self, instance_metadata, known_to_microsoft_validation, _):

tests/test_recording_http_client.py

@@ -0,0 +1,155 @@
+import json
+
+from tests import unittest
+from tests.http_client import RecordingHttpClient, MinimalResponse
+from msal.application import ConfidentialClientApplication
+
+
+class TestSovereignAuthorityForClientCredentialWithRecordingHttpClient(unittest.TestCase):
+    def test_acquire_token_for_client_on_gov_fr_should_keep_calls_on_same_host(self):
+        host = "login.sovcloud-identity.fr"
+        expected_instance_discovery_url = "https://{}/common/discovery/instance".format(host)
+        expected_instance_discovery_params = {
+            "api-version": "1.1",
+            "authorization_endpoint": (
+                "https://{}/common/oauth2/authorize".format(host)
+            ),
+        }
+
+        http_client = RecordingHttpClient()
+
+        def is_oidc_discovery(call):
+            return call["url"].startswith(
+                "https://{}/common/v2.0/.well-known/openid-configuration".format(host))
+
+        def oidc_discovery_response(_call):
+            return MinimalResponse(status_code=200, text=json.dumps({
+                "authorization_endpoint": "https://{}/common/oauth2/v2.0/authorize".format(host),
+                "token_endpoint": "https://{}/common/oauth2/v2.0/token".format(host),
+                "issuer": "https://{}/common/v2.0".format(host),
+            }))
+
+        def is_instance_discovery(call):
+            return (
+                call["url"] == expected_instance_discovery_url
+                and call["params"] == expected_instance_discovery_params
+            )
+
+        def instance_discovery_response(_call):
+            return MinimalResponse(status_code=200, text=json.dumps({
+                "tenant_discovery_endpoint": (
+                    "https://login.microsoftonline.us/"
+                    "cab8a31a-1906-4287-a0d8-4eef66b95f6e/"
+                    "v2.0/.well-known/openid-configuration"
+                ),
+                "api-version": "1.1",
+                "metadata": [
+                    {
+                        "preferred_network": "login.microsoftonline.com",
+                        "preferred_cache": "login.windows.net",
+                        "aliases": [
+                            "login.microsoftonline.com",
+                            "login.windows.net",
+                            "login.microsoft.com",
+                            "sts.windows.net",
+                        ],
+                    },
+                    {
+                        "preferred_network": "login.partner.microsoftonline.cn",
+                        "preferred_cache": "login.partner.microsoftonline.cn",
+                        "aliases": [
+                            "login.partner.microsoftonline.cn",
+                            "login.chinacloudapi.cn",
+                        ],
+                    },
+                    {
+                        "preferred_network": "login.microsoftonline.de",
+                        "preferred_cache": "login.microsoftonline.de",
+                        "aliases": ["login.microsoftonline.de"],
+                    },
+                    {
+                        "preferred_network": "login.microsoftonline.us",
+                        "preferred_cache": "login.microsoftonline.us",
+                        "aliases": [
+                            "login.microsoftonline.us",
+                            "login.usgovcloudapi.net",
+                        ],
+                    },
+                    {
+                        "preferred_network": "login-us.microsoftonline.com",
+                        "preferred_cache": "login-us.microsoftonline.com",
+                        "aliases": ["login-us.microsoftonline.com"],
+                    },
+                ],
+            }))
+
+        token_counter = {"value": 0}
+
+        def is_token_call(call):
+            return call["url"].startswith("https://{}/common/oauth2/v2.0/token".format(host))
+
+        def token_response(_call):
+            token_counter["value"] += 1
+            return MinimalResponse(status_code=200, text=json.dumps({
+                "access_token": "AT_{}".format(token_counter["value"]),
+                "expires_in": 3600,
+            }))
+
+        http_client.add_get_route(is_oidc_discovery, oidc_discovery_response)
+        http_client.add_get_route(is_instance_discovery, instance_discovery_response)
+        http_client.add_post_route(is_token_call, token_response)
+
+        app = ConfidentialClientApplication(
+            "client_id",
+            client_credential="secret",
+            authority="https://{}/common".format(host),
+            http_client=http_client,
+        )
+
+        result1 = app.acquire_token_for_client(["scope1"])
+        self.assertEqual("AT_1", result1.get("access_token"))
+
+        get_calls_after_first = list(http_client.get_calls)
+        post_calls_after_first = list(http_client.post_calls)
+
+        result2 = app.acquire_token_for_client(["scope2"])
+        self.assertEqual("AT_2", result2.get("access_token"))
+
+        post_count_after_scope2 = len(http_client.post_calls)
+        get_count_after_scope2 = len(http_client.get_calls)
+
+        cached_result1 = app.acquire_token_for_client(["scope1"])
+        self.assertEqual("AT_1", cached_result1.get("access_token"))
+
+        cached_result2 = app.acquire_token_for_client(["scope2"])
+        self.assertEqual("AT_2", cached_result2.get("access_token"))
+
+        cached_result3 = app.acquire_token_for_client(["scope1"])
+        self.assertEqual("AT_1", cached_result3.get("access_token"))
+
+        self.assertEqual(
+            post_count_after_scope2,
+            len(http_client.post_calls),
+            "Subsequent same-scope calls should be served from cache without token POST")
+        self.assertEqual(
+            get_count_after_scope2,
+            len(http_client.get_calls),
+            "Subsequent same-authority calls should not trigger additional discovery GET")
+
+        self.assertEqual(1, len(get_calls_after_first), "First acquire should trigger one discovery GET")
+        self.assertTrue(
+            get_calls_after_first[0]["url"].startswith(
+                "https://{}/common/v2.0/.well-known/openid-configuration".format(host)))
+
+        self.assertEqual(1, len(post_calls_after_first), "First acquire should trigger one token POST")
+        self.assertTrue(
+            post_calls_after_first[0]["url"].startswith("https://{}/common/oauth2/v2.0/token".format(host)))
+
+        self.assertEqual(1, len(http_client.get_calls), "Second acquire on same authority should not re-discover")
+        self.assertEqual(2, len(http_client.post_calls), "Second acquire with a different scope should request another token")
+        self.assertTrue(
+            http_client.post_calls[1]["url"].startswith("https://{}/common/oauth2/v2.0/token".format(host)))
+
+        all_urls = [c["url"] for c in http_client.get_calls + http_client.post_calls]
+        self.assertTrue(all("login.microsoftonline.com" not in url for url in all_urls))
+        self.assertTrue(all("https://{}/".format(host) in url for url in all_urls))