fix: Workload Identity token refresh issue

fix

Author: andyzhangx

deploy/csi-blob-driver.yaml

@@ -10,5 +10,7 @@ spec:
   volumeLifecycleModes:
     - Persistent
     - Ephemeral
+  RequiresRepublish: true
   tokenRequests:
     - audience: api://AzureADTokenExchange
+      expirationSeconds: 3600

pkg/blob/blob.go

@@ -23,6 +23,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"sync"
@@ -161,6 +162,7 @@ const (
 
 	DefaultTokenAudience = "api://AzureADTokenExchange" //nolint:gosec // G101 ignore this!
 
+	defaultAzureFederatedTokenDir = "/var/lib/kubelet/" + DefaultDriverName
 )
 
 var (
@@ -586,12 +588,17 @@ func (d *Driver) GetAuthEnv(ctx context.Context, volumeID, protocol string, attr
 			if err != nil {
 				return rgName, accountName, accountKey, containerName, authEnv, err
 			}
+			azureFederatedTokenFile := filepath.Join(defaultAzureFederatedTokenDir, clientID)
+			klog.V(2).Infof("write workload identity token to %s", azureFederatedTokenFile)
+			if err := os.WriteFile(azureFederatedTokenFile, []byte(workloadIdentityToken), 0644); err != nil {
+				return rgName, accountName, accountKey, containerName, authEnv, fmt.Errorf("failed to write azure federated token file %s: %v", azureFederatedTokenFile, err)
+			}
 
 			authEnv = append(authEnv, "AZURE_STORAGE_SPN_CLIENT_ID="+clientID)
 			if tenantID != "" {
 				authEnv = append(authEnv, "AZURE_STORAGE_SPN_TENANT_ID="+tenantID)
 			}
-			authEnv = append(authEnv, "WORKLOAD_IDENTITY_TOKEN="+workloadIdentityToken)
+			authEnv = append(authEnv, "AZURE_FEDERATED_TOKEN_FILE="+azureFederatedTokenFile)
 
 			return rgName, accountName, accountKey, containerName, authEnv, err
 		}

pkg/blob/nodeserver.go

@@ -327,16 +327,17 @@ func (d *Driver) NodeStageVolume(ctx context.Context, req *csi.NodeStageVolumeRe
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "Could not mount target %q: %v", targetPath, err)
 	}
-	if mnt {
-		klog.V(2).Infof("NodeStageVolume: volume %s is already mounted on %s", volumeID, targetPath)
-		return &csi.NodeStageVolumeResponse{}, nil
-	}
 
 	_, accountName, _, containerName, authEnv, err := d.GetAuthEnv(ctx, volumeID, protocol, attrib, secrets)
-	if err != nil {
+	if err != nil && !mnt {
 		return nil, status.Errorf(codes.Internal, "%v", err)
 	}
 
+	if mnt {
+		klog.V(2).Infof("NodeStageVolume: volume %s is already mounted on %s", volumeID, targetPath)
+		return &csi.NodeStageVolumeResponse{}, nil
+	}
+
 	// replace pv/pvc name namespace metadata in subDir
 	containerName = replaceWithMap(containerName, containerNameReplaceMap)
 

pkg/blob/nodeserver_test.go

@@ -243,6 +243,8 @@ func TestNodePublishVolume(t *testing.T) {
 				// Mock NodeStageVolume to return success
 				d.cloud.ResourceGroup = "rg"
 				d.enableBlobMockMount = true
+				// Create the directory for token file
+				_ = makeDir("/var/lib/kubelet/blob.csi.azure.com/")
 			},
 			req: &csi.NodePublishVolumeRequest{
 				VolumeCapability:  &csi.VolumeCapability{AccessMode: &volumeCap},