wip

Author: bryan-cox

pkg/storage/azure/azure.go

@@ -334,7 +334,7 @@ func NewDriver(ctx context.Context, c *imageregistryv1.ImageRegistryConfigStorag
 }
 
 func (d *driver) newAzClient(cfg *Azure, environment autorestazure.Environment, tagset map[string]*string) (*azureclient.Client, error) {
-	client, err := azureclient.New(&azureclient.Options{
+	clientOptions := &azureclient.Options{
 		Environment:        environment,
 		TenantID:           cfg.TenantID,
 		ClientID:           cfg.ClientID,
@@ -343,10 +343,19 @@ func (d *driver) newAzClient(cfg *Azure, environment autorestazure.Environment,
 		SubscriptionID:     cfg.SubscriptionID,
 		TagSet:             tagset,
 		Policies:           d.policies,
-	})
+	}
+
+	if cred, ok, err := d.ensureUAMICred(d.Context, environment); err != nil {
+		return nil, err
+	} else if ok {
+		clientOptions.Creds = cred
+	}
+
+	client, err := azureclient.New(clientOptions)
 	if err != nil {
 		return nil, err
 	}
+
 	return client, nil
 }
 
@@ -381,25 +390,10 @@ func (d *driver) storageAccountsClient(cfg *Azure, environment autorestazure.Env
 	// UserAssignedIdentityCredentials is specifically for managed Azure HCP
 	userAssignedIdentityCredentialsFilePath := os.Getenv("MANAGED_AZURE_HCP_CREDENTIALS_FILE_PATH")
 	if userAssignedIdentityCredentialsFilePath != "" {
-		var ok bool
-
-		// We need to only store the Azure credentials once and reuse them after that.
-		storedCreds, found := d.azureCredentials.Load(azureCredentialsKey)
-		if !found {
-			klog.V(2).Info("Using UserAssignedIdentityCredentials for Azure authentication for managed Azure HCP")
-			clientOptions := azcore.ClientOptions{
-				Cloud: cloudConfig,
-			}
-			cred, err = dataplane.NewUserAssignedIdentityCredential(context.Background(), userAssignedIdentityCredentialsFilePath, dataplane.WithClientOpts(clientOptions))
-			if err != nil {
-				return storage.AccountsClient{}, err
-			}
-			d.azureCredentials.Store(azureCredentialsKey, cred)
-		} else {
-			cred, ok = storedCreds.(azcore.TokenCredential)
-			if !ok {
-				return storage.AccountsClient{}, fmt.Errorf("expected %T to be a TokenCredential", storedCreds)
-			}
+		if c, ok, err := d.ensureUAMICred(d.Context, environment); err != nil {
+			return storage.AccountsClient{}, err
+		} else if ok {
+			cred = c
 		}
 	} else if strings.TrimSpace(cfg.ClientSecret) == "" {
 		options := azidentity.WorkloadIdentityCredentialOptions{
@@ -1237,14 +1231,22 @@ func (d *driver) RemoveStorage(cr *imageregistryv1.Config) (retry bool, err erro
 	}
 
 	if d.Config.NetworkAccess != nil && d.Config.NetworkAccess.Internal != nil && d.Config.NetworkAccess.Internal.PrivateEndpointName != "" {
-		azclient, err := azureclient.New(&azureclient.Options{
+		clientOptions := &azureclient.Options{
 			Environment:        environment,
 			TenantID:           cfg.TenantID,
 			ClientID:           cfg.ClientID,
 			ClientSecret:       cfg.ClientSecret,
 			FederatedTokenFile: cfg.FederatedTokenFile,
 			SubscriptionID:     cfg.SubscriptionID,
-		})
+		}
+
+		storedCreds, found := d.azureCredentials.Load(azureCredentialsKey)
+		if found {
+			klog.V(2).Infof("Using driver cached credential for new Azure client to remove storage")
+			clientOptions.Creds = storedCreds.(azcore.TokenCredential)
+		}
+
+		azclient, err := azureclient.New(clientOptions)
 		if err != nil {
 			util.UpdateCondition(
 				cr,
@@ -1320,3 +1322,36 @@ func (d *driver) RemoveStorage(cr *imageregistryv1.Config) (retry bool, err erro
 func (d *driver) ID() string {
 	return d.Config.Container
 }
+
+func (d *driver) ensureUAMICred(ctx context.Context, env autorestazure.Environment) (azcore.TokenCredential, bool, error) {
+	if stored, ok := d.azureCredentials.Load(azureCredentialsKey); ok {
+		if cred, ok := stored.(azcore.TokenCredential); ok {
+			klog.V(2).Infof("Loaded UAMICred from driver cache")
+			return cred, true, nil
+		}
+		return nil, false, fmt.Errorf("expected cached credential to be azcore.TokenCredential")
+	}
+	if os.Getenv("MANAGED_AZURE_HCP_CREDENTIALS_FILE_PATH") == "" {
+		return nil, false, nil
+	}
+	cloudConfig := cloud.Configuration{
+		ActiveDirectoryAuthorityHost: env.ActiveDirectoryEndpoint,
+		Services: map[cloud.ServiceName]cloud.ServiceConfiguration{
+			cloud.ResourceManager: {
+				Audience: env.TokenAudience,
+				Endpoint: env.ResourceManagerEndpoint,
+			},
+		},
+	}
+	cred, err := dataplane.NewUserAssignedIdentityCredential(
+		ctx,
+		os.Getenv("MANAGED_AZURE_HCP_CREDENTIALS_FILE_PATH"),
+		dataplane.WithClientOpts(azcore.ClientOptions{Cloud: cloudConfig}),
+	)
+	if err != nil {
+		return nil, false, err
+	}
+	d.azureCredentials.Store(azureCredentialsKey, cred)
+	klog.V(2).Infof("Storing UAMICred for driver")
+	return cred, true, nil
+}