From 0ba8458a57c400d2c9d666c547c7de0195b24fb3 Mon Sep 17 00:00:00 2001
From: Charles Lowell <10964656+chlowell@users.noreply.github.com>
Date: Tue, 4 Jan 2022 14:06:32 -0800
Subject: [PATCH] Make EndpointToScope resilient to subdomains (#16737)

---
 sdk/azcore/internal/shared/shared.go      | 18 ++++++++++++++++
 sdk/azcore/internal/shared/shared_test.go | 26 +++++++++++++++++++++--
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/sdk/azcore/internal/shared/shared.go b/sdk/azcore/internal/shared/shared.go
index 5e008b18e5..1e44f27c28 100644
--- a/sdk/azcore/internal/shared/shared.go
+++ b/sdk/azcore/internal/shared/shared.go
@@ -14,7 +14,9 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"net/url"
 	"strconv"
+	"strings"
 	"time"
 )
 
@@ -179,9 +181,25 @@ func (r *NopClosingBytesReader) Seek(offset int64, whence int) (int64, error) {
 }
 
 const defaultScope = "/.default"
+const chinaCloudARMScope = "https://management.core.chinacloudapi.cn/" + defaultScope
+const publicCloudARMScope = "https://management.core.windows.net/" + defaultScope
+const usGovCloudARMScope = "https://management.core.usgovcloudapi.net/" + defaultScope
 
 // EndpointToScope converts the provided URL endpoint to its default scope.
 func EndpointToScope(endpoint string) string {
+	parsed, err := url.Parse(endpoint)
+	if err == nil {
+		host := parsed.Hostname()
+		switch {
+		case strings.HasSuffix(host, "management.azure.com"):
+			return publicCloudARMScope
+		case strings.HasSuffix(host, "management.usgovcloudapi.net"):
+			return usGovCloudARMScope
+		case strings.HasSuffix(host, "management.chinacloudapi.cn"):
+			return chinaCloudARMScope
+		}
+	}
+	// fall back to legacy behavior when endpoint doesn't parse or match a known cloud's ARM endpoint
 	if endpoint[len(endpoint)-1] != '/' {
 		endpoint += "/"
 	}
diff --git a/sdk/azcore/internal/shared/shared_test.go b/sdk/azcore/internal/shared/shared_test.go
index b40818b34d..e8422eac29 100644
--- a/sdk/azcore/internal/shared/shared_test.go
+++ b/sdk/azcore/internal/shared/shared_test.go
@@ -98,8 +98,30 @@ func TestHasStatusCode(t *testing.T) {
 }
 
 func TestEndpointToScope(t *testing.T) {
-	if s := EndpointToScope("https://management.usgovcloudapi.net"); s != "https://management.usgovcloudapi.net//.default" {
-		t.Fatalf("unexpected scope %s", s)
+	knownClouds := map[string][]string{
+		chinaCloudARMScope:  {"https://foo.management.chinacloudapi.cn", "https://management.chinacloudapi.cn"},
+		publicCloudARMScope: {"https://centraluseuap.management.azure.com", "https://management.azure.com"},
+		usGovCloudARMScope:  {"https://foo.management.usgovcloudapi.net", "https://management.usgovcloudapi.net"},
+	}
+	for expected, endpoints := range knownClouds {
+		for _, endpoint := range endpoints {
+			if actual := EndpointToScope(endpoint); actual != expected {
+				t.Fatalf(`unexpected scope "%s" for endpoint "%s"`, actual, endpoint)
+			}
+			if actual := EndpointToScope(endpoint + "/"); actual != expected {
+				t.Fatalf(`unexpected scope "%s" for endpoint "%s"/`, actual, endpoint)
+			}
+		}
+	}
+
+	// legacy behavior for unknown clouds: add "//.default" suffix to endpoint
+	for _, endpoint := range []string{"localhost", "http://foo.bar"} {
+		if actual := EndpointToScope(endpoint); actual != endpoint+"//.default" {
+			t.Fatalf(`unexpected scope "%s" for endpoint "%s"`, actual, endpoint)
+		}
+		if actual := EndpointToScope(endpoint + "/"); actual != endpoint+"//.default" {
+			t.Fatalf(`unexpected scope "%s" for endpoint "%s"/`, actual, endpoint)
+		}
 	}
 }