fix: auxiliary token refresh and use sig in acceptance tests (#991)

* initial implementation, start of tests

* Add tests + update test env to use sig

* ci fixes + small improvements

* Update tests to directly call UseAuxiliaryTokenPolicy, incorporate small feedback

Author: rakechill

pkg/auth/auxiliarytoken.go

@@ -17,10 +17,11 @@ limitations under the License.
 package auth
 
 import (
-	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"sync"
+	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
@@ -36,22 +37,47 @@ var _ policy.Policy = &AuxiliaryTokenPolicy{}
 // AuxiliaryTokenPolicy provides a custom policy used to authenticate
 // with shared node image galleries.
 type AuxiliaryTokenPolicy struct {
-	Token string
+	Token  azcore.AccessToken
+	url    string
+	scope  string
+	client AuxiliaryTokenServer
+	lock   sync.Mutex
+}
+
+func (p *AuxiliaryTokenPolicy) GetAuxiliaryToken() error {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+	// If the token is uninitialized or close to expiration, fetch a new one
+	currentTime := time.Now()
+	if p.Token.ExpiresOn.IsZero() || p.Token.RefreshOn.Before(currentTime) || p.Token.ExpiresOn.Before(currentTime.Add(5*time.Minute)) {
+		newToken, err := getAuxiliaryToken(p.client, p.url, p.scope)
+		if err != nil {
+			return err
+		}
+		p.Token = newToken
+	}
+	return nil
 }
 
 func (p *AuxiliaryTokenPolicy) Do(req *policy.Request) (*http.Response, error) {
-	req.Raw().Header.Add("x-ms-authorization-auxiliary", "Bearer "+p.Token)
+	err := p.GetAuxiliaryToken()
+	if err != nil {
+		log.FromContext(req.Raw().Context()).Error(err, "Failed to get auxiliary token")
+		return nil, err
+	}
+	req.Raw().Header.Add("x-ms-authorization-auxiliary", "Bearer "+p.Token.Token)
 	return req.Next()
 }
 
-func NewAuxiliaryTokenPolicy(ctx context.Context, client AuxiliaryTokenServer, url string, scope string) (*AuxiliaryTokenPolicy, error) {
-	token, err := getAuxiliaryToken(client, url, scope)
-	if err != nil {
-		return nil, fmt.Errorf("failed to get auxiliary token: %w", err)
+func NewAuxiliaryTokenPolicy(client AuxiliaryTokenServer, url string, scope string) *AuxiliaryTokenPolicy {
+	auxPolicy := AuxiliaryTokenPolicy{
+		Token:  azcore.AccessToken{},
+		url:    url,
+		scope:  scope,
+		client: client,
+		lock:   sync.Mutex{},
 	}
-	auxPolicy := AuxiliaryTokenPolicy{Token: token.Token}
-	log.FromContext(ctx).V(1).Info("Will use auxiliary token policy for creating virtual machines")
-	return &auxPolicy, nil
+	return &auxPolicy
 }
 
 func getAuxiliaryToken(client AuxiliaryTokenServer, url string, scope string) (azcore.AccessToken, error) {

pkg/cloudprovider/suite_test.go

@@ -53,6 +53,7 @@ import (
 )
 
 var ctx context.Context
+var testOptions *options.Options
 var stop context.CancelFunc
 var env *coretest.Environment
 var azureEnv *test.Environment
@@ -75,7 +76,8 @@ func TestCloudProvider(t *testing.T) {
 var _ = BeforeSuite(func() {
 	env = coretest.NewEnvironment(coretest.WithCRDs(apis.CRDs...), coretest.WithCRDs(v1alpha1.CRDs...), coretest.WithFieldIndexers(coretest.NodeProviderIDFieldIndexer(ctx)))
 	ctx = coreoptions.ToContext(ctx, coretest.Options())
-	ctx = options.ToContext(ctx, test.Options())
+	testOptions = test.Options()
+	ctx = options.ToContext(ctx, testOptions)
 	ctx, stop = context.WithCancel(ctx)
 	azureEnv = test.NewEnvironment(ctx, env)
 	fakeClock = clock.NewFakeClock(time.Now())
@@ -91,11 +93,12 @@ var _ = AfterSuite(func() {
 })
 
 var _ = BeforeEach(func() {
+	testOptions = test.Options()
 	ctx = coreoptions.ToContext(ctx, coretest.Options())
-	ctx = options.ToContext(ctx, test.Options())
+	ctx = options.ToContext(ctx, testOptions)
 
 	nodeClass = test.AKSNodeClass()
-	test.ApplyDefaultStatus(nodeClass, env)
+	test.ApplyDefaultStatus(nodeClass, env, testOptions.UseSIG)
 
 	nodePool = coretest.NodePool(karpv1.NodePool{
 		Spec: karpv1.NodePoolSpec{

pkg/controllers/nodeclaim/garbagecollection/suite_test.go

@@ -55,6 +55,7 @@ import (
 )
 
 var ctx context.Context
+var testOptions *options.Options
 var env *coretest.Environment
 var azureEnv *test.Environment
 var fakeClock *clock.FakeClock
@@ -74,7 +75,8 @@ func TestAPIs(t *testing.T) {
 
 var _ = BeforeSuite(func() {
 	ctx = coreoptions.ToContext(ctx, coretest.Options())
-	ctx = options.ToContext(ctx, test.Options())
+	testOptions = test.Options()
+	ctx = options.ToContext(ctx, testOptions)
 	env = coretest.NewEnvironment(coretest.WithCRDs(apis.CRDs...), coretest.WithCRDs(v1alpha1.CRDs...))
 	//	ctx, stop = context.WithCancel(ctx)
 	azureEnv = test.NewEnvironment(ctx, env)
@@ -94,7 +96,7 @@ var _ = AfterSuite(func() {
 
 var _ = BeforeEach(func() {
 	nodeClass = test.AKSNodeClass()
-	test.ApplyDefaultStatus(nodeClass, env)
+	test.ApplyDefaultStatus(nodeClass, env, testOptions.UseSIG)
 
 	nodePool = coretest.NodePool(karpv1.NodePool{
 		Spec: karpv1.NodePoolSpec{

pkg/fake/auxiliarytokenserver.go

@@ -40,6 +40,26 @@ var _ auth.AuxiliaryTokenServer = &AuxiliaryTokenServer{}
 
 type AuxiliaryTokenServer struct {
 	AuxiliaryTokenBehavior
+	Token azcore.AccessToken
+}
+
+// NewAuxiliaryTokenServer creates a new AuxiliaryTokenServer with the given token.
+func NewAuxiliaryTokenServer(token string, expiresOn time.Time, refreshOn time.Time) *AuxiliaryTokenServer {
+	return &AuxiliaryTokenServer{
+		Token: azcore.AccessToken{
+			Token:     token,
+			ExpiresOn: expiresOn,
+			RefreshOn: refreshOn,
+		},
+	}
+}
+
+func (c *AuxiliaryTokenServer) SetToken(token string, expiresOn time.Time, refreshOn time.Time) {
+	c.Token = azcore.AccessToken{
+		Token:     token,
+		ExpiresOn: expiresOn,
+		RefreshOn: refreshOn,
+	}
 }
 
 // Reset must be called between tests otherwise tests will pollute each other.
@@ -61,10 +81,7 @@ func (c *AuxiliaryTokenServer) Do(req *http.Request) (*http.Response, error) {
 			return resp, nil
 		}
 
-		token := azcore.AccessToken{
-			Token:     "fake-token",
-			ExpiresOn: time.Now().Add(1 * time.Hour),
-		}
+		token := c.Token
 		tokenBytes, _ := json.Marshal(token)
 		resp.StatusCode = http.StatusOK
 		resp.Body = io.NopCloser(bytes.NewReader(tokenBytes))

pkg/fake/auxiliarytokenserver_test.go

@@ -17,62 +17,55 @@ limitations under the License.
 package fake
 
 import (
-	"context"
+	"net/http"
+	"net/url"
 	"testing"
+	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/karpenter-provider-azure/pkg/auth"
-	armopts "github.com/Azure/karpenter-provider-azure/pkg/utils/opts"
 	"github.com/stretchr/testify/assert"
 )
 
 func Test_AddAuxiliaryTokenPolicyClientOptions(t *testing.T) {
+	defaultToken := azcore.AccessToken{
+		Token:     "test-token",
+		ExpiresOn: time.Now().Add(1 * time.Hour),
+		RefreshOn: time.Now().Add(5 * time.Second),
+	}
 	tests := []struct {
-		name      string
-		expected  azcore.AccessToken
-		wantErr   bool
-		errString string
-		url       string
-		scope     string
+		name       string
+		userAgent  string
+		statusCode int
 	}{
 		{
-			name:      "url is not set",
-			wantErr:   true,
-			errString: "access token server URL is not set",
-			url:       "",
-			scope:     "anything",
-		},
-		{
-			name:      "scope is not set",
-			wantErr:   true,
-			errString: "access token scope is not set",
-			url:       "anything",
-			scope:     "",
+			name:       "default",
+			userAgent:  auth.GetUserAgentExtension(),
+			statusCode: http.StatusOK,
 		},
 		{
-			name:    "default",
-			wantErr: false,
-			url:     "http://test-url.com",
-			scope:   "test-scope",
+			name:       "wrong user agent",
+			userAgent:  "wrong-user-agent",
+			statusCode: http.StatusUnauthorized,
 		},
 	}
-	tokenServer := &AuxiliaryTokenServer{}
+	tokenServer := &AuxiliaryTokenServer{Token: defaultToken}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			clientOpts := armopts.DefaultArmOpts()
-			vmClientOpts := *clientOpts
-			auxPolicy, err := auth.NewAuxiliaryTokenPolicy(context.Background(), tokenServer, tt.url, tt.scope)
-			if (err != nil) != tt.wantErr {
-				t.Errorf("getAuxiliaryToken() error = %v, wantErr: %v", err, tt.wantErr)
-				return
+			request := &http.Request{
+				Method: http.MethodGet,
+				URL:    &url.URL{Path: "/"},
+				Header: http.Header{
+					"User-Agent": []string{tt.userAgent},
+				},
 			}
-			vmClientOpts.ClientOptions.PerRetryPolicies = append(vmClientOpts.ClientOptions.PerRetryPolicies, auxPolicy)
-			if tt.wantErr {
-				assert.ErrorContains(t, err, tt.errString)
-			} else {
-				assert.NotEqual(t, clientOpts.ClientOptions.PerRetryPolicies, vmClientOpts.ClientOptions.PerRetryPolicies)
+			resp, err := tokenServer.Do(request)
+			if err != nil {
+				t.Errorf("Unexpected error %v", err)
+				return
 			}
+			assert.Equal(t, tt.statusCode, resp.StatusCode, "Expected status code %d, got %d", tt.statusCode, resp.StatusCode)
+			tokenServer.Reset()
 		})
-		tokenServer.Reset()
 	}
 }

pkg/fake/virtualmachinesapi.go

@@ -29,6 +29,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute"
+	"github.com/Azure/karpenter-provider-azure/pkg/auth"
 	"github.com/Azure/karpenter-provider-azure/pkg/providers/instance"
 	"github.com/samber/lo"
 )
@@ -74,6 +75,7 @@ type VirtualMachinesAPI struct {
 	// TODO: document the implications of embedding vs. not embedding the interface here
 	// instance.VirtualMachinesAPI // - this is the interface we are mocking.
 	VirtualMachinesBehavior
+	AuxiliaryTokenPolicy *auth.AuxiliaryTokenPolicy
 }
 
 // Reset must be called between tests otherwise tests will pollute each other.
@@ -88,7 +90,24 @@ func (c *VirtualMachinesAPI) Reset() {
 	})
 }
 
-func (c *VirtualMachinesAPI) BeginCreateOrUpdate(_ context.Context, resourceGroupName string, vmName string, parameters armcompute.VirtualMachine, options *armcompute.VirtualMachinesClientBeginCreateOrUpdateOptions) (*runtime.Poller[armcompute.VirtualMachinesClientCreateOrUpdateResponse], error) {
+// UseAuxiliaryTokenPolicy simulates AuxiliaryTokenPolicy.Do() method being called at the beginning of each API call
+// This is useful for testing scenarios where the auxiliary token is required for the API call to succeed.
+// If the AuxiliaryTokenPolicy is not set (USE_SIG: false), this method does nothing and returns nil.
+func (c *VirtualMachinesAPI) UseAuxiliaryTokenPolicy() error {
+	if c.AuxiliaryTokenPolicy != nil {
+		request, _ := runtime.NewRequest(context.Background(), "GET", "http://example.com")
+		if _, err := c.AuxiliaryTokenPolicy.Do(request); err != nil {
+			// req.Next() returns this if there are no more policies.
+			if err.Error() == "no more policies" {
+				return nil
+			}
+			return err
+		}
+	}
+	return nil
+}
+
+func (c *VirtualMachinesAPI) BeginCreateOrUpdate(ctx context.Context, resourceGroupName string, vmName string, parameters armcompute.VirtualMachine, options *armcompute.VirtualMachinesClientBeginCreateOrUpdateOptions) (*runtime.Poller[armcompute.VirtualMachinesClientCreateOrUpdateResponse], error) {
 	// gather input parameters (may get rid of this with multiple mocked function signatures to reflect common patterns)
 	input := &VirtualMachineCreateOrUpdateInput{
 		ResourceGroupName: resourceGroupName,
@@ -99,6 +118,9 @@ func (c *VirtualMachinesAPI) BeginCreateOrUpdate(_ context.Context, resourceGrou
 	// BeginCreateOrUpdate should fail, if the vm exists in the cache, and we are attempting to change properties for zone
 
 	return c.VirtualMachineCreateOrUpdateBehavior.Invoke(input, func(input *VirtualMachineCreateOrUpdateInput) (*armcompute.VirtualMachinesClientCreateOrUpdateResponse, error) {
+		if err := c.UseAuxiliaryTokenPolicy(); err != nil {
+			return nil, getAuthTokenError(err)
+		}
 		//if input.ResourceGroupName == "" {
 		//	return nil, errors.New("ResourceGroupName is required")
 		//}
@@ -160,6 +182,9 @@ func (c *VirtualMachinesAPI) BeginUpdate(_ context.Context, resourceGroupName st
 		Options:           options,
 	}
 	return c.VirtualMachineUpdateBehavior.Invoke(input, func(input *VirtualMachineUpdateInput) (*armcompute.VirtualMachinesClientUpdateResponse, error) {
+		if err := c.UseAuxiliaryTokenPolicy(); err != nil {
+			return nil, getAuthTokenError(err)
+		}
 		id := MkVMID(input.ResourceGroupName, input.VMName)
 
 		instance, ok := c.Instances.Load(id)
@@ -203,6 +228,9 @@ func (c *VirtualMachinesAPI) Get(_ context.Context, resourceGroupName string, vm
 		Options:           options,
 	}
 	return c.VirtualMachineGetBehavior.Invoke(input, func(input *VirtualMachineGetInput) (armcompute.VirtualMachinesClientGetResponse, error) {
+		if err := c.UseAuxiliaryTokenPolicy(); err != nil {
+			return armcompute.VirtualMachinesClientGetResponse{}, getAuthTokenError(err)
+		}
 		instance, ok := c.Instances.Load(MkVMID(input.ResourceGroupName, input.VMName))
 		if !ok {
 			return armcompute.VirtualMachinesClientGetResponse{}, &azcore.ResponseError{ErrorCode: errors.ResourceNotFound}
@@ -220,6 +248,9 @@ func (c *VirtualMachinesAPI) BeginDelete(_ context.Context, resourceGroupName st
 		Options:           options,
 	}
 	return c.VirtualMachineDeleteBehavior.Invoke(input, func(input *VirtualMachineDeleteInput) (*armcompute.VirtualMachinesClientDeleteResponse, error) {
+		if err := c.UseAuxiliaryTokenPolicy(); err != nil {
+			return &armcompute.VirtualMachinesClientDeleteResponse{}, getAuthTokenError(err)
+		}
 		c.Instances.Delete(MkVMID(input.ResourceGroupName, input.VMName))
 		return &armcompute.VirtualMachinesClientDeleteResponse{}, nil
 	})
@@ -233,3 +264,12 @@ func MkVMID(resourceGroupName string, vmName string) string {
 	const idFormat = "/subscriptions/subscriptionID/resourceGroups/%s/providers/Microsoft.Compute/virtualMachines/%s"
 	return fmt.Sprintf(idFormat, resourceGroupName, vmName)
 }
+
+func getAuthTokenError(err error) *azcore.ResponseError {
+	return &azcore.ResponseError{
+		ErrorCode: "AuthenticationFailed",
+		RawResponse: &http.Response{
+			Body: createSDKErrorBody("AuthenticationFailed", err.Error()),
+		},
+	}
+}

pkg/providers/imagefamily/nodeimage_test.go

@@ -91,6 +91,7 @@ func getExpectedTestSIGImages(imageFamily string, version string, kubernetesVers
 
 var _ = Describe("NodeImageProvider tests", func() {
 	var (
+		testOptions               *options.Options
 		communityImageVersionsAPI *fake.CommunityGalleryImageVersionsAPI
 
 		nodeImageProvider imagefamily.NodeImageProvider
@@ -100,7 +101,8 @@ var _ = Describe("NodeImageProvider tests", func() {
 
 	BeforeEach(func() {
 		ctx = coreoptions.ToContext(ctx, coretest.Options())
-		ctx = options.ToContext(ctx, test.Options())
+		testOptions = test.Options()
+		ctx = options.ToContext(ctx, testOptions)
 
 		communityImageVersionsAPI = &fake.CommunityGalleryImageVersionsAPI{}
 		cigImageVersionTest := cigImageVersion
@@ -110,7 +112,7 @@ var _ = Describe("NodeImageProvider tests", func() {
 		kubernetesVersion = lo.Must(env.KubernetesInterface.Discovery().ServerVersion()).String()
 
 		nodeClass = test.AKSNodeClass()
-		test.ApplyDefaultStatus(nodeClass, env)
+		test.ApplyDefaultStatus(nodeClass, env, testOptions.UseSIG)
 	})
 
 	Context("List CIG Images", func() {
@@ -161,7 +163,7 @@ var _ = Describe("NodeImageProvider tests", func() {
 
 	Context("List SIG Images", func() {
 		BeforeEach(func() {
-			testOptions := options.FromContext(ctx)
+			testOptions = options.FromContext(ctx)
 			testOptions.UseSIG = true
 			testOptions.SIGSubscriptionID = sigSubscription
 			testOptions.SIGAccessTokenScope = "http://valid-scope.com/.default"