feat: Add auxiliary token policy and update VMsClient to use in client options

pkg/auth/cred.go

@@ -18,6 +18,10 @@ package auth
 
 import (
 	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
@@ -33,6 +37,38 @@ type expireEarlyTokenCredential struct {
 	cred azcore.TokenCredential
 }
 
+func GetAuxiliaryToken(ctx context.Context, scope string) (azcore.AccessToken, error) {
+	client := &http.Client{}
+	var token azcore.AccessToken
+	req, err := http.NewRequest("GET", "msi-connector.msi-connector.svc.cluster.local/karpenter", nil)
+	if err != nil {
+		return token, err
+	}
+	q := req.URL.Query()
+	q.Add("scope", scope)
+	req.URL.RawQuery = q.Encode()
+
+	// Send the request
+	resp, err := client.Do(req)
+	if err != nil {
+		return token, fmt.Errorf("error making request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	// Read the response body
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return token, fmt.Errorf("error reading response: %w", err)
+	}
+
+	// Unmarshal the response body into the AccessToken struct
+	err = json.Unmarshal(body, &token)
+	if err != nil {
+		return token, fmt.Errorf("error unmarshalling response: %w", err)
+	}
+	return token, nil
+}
+
 func NewTokenWrapper(cred azcore.TokenCredential) azcore.TokenCredential {
 	return &expireEarlyTokenCredential{
 		cred: cred,

pkg/auth/policy.go

@@ -0,0 +1,39 @@
+/*
+Portions Copyright (c) Microsoft Corporation.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package auth
+
+import (
+	"net/http"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+)
+
+// AuxiliaryTokenPolicy provides a custom policy used to authenticate
+// with shared node image galleries.
+type AuxiliaryTokenPolicy struct {
+	Token string
+}
+
+func (p *AuxiliaryTokenPolicy) Do(req *policy.Request) (*http.Response, error) {
+	req.Raw().Header.Add("x-ms-authorization-auxiliary", "Bearer "+p.Token)
+	return req.Next()
+}
+
+func NewAuxiliaryTokenPolicy(token azcore.AccessToken) AuxiliaryTokenPolicy {
+	return AuxiliaryTokenPolicy{Token: token.Token}
+}

pkg/providers/instance/azure_client.go

@@ -18,7 +18,10 @@ package instance
 
 import (
 	"context"
+	"fmt"
 
+	armpolicy "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute"
@@ -128,8 +131,12 @@ func NewAZClient(ctx context.Context, cfg *auth.Config, env *azure.Environment)
 		return nil, err
 	}
 	klog.V(5).Infof("Created network interface client %v using token credential", interfacesClient)
-
-	virtualMachinesClient, err := armcompute.NewVirtualMachinesClient(cfg.SubscriptionID, cred, opts)
+	// TODO: Once ListImageVersions work is merged, gate this behind ManagedKarpenter
+	clientOptions, err := getVirtualMachinesClientOptions(ctx)
+	if err != nil {
+		return nil, err
+	}
+	virtualMachinesClient, err := armcompute.NewVirtualMachinesClient(cfg.SubscriptionID, cred, clientOptions)
 	if err != nil {
 		return nil, err
 	}
@@ -167,3 +174,18 @@ func NewAZClient(ctx context.Context, cfg *auth.Config, env *azure.Environment)
 		nodeImageVersionsClient,
 		skuClient), nil
 }
+
+func getVirtualMachinesClientOptions(ctx context.Context) (*armpolicy.ClientOptions, error) {
+	// TODO: pass in custom scope based on cloud env
+	token, err := auth.GetAuxiliaryToken(ctx, "https://management.azure.com/.default")
+	if err != nil {
+		return &armpolicy.ClientOptions{}, fmt.Errorf("failed to get auxiliary token: %w", err)
+	}
+	auxPolicy := auth.NewAuxiliaryTokenPolicy(token)
+
+	return &armpolicy.ClientOptions{
+		ClientOptions: policy.ClientOptions{
+			PerRetryPolicies: []policy.Policy{&auxPolicy},
+		},
+	}, nil
+}