fixup! kola: azure: Implement support for managed identities

auth/azure.go

@@ -124,21 +124,45 @@ func (ap *AzureProfile) AsOptions() []Options {
 	return o
 }
 
+type SubFilter struct {
+	name string
+	id   string
+}
+
+func FilterByName(name string) SubFilter {
+	return SubFilter{name: name}
+}
+func FilterByID(id string) SubFilter {
+	return SubFilter{id: id}
+}
+func (s *SubFilter) IsEmpty() bool {
+	return s.name == "" && s.id == ""
+}
+func (s *SubFilter) Matches(opts *Options) bool {
+	if s.name != "" && opts.SubscriptionName == s.name {
+		return true
+	}
+	if s.id != "" && opts.SubscriptionID == s.id {
+		return true
+	}
+	return false
+}
+
 // SubscriptionOptions returns the name subscription in the Azure profile as a Options struct.
 // If the subscription name is "", the first subscription is returned.
 // If there are no subscriptions or the named subscription is not found, SubscriptionOptions returns nil.
-func (ap *AzureProfile) SubscriptionOptions(name string) *Options {
+func (ap *AzureProfile) SubscriptionOptions(filter SubFilter) *Options {
 	opts := ap.AsOptions()
 
 	if len(opts) == 0 {
 		return nil
 	}
 
-	if name == "" {
+	if filter.IsEmpty() {
 		return &opts[0]
 	} else {
 		for _, o := range ap.AsOptions() {
-			if o.SubscriptionName == name {
+			if filter.Matches(&o) {
 				return &o
 			}
 		}

platform/api/azure/api.go

@@ -67,11 +67,6 @@ func setOptsFromProfile(opts *Options) error {
 		return fmt.Errorf("couldn't read Azure profile: %v", err)
 	}
 
-	subOpts := profiles.SubscriptionOptions(opts.AzureSubscription)
-	if subOpts == nil {
-		return fmt.Errorf("Azure subscription named %q doesn't exist in %q", opts.AzureSubscription, opts.AzureProfile)
-	}
-
 	if os.Getenv("AZURE_AUTH_LOCATION") == "" {
 		if opts.AzureAuthLocation == "" {
 			user, err := user.Current()
@@ -84,6 +79,20 @@ func setOptsFromProfile(opts *Options) error {
 		os.Setenv("AZURE_AUTH_LOCATION", opts.AzureAuthLocation)
 	}
 
+	var subOpts *internalAuth.Options
+	if opts.AzureSubscription == "" {
+		settings, err := auth.GetSettingsFromFile()
+		if err != nil {
+			return err
+		}
+		subOpts = profiles.SubscriptionOptions(internalAuth.FilterByID(settings.GetSubscriptionID()))
+	} else {
+		subOpts = profiles.SubscriptionOptions(internalAuth.FilterByName(opts.AzureSubscription))
+	}
+	if subOpts == nil {
+		return fmt.Errorf("Azure subscription named %q doesn't exist in %q", opts.AzureSubscription, opts.AzureProfile)
+	}
+
 	if opts.SubscriptionID == "" {
 		opts.SubscriptionID = subOpts.SubscriptionID
 	}
@@ -127,6 +136,12 @@ func New(opts *Options) (*API, error) {
 		if err != nil {
 			return nil, fmt.Errorf("failed to get options from azure profile: %w", err)
 		}
+	} else {
+		subid, err := msiGetSubscriptionID()
+		if err != nil {
+			return nil, fmt.Errorf("failed to query subscription id: %w", err)
+		}
+		opts.SubscriptionID = subid
 	}
 
 	var client management.Client
@@ -163,15 +178,7 @@ func (a *API) newAuthorizer(baseURI string) (autorest.Authorizer, error) {
 	return settings.GetMSI().Authorizer()
 }
 
-func (a *API) getSubscriptionID(auther autorest.Authorizer) (string, error) {
-	if !a.Opts.UseIdentity {
-		settings, err := auth.GetSettingsFromFile()
-		if err != nil {
-			return "", err
-		}
-		return settings.GetSubscriptionID(), nil
-	}
-
+func msiGetSubscriptionID() (string, error) {
 	settings, err := auth.GetSettingsFromEnvironment()
 	if err != nil {
 		return "", err
@@ -180,6 +187,10 @@ func (a *API) getSubscriptionID(auther autorest.Authorizer) (string, error) {
 	if subid != "" {
 		return subid, nil
 	}
+	auther, err := settings.GetMSI().Authorizer()
+	if err != nil {
+		return "", err
+	}
 	client := subscriptions.NewClient()
 	client.Authorizer = auther
 	iter, err := client.ListComplete(context.Background())
@@ -208,10 +219,8 @@ func (a *API) SetupClients() error {
 	if err != nil {
 		return err
 	}
-	subid, err := a.getSubscriptionID(auther)
-	if err != nil {
-		return err
-	}
+	subid := a.Opts.SubscriptionID
+
 	a.rgClient = resources.NewGroupsClient(subid)
 	a.rgClient.Authorizer = auther