Expensive IMDS call

Author: rayluo

msal/application.py

@@ -9,7 +9,6 @@
 import sys
 import warnings
 from threading import Lock
-import os
 
 import requests
 
@@ -20,6 +19,7 @@
 from .wstrust_response import *
 from .token_cache import TokenCache
 import msal.telemetry
+from .region import _detect_region
 
 
 # The __init__.py will import this. Not the other way around.
@@ -261,7 +261,7 @@ def __init__(
             # Requests, does not support session - wide timeout
             # But you can patch that (https://github.com/psf/requests/issues/3341):
             self.http_client.request = functools.partial(
-                self.http_client.request, timeout=timeout)
+                self.http_client.request, timeout=timeout or 2)
 
             # Enable a minimal retry. Better than nothing.
             # https://github.com/psf/requests/blob/v2.25.1/requests/adapters.py#L94-L108
@@ -290,11 +290,8 @@ 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()
+        self._region_detected = self._region_detected or _detect_region(self.http_client)
         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)))
@@ -369,27 +366,28 @@ def _build_client(self, client_credential, authority):
             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)
+        if client_credential:  # Currently regional endpoint only serves some CCA flows
+            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(

msal/region.py

@@ -0,0 +1,26 @@
+import os
+import json
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def _detect_region(http_client):
+    return _detect_region_of_azure_function() or _detect_region_of_azure_vm(http_client)
+
+
+def _detect_region_of_azure_function():
+    return os.environ.get("REGION_NAME")
+
+
+def _detect_region_of_azure_vm(http_client):
+    url = "http://169.254.169.254/metadata/instance?api-version=2021-01-01"
+    try:
+        # https://docs.microsoft.com/en-us/azure/virtual-machines/windows/instance-metadata-service?tabs=linux#instance-metadata
+        resp = http_client.get(url, headers={"Metadata": "true"})
+    except:
+        logger.info("IMDS {} unavailable. Perhaps not running in Azure VM?".format(url))
+        return None
+    else:
+        return json.loads(resp.text)["compute"]["location"]
+

tests/test_e2e.py

@@ -27,7 +27,7 @@ def _get_app_and_auth_code(
         app = msal.ConfidentialClientApplication(
             client_id,
             client_credential=client_secret,
-            authority=authority, http_client=MinimalHttpClient())
+            authority=authority, http_client=MinimalHttpClient(timeout=2))
     else:
         app = msal.PublicClientApplication(
             client_id, authority=authority, http_client=MinimalHttpClient())
@@ -292,7 +292,7 @@ def test_client_secret(self):
             self.config["client_id"],
             client_credential=self.config.get("client_secret"),
             authority=self.config.get("authority"),
-            http_client=MinimalHttpClient())
+            http_client=MinimalHttpClient(timeout=2))
         scope = self.config.get("scope", [])
         result = self.app.acquire_token_for_client(scope)
         self.assertIn('access_token', result)
@@ -307,7 +307,7 @@ def test_client_certificate(self):
         self.app = msal.ConfidentialClientApplication(
             self.config['client_id'],
             {"private_key": private_key, "thumbprint": client_cert["thumbprint"]},
-            http_client=MinimalHttpClient())
+            http_client=MinimalHttpClient(timeout=2))
         scope = self.config.get("scope", [])
         result = self.app.acquire_token_for_client(scope)
         self.assertIn('access_token', result)
@@ -330,7 +330,7 @@ def test_subject_name_issuer_authentication(self):
                 "thumbprint": self.config["thumbprint"],
                 "public_certificate": public_certificate,
                 },
-            http_client=MinimalHttpClient())
+            http_client=MinimalHttpClient(timeout=2))
         scope = self.config.get("scope", [])
         result = self.app.acquire_token_for_client(scope)
         self.assertIn('access_token', result)
@@ -379,14 +379,17 @@ def get_lab_app(
             client_id,
             client_credential=client_secret,
             authority=authority,
-            http_client=MinimalHttpClient(),
+            http_client=MinimalHttpClient(timeout=2),
             **kwargs)
 
 def get_session(lab_app, scopes):  # BTW, this infrastructure tests the confidential client flow
     logger.info("Creating session")
-    lab_token = lab_app.acquire_token_for_client(scopes)
+    result = lab_app.acquire_token_for_client(scopes)
+    assert result.get("access_token"), \
+        "Unable to obtain token for lab. Encountered {}: {}".format(
+            result.get("error"), result.get("error_description"))
     session = requests.Session()
-    session.headers.update({"Authorization": "Bearer %s" % lab_token["access_token"]})
+    session.headers.update({"Authorization": "Bearer %s" % result["access_token"]})
     session.hooks["response"].append(lambda r, *args, **kwargs: r.raise_for_status())
     return session
 
@@ -525,7 +528,7 @@ def _test_acquire_token_obo(self, config_pca, config_cca):
             config_cca["client_id"],
             client_credential=config_cca["client_secret"],
             authority=config_cca["authority"],
-            http_client=MinimalHttpClient(),
+            http_client=MinimalHttpClient(timeout=2),
             # token_cache= ...,  # Default token cache is all-tokens-store-in-memory.
                 # That's fine if OBO app uses short-lived msal instance per session.
                 # Otherwise, the OBO app need to implement a one-cache-per-user setup.
@@ -553,7 +556,7 @@ def _test_acquire_token_by_client_secret(
         assert client_id and client_secret and authority and scope
         app = msal.ConfidentialClientApplication(
             client_id, client_credential=client_secret, authority=authority,
-            http_client=MinimalHttpClient())
+            http_client=MinimalHttpClient(timeout=2))
         result = app.acquire_token_for_client(scope)
         self.assertIsNotNone(result.get("access_token"), "Got %s instead" % result)
 
@@ -852,7 +855,7 @@ def test_acquire_token_silent_with_an_empty_cache_should_return_none(self):
             usertype="cloud", azureenvironment=self.environment, publicClient="no")
         app = msal.ConfidentialClientApplication(
             config['client_id'], authority=config['authority'],
-            http_client=MinimalHttpClient())
+            http_client=MinimalHttpClient(timeout=2))
         result = app.acquire_token_silent(scopes=config['scope'], account=None)
         self.assertEqual(result, None)
         # Note: An alias in this region is no longer accepting HTTPS traffic.