Merge pull request #1304 from k8s-infra-cherrypick-robot/cherry-pick-1303-to-release-1.24

[release-1.24] chore: refine sastoken cache

Author: k8s-ci-robot

pkg/blob/controllerserver.go

@@ -843,46 +843,40 @@ func (d *Driver) authorizeAzcopyWithIdentity() ([]string, error) {
 // 3. azcopy returns AuthorizationPermissionMismatch error when using service principal or managed identity
 func (d *Driver) getAzcopyAuth(ctx context.Context, accountName, accountKey, storageEndpointSuffix string, accountOptions *azure.AccountOptions, secrets map[string]string, secretName, secretNamespace string) (string, []string, error) {
 	var authAzcopyEnv []string
+	var err error
 	useSasToken := false
 	if !d.useDataPlaneAPI("", accountName) && len(secrets) == 0 && len(secretName) == 0 {
-		var err error
+		// search in cache first
+		if cache, err := d.azcopySasTokenCache.Get(accountName, azcache.CacheReadTypeDefault); err == nil && cache != nil {
+			klog.V(2).Infof("use sas token for account(%s) since this account is found in azcopySasTokenCache", accountName)
+			return cache.(string), nil, nil
+		}
+
 		authAzcopyEnv, err = d.authorizeAzcopyWithIdentity()
 		if err != nil {
 			klog.Warningf("failed to authorize azcopy with identity, error: %v", err)
 		} else {
 			if len(authAzcopyEnv) > 0 {
-				// search in cache first
-				cache, err := d.azcopySasTokenCache.Get(accountName, azcache.CacheReadTypeDefault)
-				if err != nil {
-					return "", nil, fmt.Errorf("get(%s) from azcopySasTokenCache failed with error: %v", accountName, err)
+				out, testErr := d.azcopy.TestListJobs(accountName, storageEndpointSuffix, authAzcopyEnv)
+				if testErr != nil {
+					return "", nil, fmt.Errorf("azcopy list command failed with error(%v): %v", testErr, out)
 				}
-				if cache != nil {
-					klog.V(2).Infof("use sas token for account(%s) since this account is found in azcopySasTokenCache", accountName)
+				if strings.Contains(out, authorizationPermissionMismatch) {
+					klog.Warningf("azcopy list failed with AuthorizationPermissionMismatch error, should assign \"Storage Blob Data Contributor\" role to controller identity, fall back to use sas token, original output: %v", out)
 					useSasToken = true
-				} else {
-					out, testErr := d.azcopy.TestListJobs(accountName, storageEndpointSuffix, authAzcopyEnv)
-					if testErr != nil {
-						return "", nil, fmt.Errorf("azcopy list command failed with error(%v): %v", testErr, out)
-					}
-					if strings.Contains(out, authorizationPermissionMismatch) {
-						klog.Warningf("azcopy list failed with AuthorizationPermissionMismatch error, should assign \"Storage Blob Data Contributor\" role to controller identity, fall back to use sas token, original output: %v", out)
-						d.azcopySasTokenCache.Set(accountName, "")
-						useSasToken = true
-					}
 				}
 			}
 		}
 	}
 
 	if len(secrets) > 0 || len(secretName) > 0 || len(authAzcopyEnv) == 0 || useSasToken {
-		var err error
 		if accountKey == "" {
 			if _, accountKey, err = d.GetStorageAccesskey(ctx, accountOptions, secrets, secretName, secretNamespace); err != nil {
 				return "", nil, err
 			}
 		}
 		klog.V(2).Infof("generate sas token for account(%s)", accountName)
-		sasToken, err := generateSASToken(accountName, accountKey, storageEndpointSuffix, d.sasTokenExpirationMinutes)
+		sasToken, err := d.generateSASToken(accountName, accountKey, storageEndpointSuffix, d.sasTokenExpirationMinutes)
 		return sasToken, nil, err
 	}
 	return "", authAzcopyEnv, nil
@@ -914,7 +908,17 @@ func parseDays(dayStr string) (int32, error) {
 }
 
 // generateSASToken generate a sas token for storage account
-func generateSASToken(accountName, accountKey, storageEndpointSuffix string, expiryTime int) (string, error) {
+func (d *Driver) generateSASToken(accountName, accountKey, storageEndpointSuffix string, expiryTime int) (string, error) {
+	// search in cache first
+	cache, err := d.azcopySasTokenCache.Get(accountName, azcache.CacheReadTypeDefault)
+	if err != nil {
+		return "", fmt.Errorf("get(%s) from azcopySasTokenCache failed with error: %v", accountName, err)
+	}
+	if cache != nil {
+		klog.V(2).Infof("use sas token for account(%s) since this account is found in azcopySasTokenCache", accountName)
+		return cache.(string), nil
+	}
+
 	credential, err := azblob.NewSharedKeyCredential(accountName, accountKey)
 	if err != nil {
 		return "", status.Errorf(codes.Internal, fmt.Sprintf("failed to generate sas token in creating new shared key credential, accountName: %s, err: %s", accountName, err.Error()))
@@ -936,5 +940,7 @@ func generateSASToken(accountName, accountKey, storageEndpointSuffix string, exp
 	if err != nil {
 		return "", err
 	}
-	return "?" + u.RawQuery, nil
+	sasToken := "?" + u.RawQuery
+	d.azcopySasTokenCache.Set(accountName, sasToken)
+	return sasToken, nil
 }

pkg/blob/controllerserver_test.go

@@ -1808,6 +1808,7 @@ func TestParseDays(t *testing.T) {
 }
 
 func TestGenerateSASToken(t *testing.T) {
+	d := NewFakeDriver()
 	storageEndpointSuffix := "core.windows.net"
 	tests := []struct {
 		name        string
@@ -1833,7 +1834,7 @@ func TestGenerateSASToken(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			sas, err := generateSASToken(tt.accountName, tt.accountKey, storageEndpointSuffix, 30)
+			sas, err := d.generateSASToken(tt.accountName, tt.accountKey, storageEndpointSuffix, 30)
 			if !reflect.DeepEqual(err, tt.expectedErr) {
 				t.Errorf("generateSASToken error = %v, expectedErr %v, sas token = %v, want %v", err, tt.expectedErr, sas, tt.want)
 				return