From b1b5e94c81387d4d2ffcbe1845931bbf1395a88b Mon Sep 17 00:00:00 2001
From: Ray Luo <rayluo@microsoft.com>
Date: Wed, 4 Dec 2024 12:34:56 -0800
Subject: [PATCH] Adds dSTS authority (as if it were an oidc_authority)

---
 msal/application.py     |  4 ++++
 tests/test_authority.py | 51 +++++++++++++++++++++++++++++++++--------
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/msal/application.py b/msal/application.py
index e46d7484..89b9b91a 100644
--- a/msal/application.py
+++ b/msal/application.py
@@ -6,6 +6,7 @@
 import warnings
 from threading import Lock
 from typing import Optional  # Needed in Python 3.7 & 3.8
+from urllib.parse import urlparse
 import os
 
 from .oauth2cli import Client, JwtAssertionCreator
@@ -623,6 +624,9 @@ def __init__(
         # Here the self.authority will not be the same type as authority in input
         if oidc_authority and authority:
             raise ValueError("You can not provide both authority and oidc_authority")
+        if isinstance(authority, str) and urlparse(authority).path.startswith(
+            "/dstsv2"):  # dSTS authority's path always starts with "/dstsv2"
+            oidc_authority = authority  # So we treat it as if an oidc_authority
         try:
             authority_to_use = authority or "https://{}/common/".format(WORLD_WIDE)
             self.authority = Authority(
diff --git a/tests/test_authority.py b/tests/test_authority.py
index 0d6c790f..3fd1fce1 100644
--- a/tests/test_authority.py
+++ b/tests/test_authority.py
@@ -104,32 +104,63 @@ def test_authority_with_path_should_be_used_as_is(self, oidc_discovery):
     "authorization_endpoint": "https://contoso.com/authorize",
     "token_endpoint": "https://contoso.com/token",
     })
-class TestOidcAuthority(unittest.TestCase):
+class OidcAuthorityTestCase(unittest.TestCase):
+    authority = "https://contoso.com/tenant"
+
+    def setUp(self):
+        # setUp() gives subclass a dynamic setup based on their authority
+        self.oidc_discovery_endpoint = (
+            # MSAL Python always does OIDC Discovery,
+            # not to be confused with Instance Discovery
+            # Here the test is to confirm the OIDC endpoint contains no "/v2.0"
+            self.authority + "/.well-known/openid-configuration")
+
     def test_authority_obj_should_do_oidc_discovery_and_skip_instance_discovery(
             self, oidc_discovery, instance_discovery):
         c = MinimalHttpClient()
-        a = Authority(None, c, oidc_authority_url="https://contoso.com/tenant")
+        a = Authority(None, c, oidc_authority_url=self.authority)
         instance_discovery.assert_not_called()
-        oidc_discovery.assert_called_once_with(
-            "https://contoso.com/tenant/.well-known/openid-configuration", c)
+        oidc_discovery.assert_called_once_with(self.oidc_discovery_endpoint, c)
         self.assertEqual(a.authorization_endpoint, 'https://contoso.com/authorize')
         self.assertEqual(a.token_endpoint, 'https://contoso.com/token')
 
     def test_application_obj_should_do_oidc_discovery_and_skip_instance_discovery(
             self, oidc_discovery, instance_discovery):
         app = msal.ClientApplication(
-            "id",
-            authority=None,
-            oidc_authority="https://contoso.com/tenant",
-            )
+            "id", authority=None, oidc_authority=self.authority)
         instance_discovery.assert_not_called()
         oidc_discovery.assert_called_once_with(
-            "https://contoso.com/tenant/.well-known/openid-configuration",
-            app.http_client)
+            self.oidc_discovery_endpoint, app.http_client)
         self.assertEqual(
             app.authority.authorization_endpoint, 'https://contoso.com/authorize')
         self.assertEqual(app.authority.token_endpoint, 'https://contoso.com/token')
 
+
+class DstsAuthorityTestCase(OidcAuthorityTestCase):
+    # Inherits OidcAuthority's test cases and run them with a dSTS authority
+    authority = (  # dSTS is single tenanted with a tenant placeholder
+        'https://test-instance1-dsts.dsts.core.azure-test.net/dstsv2/common')
+    authorization_endpoint = (
+        "https://some.url.dsts.core.azure-test.net/dstsv2/common/oauth2/authorize")
+    token_endpoint = (
+        "https://some.url.dsts.core.azure-test.net/dstsv2/common/oauth2/token")
+
+    @patch("msal.authority._instance_discovery")
+    @patch("msal.authority.tenant_discovery", return_value={
+        "authorization_endpoint": authorization_endpoint,
+        "token_endpoint": token_endpoint,
+    })  # We need to create new patches (i.e. mocks) for non-inherited test cases
+    def test_application_obj_should_accept_dsts_url_as_an_authority(
+            self, oidc_discovery, instance_discovery):
+        app = msal.ClientApplication("id", authority=self.authority)
+        instance_discovery.assert_not_called()
+        oidc_discovery.assert_called_once_with(
+            self.oidc_discovery_endpoint, app.http_client)
+        self.assertEqual(
+            app.authority.authorization_endpoint, self.authorization_endpoint)
+        self.assertEqual(app.authority.token_endpoint, self.token_endpoint)
+
+
 class TestAuthorityInternalHelperCanonicalize(unittest.TestCase):
 
     def test_canonicalize_tenant_followed_by_extra_paths(self):