Merge pull request #96355 from andyzhangx/acr-mi-fix2

fix pull image error from multiple ACRs using azure managed identity

pkg/credentialprovider/azure/BUILD

@@ -16,6 +16,7 @@ go_library(
     importpath = "k8s.io/kubernetes/pkg/credentialprovider/azure",
     deps = [
         "//pkg/credentialprovider:go_default_library",
+        "//staging/src/k8s.io/client-go/tools/cache:go_default_library",
         "//staging/src/k8s.io/legacy-cloud-providers/azure/auth:go_default_library",
         "//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:go_default_library",
@@ -32,6 +33,7 @@ go_test(
     srcs = ["azure_credentials_test.go"],
     embed = [":go_default_library"],
     deps = [
+        "//staging/src/k8s.io/client-go/tools/cache:go_default_library",
         "//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",

pkg/credentialprovider/azure/azure_credentials.go

@@ -34,6 +34,7 @@ import (
 	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/spf13/pflag"
 
+	"k8s.io/client-go/tools/cache"
 	"k8s.io/klog/v2"
 	"k8s.io/kubernetes/pkg/credentialprovider"
 	"k8s.io/legacy-cloud-providers/azure/auth"
@@ -56,12 +57,30 @@ var (
 // init registers the various means by which credentials may
 // be resolved on Azure.
 func init() {
-	credentialprovider.RegisterCredentialProvider("azure",
+	credentialprovider.RegisterCredentialProvider(
-		&credentialprovider.CachingDockerConfigProvider{
+		"azure",
-			Provider:    NewACRProvider(flagConfigFile),
+		NewACRProvider(flagConfigFile),
-			Lifetime:    1 * time.Minute,
+	)
-			ShouldCache: func(d credentialprovider.DockerConfig) bool { return len(d) > 0 },
+}
-		})
+
+type cacheEntry struct {
+	expiresAt   time.Time
+	credentials credentialprovider.DockerConfigEntry
+	registry    string
+}
+
+// acrExpirationPolicy implements ExpirationPolicy from client-go.
+type acrExpirationPolicy struct{}
+
+// stringKeyFunc returns the cache key as a string
+func stringKeyFunc(obj interface{}) (string, error) {
+	key := obj.(*cacheEntry).registry
+	return key, nil
+}
+
+// IsExpired checks if the ACR credentials are expired.
+func (p *acrExpirationPolicy) IsExpired(entry *cache.TimestampedEntry) bool {
+	return time.Now().After(entry.Obj.(*cacheEntry).expiresAt)
 }
 
 // RegistriesClient is a testable interface for the ACR client List operation.
@@ -106,6 +125,7 @@ func (az *azRegistriesClient) List(ctx context.Context) ([]containerregistry.Reg
 func NewACRProvider(configFile *string) credentialprovider.DockerConfigProvider {
 	return &acrProvider{
 		file:  configFile,
+		cache: cache.NewExpirationStore(stringKeyFunc, &acrExpirationPolicy{}),
 	}
 }
 
@@ -115,6 +135,7 @@ type acrProvider struct {
 	environment           *azure.Environment
 	registryClient        RegistriesClient
 	servicePrincipalToken *adal.ServicePrincipalToken
+	cache                 cache.Store
 }
 
 // ParseConfig returns a parsed configuration for an Azure cloudprovider config file
@@ -185,25 +206,63 @@ func (a *acrProvider) Enabled() bool {
 	return true
 }
 
-func (a *acrProvider) Provide(image string) credentialprovider.DockerConfig {
+// getFromCache attempts to get credentials from the cache
-	klog.V(4).Infof("try to provide secret for image %s", image)
+func (a *acrProvider) getFromCache(loginServer string) (credentialprovider.DockerConfig, bool) {
 	cfg := credentialprovider.DockerConfig{}
-
+	obj, exists, err := a.cache.GetByKey(loginServer)
-	defaultConfigEntry := credentialprovider.DockerConfigEntry{
+	if err != nil {
-		Username: "",
+		klog.Errorf("error getting ACR credentials from cache: %v", err)
-		Password: "",
+		return cfg, false
-		Email:    dummyRegistryEmail,
+	}
+	if !exists {
+		return cfg, false
 	}
 
-	if a.config.UseManagedIdentityExtension {
+	entry := obj.(*cacheEntry)
-		if loginServer := a.parseACRLoginServerFromImage(image); loginServer == "" {
+	cfg[entry.registry] = entry.credentials
-			klog.V(4).Infof("image(%s) is not from ACR, skip MSI authentication", image)
+	return cfg, true
-		} else {
+}
-			if cred, err := getACRDockerEntryFromARMToken(a, loginServer); err == nil {
+
+// getFromACR gets credentials from ACR since they are not in the cache
+func (a *acrProvider) getFromACR(loginServer string) (credentialprovider.DockerConfig, error) {
+	cfg := credentialprovider.DockerConfig{}
+	cred, err := getACRDockerEntryFromARMToken(a, loginServer)
+	if err != nil {
+		return cfg, err
+	}
+
+	entry := &cacheEntry{
+		expiresAt:   time.Now().Add(10 * time.Minute),
+		credentials: *cred,
+		registry:    loginServer,
+	}
+	if err := a.cache.Add(entry); err != nil {
+		return cfg, err
+	}
 	cfg[loginServer] = *cred
+	return cfg, nil
+}
+
+func (a *acrProvider) Provide(image string) credentialprovider.DockerConfig {
+	loginServer := a.parseACRLoginServerFromImage(image)
+	if loginServer == "" {
+		klog.V(2).Infof("image(%s) is not from ACR, return empty authentication", image)
+		return credentialprovider.DockerConfig{}
+	}
+
+	cfg := credentialprovider.DockerConfig{}
+	if a.config != nil && a.config.UseManagedIdentityExtension {
+		var exists bool
+		cfg, exists = a.getFromCache(loginServer)
+		if exists {
+			klog.V(4).Infof("Got ACR credentials from cache for %s", loginServer)
+		} else {
+			klog.V(2).Info("unable to get ACR credentials from cache for %s, checking ACR API", loginServer)
+			var err error
+			cfg, err = a.getFromACR(loginServer)
+			if err != nil {
+				klog.Errorf("error getting credentials from ACR for %s %v", loginServer, err)
 			}
-			// add ACR anonymous repo support: use empty username and password for anonymous access
-			cfg["*.azurecr.*"] = defaultConfigEntry
 		}
 	} else {
 		// Add our entry for each of the supported container registry URLs
@@ -237,10 +296,15 @@ func (a *acrProvider) Provide(image string) credentialprovider.DockerConfig {
 				cfg[customAcrSuffix] = *cred
 			}
 		}
+	}
 
 	// add ACR anonymous repo support: use empty username and password for anonymous access
-		cfg["*.azurecr.*"] = defaultConfigEntry
+	defaultConfigEntry := credentialprovider.DockerConfigEntry{
+		Username: "",
+		Password: "",
+		Email:    dummyRegistryEmail,
 	}
+	cfg["*.azurecr.*"] = defaultConfigEntry
 	return cfg
 }
 

pkg/credentialprovider/azure/azure_credentials_test.go

@@ -26,6 +26,7 @@ import (
 	"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"
+	"k8s.io/client-go/tools/cache"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -76,10 +77,11 @@ func Test(t *testing.T) {
 
 	provider := &acrProvider{
 		registryClient: fakeClient,
+		cache:          cache.NewExpirationStore(stringKeyFunc, &acrExpirationPolicy{}),
 	}
 	provider.loadConfig(bytes.NewBufferString(configStr))
 
-	creds := provider.Provide("")
+	creds := provider.Provide("foo.azurecr.io/nginx:v1")
 
 	if len(creds) != len(result)+1 {
 		t.Errorf("Unexpected list: %v, expected length %d", creds, len(result)+1)
@@ -103,11 +105,13 @@ func Test(t *testing.T) {
 func TestProvide(t *testing.T) {
 	testCases := []struct {
 		desc                string
+		image               string
 		configStr           string
 		expectedCredsLength int
 	}{
 		{
 			desc:  "return multiple credentials using Service Principal",
+			image: "foo.azurecr.io/bar/image:v1",
 			configStr: `
     {
         "aadClientId": "foo",
@@ -117,6 +121,7 @@ func TestProvide(t *testing.T) {
 		},
 		{
 			desc:  "retuen 0 credential for non-ACR image using Managed Identity",
+			image: "busybox",
 			configStr: `
     {
 	"UseManagedIdentityExtension": true
@@ -128,10 +133,11 @@ func TestProvide(t *testing.T) {
 	for i, test := range testCases {
 		provider := &acrProvider{
 			registryClient: &fakeClient{},
+			cache:          cache.NewExpirationStore(stringKeyFunc, &acrExpirationPolicy{}),
 		}
 		provider.loadConfig(bytes.NewBufferString(test.configStr))
 
-		creds := provider.Provide("busybox")
+		creds := provider.Provide(test.image)
 		assert.Equal(t, test.expectedCredsLength, len(creds), "TestCase[%d]: %s", i, test.desc)
 	}
 }