From 6cb5bd3469ef3a9cdabcd5f164db48230966ba4e Mon Sep 17 00:00:00 2001
From: David Parks <DavidParks8@outlook.com>
Date: Thu, 23 Apr 2020 16:15:47 -0700
Subject: [PATCH] fix: ACR auth fails in private azure clouds

---
 pkg/credentialprovider/azure/BUILD            |  1 +
 .../azure/azure_credentials.go                | 41 ++++++++++++++++++-
 .../azure/azure_credentials_test.go           | 31 +++++++++++++-
 3 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/pkg/credentialprovider/azure/BUILD b/pkg/credentialprovider/azure/BUILD
index e3c4c814f4f6..18325bef3946 100644
--- a/pkg/credentialprovider/azure/BUILD
+++ b/pkg/credentialprovider/azure/BUILD
@@ -32,6 +32,7 @@ go_test(
     embed = [":go_default_library"],
     deps = [
         "//vendor/github.com/Azure/azure-sdk-for-go/services/containerregistry/mgmt/2019-05-01/containerregistry:go_default_library",
+        "//vendor/github.com/Azure/go-autorest/autorest/azure:go_default_library",
         "//vendor/github.com/Azure/go-autorest/autorest/to:go_default_library",
     ],
 )
diff --git a/pkg/credentialprovider/azure/azure_credentials.go b/pkg/credentialprovider/azure/azure_credentials.go
index 1cf5f908d181..aaa50b8eb9cf 100644
--- a/pkg/credentialprovider/azure/azure_credentials.go
+++ b/pkg/credentialprovider/azure/azure_credentials.go
@@ -23,6 +23,7 @@ import (
 	"io/ioutil"
 	"os"
 	"regexp"
+	"strings"
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/services/containerregistry/mgmt/2019-05-01/containerregistry"
@@ -186,7 +187,7 @@ func (a *acrProvider) Provide(image string) credentialprovider.DockerConfig {
 	cfg := credentialprovider.DockerConfig{}
 
 	if a.config.UseManagedIdentityExtension {
-		if loginServer := parseACRLoginServerFromImage(image); loginServer == "" {
+		if loginServer := a.parseACRLoginServerFromImage(image); loginServer == "" {
 			klog.V(4).Infof("image(%s) is not from ACR, skip MSI authentication", image)
 		} else {
 			if cred, err := getACRDockerEntryFromARMToken(a, loginServer); err == nil {
@@ -203,6 +204,28 @@ func (a *acrProvider) Provide(image string) credentialprovider.DockerConfig {
 			}
 			cfg[url] = *cred
 		}
+
+		// Handle the custom cloud case
+		// In clouds where ACR is not yet deployed, the string will be empty
+		if a.environment != nil && strings.Contains(a.environment.ContainerRegistryDNSSuffix, ".azurecr.") {
+			customAcrSuffix := "*" + a.environment.ContainerRegistryDNSSuffix
+			hasBeenAdded := false
+			for _, url := range containerRegistryUrls {
+				if strings.EqualFold(url, customAcrSuffix) {
+					hasBeenAdded = true
+					break
+				}
+			}
+
+			if !hasBeenAdded {
+				cred := &credentialprovider.DockerConfigEntry{
+					Username: a.config.AADClientID,
+					Password: a.config.AADClientSecret,
+					Email:    dummyRegistryEmail,
+				}
+				cfg[customAcrSuffix] = *cred
+			}
+		}
 	}
 
 	// add ACR anonymous repo support: use empty username and password for anonymous access
@@ -252,10 +275,24 @@ func getACRDockerEntryFromARMToken(a *acrProvider, loginServer string) (*credent
 // parseACRLoginServerFromImage takes image as parameter and returns login server of it.
 // Parameter `image` is expected in following format: foo.azurecr.io/bar/imageName:version
 // If the provided image is not an acr image, this function will return an empty string.
-func parseACRLoginServerFromImage(image string) string {
+func (a *acrProvider) parseACRLoginServerFromImage(image string) string {
 	match := acrRE.FindAllString(image, -1)
 	if len(match) == 1 {
 		return match[0]
 	}
+
+	// handle the custom cloud case
+	if a != nil && a.environment != nil {
+		cloudAcrSuffix := a.environment.ContainerRegistryDNSSuffix
+		cloudAcrSuffixLength := len(cloudAcrSuffix)
+		if cloudAcrSuffixLength > 0 {
+			customAcrSuffixIndex := strings.Index(image, cloudAcrSuffix)
+			if customAcrSuffixIndex != -1 {
+				endIndex := customAcrSuffixIndex + cloudAcrSuffixLength
+				return image[0:endIndex]
+			}
+		}
+	}
+
 	return ""
 }
diff --git a/pkg/credentialprovider/azure/azure_credentials_test.go b/pkg/credentialprovider/azure/azure_credentials_test.go
index 88962e260f6a..fffddea3f64f 100644
--- a/pkg/credentialprovider/azure/azure_credentials_test.go
+++ b/pkg/credentialprovider/azure/azure_credentials_test.go
@@ -22,6 +22,7 @@ import (
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/services/containerregistry/mgmt/2019-05-01/containerregistry"
+	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/Azure/go-autorest/autorest/to"
 )
 
@@ -96,6 +97,30 @@ func Test(t *testing.T) {
 }
 
 func TestParseACRLoginServerFromImage(t *testing.T) {
+	configStr := `
+    {
+        "aadClientId": "foo",
+        "aadClientSecret": "bar"
+    }`
+	result := []containerregistry.Registry{
+		{
+			Name: to.StringPtr("foo"),
+			RegistryProperties: &containerregistry.RegistryProperties{
+				LoginServer: to.StringPtr("*.azurecr.io"),
+			},
+		},
+	}
+	fakeClient := &fakeClient{
+		results: result,
+	}
+
+	provider := &acrProvider{
+		registryClient: fakeClient,
+	}
+	provider.loadConfig(bytes.NewBufferString(configStr))
+	provider.environment = &azure.Environment{
+		ContainerRegistryDNSSuffix: ".azurecr.my.cloud",
+	}
 	tests := []struct {
 		image    string
 		expected string
@@ -124,9 +149,13 @@ func TestParseACRLoginServerFromImage(t *testing.T) {
 			image:    "foo.azurecr.us/bar/image:version",
 			expected: "foo.azurecr.us",
 		},
+		{
+			image:    "foo.azurecr.my.cloud/bar/image:version",
+			expected: "foo.azurecr.my.cloud",
+		},
 	}
 	for _, test := range tests {
-		if loginServer := parseACRLoginServerFromImage(test.image); loginServer != test.expected {
+		if loginServer := provider.parseACRLoginServerFromImage(test.image); loginServer != test.expected {
 			t.Errorf("function parseACRLoginServerFromImage returns \"%s\" for image %s, expected \"%s\"", loginServer, test.image, test.expected)
 		}
 	}