refactor(client): allow specifying RPC deadlines and streamline retry logic (#162)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: cameronmeissner <24923771+cameronmeissner@users.noreply.github.com>

Author: cameronmeissner

.pipelines/client/e2e.yaml

@@ -82,7 +82,7 @@ stages:
           ADO_PAT: $(PAT-aksdevassistant)
           ORGANIZATION: $(ORGANIZATION)
           PROJECT: $(PROJECT)
-        condition: succeeded()
+        condition: and(succeeded(), ne(variables.SKIP_E2E, 'True'))
         displayName: ADO Login
 
       - bash: /bin/bash .pipelines/client/scripts/run-pipeline.sh
@@ -92,7 +92,7 @@ stages:
           ADO_ORGANIZATION: $(ORGANIZATION)
           ADO_PROJECT: $(PROJECT)
           SUITE_ID: $(SECURE_TLS_BOOTSTRAPPING_CHECKIN_SUITE_ID)
-          ADDITIONAL_VARS: "VSTS_TOGGLE_OVERRIDES=$(ENABLE_SECURE_TLS_BOOTSTRAPPING_TOGGLE_SETTINGS),$(TOGGLE_OVERRIDES)"
-        condition: succeeded()
+          ADDITIONAL_VARS: "$(TESTS_TO_RUN) VSTS_TOGGLE_OVERRIDES=$(ENABLE_SECURE_TLS_BOOTSTRAPPING_TOGGLE_SETTINGS),$(TOGGLE_OVERRIDES)"
+        condition: and(succeeded(), ne(variables.SKIP_E2E, 'True'))
         displayName: Run Secure TLS Bootstrapping Check-in Tests
         timeoutInMinutes: 0

client/cmd/client/main.go

@@ -15,7 +15,6 @@ import (
 
 	"github.com/Azure/aks-secure-tls-bootstrap/client/internal/bootstrap"
 	"github.com/Azure/aks-secure-tls-bootstrap/client/internal/build"
-	"github.com/Azure/aks-secure-tls-bootstrap/client/internal/kubeconfig"
 	"github.com/Azure/aks-secure-tls-bootstrap/client/internal/log"
 	"github.com/Azure/aks-secure-tls-bootstrap/client/internal/telemetry"
 	"go.uber.org/zap"
@@ -45,7 +44,13 @@ func init() {
 	flag.StringVar(&config.TLSMinVersion, "tls-min-version", "", "the minimum TLS version used to communicate with control plane")
 	flag.BoolVar(&config.InsecureSkipTLSVerify, "insecure-skip-tls-verify", false, "skip TLS verification when connecting to the control plane")
 	flag.BoolVar(&config.EnsureAuthorizedClient, "ensure-authorized", false, "ensure the specified kubeconfig contains an authorized clientset before bootstrapping")
-	flag.DurationVar(&config.Deadline, "deadline", 0, "the deadline within which bootstrapping must succeed")
+	flag.DurationVar(&config.ValidateKubeconfigTimeout, "validate-kubeconfig-timeout", 0, "timeout applied to existing kubeconfig validation")
+	flag.DurationVar(&config.GetAccessTokenTimeout, "get-access-token-timeout", 0, "timeout applied to the get access token RPC")
+	flag.DurationVar(&config.GetInstanceDataTimeout, "get-instance-data-timeout", 0, "timeout applied to the get instance data RPC")
+	flag.DurationVar(&config.GetNonceTimeout, "get-nonce-timeout", 0, "timeout applied to the get nonce RPC")
+	flag.DurationVar(&config.GetAttestedDataTimeout, "get-attested-data-timeout", 0, "timeout applied to the get attested data RPC")
+	flag.DurationVar(&config.GetCredentialTimeout, "get-credential-timeout", 0, "timeout applied to the get credential RPC")
+	flag.DurationVar(&config.Deadline, "deadline", 0, "the deadline within which bootstrapping must succeed - DEPRECATED, use RPC timeouts instead")
 
 	flag.Usage = func() {
 		_, _ = fmt.Fprintf(os.Stderr, "Usage of %s - %s:\n", os.Args[0], build.GetVersion())
@@ -83,61 +88,45 @@ func run(ctx context.Context) int {
 
 	ctx = log.WithLogger(telemetry.WithTracing(ctx), logger)
 
-	var endTime time.Time
+	logger.Info("running with config", config.ToZapFields()...)
+
+	var bootstrapErr error
+	var startTime, endTime time.Time
 	result := &bootstrap.Result{
 		Status: bootstrap.StatusSuccess,
 	}
-
-	startTime := time.Now()
-	deadline := startTime.Add(config.Deadline)
-	bootstrapCtx, cancel := context.WithDeadline(ctx, deadline)
-	defer cancel()
-	logger.Info("set bootstrap deadline", zap.Time("deadline", deadline))
-
-	kubeconfigPath := config.KubeconfigPath
-	err = kubeconfig.NewValidator().Validate(bootstrapCtx, kubeconfigPath, config.EnsureAuthorizedClient)
-	if err == nil {
-		logger.Info("existing kubeconfig is valid, will not bootstrap a new kubelet client credential", zap.String("kubeconfig", kubeconfigPath))
-		endTime = time.Now()
+	defer func() {
+		if bootstrapErr != nil {
+			result.Status = bootstrap.StatusFailure
+			result.FinalErrorType = bootstrap.GetErrorType(bootstrapErr)
+			result.FinalError = bootstrapErr.Error()
+		}
+		result.Trace = telemetry.GetTrace(ctx)
 		emitGuestAgentEvent(logger, startTime, endTime, result)
-		return 0
-	}
-	logger.Info("failed to validate existing kubeconfig, will bootstrap a new kubelet client credential", zap.String("kubeconfig", kubeconfigPath), zap.Error(err))
+	}()
 
-	errLog, traces, err := bootstrap.Bootstrap(bootstrapCtx, config)
+	startTime = time.Now()
+	bootstrapErr = bootstrap.Bootstrap(ctx, config)
 	endTime = time.Now()
-	result.Errors = errLog
-	result.Traces = traces.GetLastNTraces(5) // only keep the last 5 traces to avoid truncating guest agent event data
-	result.TraceSummary = traces.GetTraceSummary()
 
 	var exitCode int
-	if err != nil {
-		result.Status = bootstrap.StatusFailure
+	if bootstrapErr != nil {
 		switch {
-		case errors.Is(err, context.Canceled):
-			logger.Error("context was cancelled before bootstrapping could complete")
-		case errors.Is(err, context.DeadlineExceeded):
-			err = errors.Unwrap(err)
-			logger.Error(
-				"failed to successfully bootstrap before the specified deadline",
-				zap.Error(err),
-				zap.Time("deadline", deadline),
-				zap.Duration("deadlineDuration", config.Deadline),
-			)
+		case errors.Is(bootstrapErr, context.Canceled):
+			logger.Error("context was canceled before bootstrapping could complete")
+		case errors.Is(bootstrapErr, context.DeadlineExceeded):
+			logger.Error("failed to bootstrap due to exceeding context deadline", zap.Error(bootstrapErr))
 		default:
-			logger.Error("failed to bootstrap", zap.Error(err))
+			logger.Error("failed to bootstrap", zap.Error(bootstrapErr))
 		}
-		result.FinalError = err.Error()
 		exitCode = 1
 	}
 
-	emitGuestAgentEvent(logger, startTime, endTime, result)
 	return exitCode
 }
 
 func emitGuestAgentEvent(logger *zap.Logger, startTime, endTime time.Time, result *bootstrap.Result) {
 	result.ElapsedMilliseconds = endTime.Sub(startTime).Milliseconds()
-
 	bootstrapEvent := &bootstrap.Event{
 		Start: startTime,
 		End:   endTime,

client/go.mod

@@ -6,7 +6,6 @@ require (
 	github.com/Azure/aks-secure-tls-bootstrap/service v1.0.2
 	github.com/Azure/go-autorest/autorest v0.11.29
 	github.com/Azure/go-autorest/autorest/adal v0.9.22
-	github.com/avast/retry-go/v4 v4.6.1
 	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.2
 	github.com/hashicorp/go-retryablehttp v0.7.7
 	github.com/stretchr/testify v1.10.0

client/go.sum

@@ -19,8 +19,6 @@ github.com/Azure/go-autorest/logger v0.2.1 h1:IG7i4p/mDa2Ce4TRyAO8IHnVhAVF3RFU+Z
 github.com/Azure/go-autorest/logger v0.2.1/go.mod h1:T9E3cAhj2VqvPOtCYAvby9aBXkZmbF5NWuPV8+WeEW8=
 github.com/Azure/go-autorest/tracing v0.6.0 h1:TYi4+3m5t6K48TGI9AUdb+IzbnSxvnvUMfuitfgcfuo=
 github.com/Azure/go-autorest/tracing v0.6.0/go.mod h1:+vhtPC754Xsa23ID7GlGsrdKBpUA79WCAKPPZVC2DeU=
-github.com/avast/retry-go/v4 v4.6.1 h1:VkOLRubHdisGrHnTu89g08aQEWEgRU7LVEop3GbIcMk=
-github.com/avast/retry-go/v4 v4.6.1/go.mod h1:V6oF8njAwxJ5gRo1Q7Cxab24xs5NCWZBeaHHBklR8mA=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=

client/hack/linux/install.sh

@@ -2,7 +2,7 @@
 set -euxo pipefail
 
 # this script can be used to install an arbitrary version of aks-secure-tls-bootstrap-client
-# on a running AKS node for development/testing.
+# on a running AKS Linux node for development/testing.
 
 # download and usage:
 # 1. $ curl -o install-aks-secure-tls-bootstrap-client.sh https://raw.githubusercontent.com/Azure/aks-secure-tls-bootstrap/refs/heads/main/client/hack/linux/install.sh
@@ -14,12 +14,18 @@ VERSION="${VERSION:-}"
 [ -z "$VERSION" ] && echo "VERSION must be specified" && exit 1
 [ -z "$STORAGE_ACCOUNT_NAME" ] && echo "STORAGE_ACCOUNT_NAME must be specified" && exit 1
 
-curl -fsSL https://${STORAGE_ACCOUNT_NAME}.z22.web.core.windows.net/client/linux/amd64/${VERSION} -o linux-amd64.tar.gz
+curl -fsSL https://${STORAGE_ACCOUNT_NAME}.z22.web.core.windows.net/client/linux/amd64/${VERSION}.tar.gz -o linux-amd64.tar.gz
 mkdir -p client
 tar -xvzf linux-amd64.tar.gz -C client
-rm /usr/local/bin/aks-secure-tls-bootstrap-client
 chmod +x client/aks-secure-tls-bootstrap-client
-mv client/aks-secure-tls-bootstrap-client /usr/local/bin/aks-secure-tls-bootstrap-client
+
+rm -f /usr/local/bin/aks-secure-tls-bootstrap-client
+rm -f /opt/bin/aks-secure-tls-bootstrap-client
+mv client/aks-secure-tls-bootstrap-client /opt/bin/aks-secure-tls-bootstrap-client
+cp /opt/bin/aks-secure-tls-bootstrap-client /usr/local/bin/aks-secure-tls-bootstrap-client
+
 rm -rf client
-rm linux-amd64.tar.gz
-stat /usr/local/bin/aks-secure-tls-bootstrap-client
+rm -f linux-amd64.tar.gz
+
+stat /opt/bin/aks-secure-tls-bootstrap-client
+/opt/bin/aks-secure-tls-bootstrap-client -h

client/hack/windows/install.ps1

@@ -17,7 +17,7 @@ param(
 
 Write-Host "Downloading aks-secure-tls-bootstrap-client version $Version from storage account $StorageAccountName"
 
-$downloadUrl = "https://$StorageAccountName.z22.web.core.windows.net/client/windows/amd64/$Version"
+$downloadUrl = "https://$StorageAccountName.z22.web.core.windows.net/client/windows/amd64/$Version.zip"
 $archivePath = "windows-amd64.zip"
 
 Write-Host "Downloading from: $downloadUrl"

client/internal/bootstrap/auth.go

@@ -7,75 +7,73 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
-	"net/http"
 	"strings"
 
 	"github.com/Azure/aks-secure-tls-bootstrap/client/internal/cloud"
-	"github.com/Azure/aks-secure-tls-bootstrap/client/internal/imds"
 	"github.com/Azure/aks-secure-tls-bootstrap/client/internal/log"
-	"github.com/Azure/aks-secure-tls-bootstrap/client/internal/telemetry"
 	"github.com/Azure/go-autorest/autorest/adal"
 	"github.com/Azure/go-autorest/autorest/azure"
 	"go.uber.org/zap"
 )
 
 const (
+	// service principal secrets containing this prefix are PFX certificates which need to be decoded,
+	// rather than raw password / secret strings
 	certificateSecretPrefix = "certificate:"
 )
 
 const (
+	// this will be the exact value of the "userAssignedIdentityID" field of the cloud provider config (azure.json)
+	// when the node is using a (user-assigned) managed identity, rather than a service principal
 	clientIDForMSI = "msi"
 )
 
 const (
-	maxMSIRefreshAttempts = 3
+	// to avoid falling too deep into exponential backoff implemented by adal, which follows the public retry guidance:
+	// https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/how-to-use-vm-token#retry-guidance
+	maxMSIRefreshAttempts = 5
 )
 
 // extractAccessTokenFunc extracts an oauth access token from the specified service principal token after a refresh, fake implementations given in unit tests.
-type extractAccessTokenFunc func(token *adal.ServicePrincipalToken, isMSI bool) (string, error)
+type extractAccessTokenFunc func(ctx context.Context, token *adal.ServicePrincipalToken, isMSI bool) (string, error)
 
-func extractAccessToken(token *adal.ServicePrincipalToken, isMSI bool) (string, error) {
-	if err := token.Refresh(); err != nil {
-		return "", tokenRefreshErrorToGetAccessTokenFailure(err, isMSI)
+func extractAccessToken(ctx context.Context, token *adal.ServicePrincipalToken, isMSI bool) (string, error) {
+	if err := token.RefreshWithContext(ctx); err != nil {
+		return "", err
 	}
 	return token.OAuthToken(), nil
 }
 
-func (c *client) getAccessToken(ctx context.Context, userAssignedIdentityID, resource string, cloudProviderConfig *cloud.ProviderConfig) (string, error) {
-	endSpan := telemetry.StartSpan(ctx, "GetAccessToken")
-	defer endSpan()
-
+func (c *client) getToken(ctx context.Context, config *Config) (string, error) {
 	logger := log.MustGetLogger(ctx)
 
-	userAssignedID := cloudProviderConfig.UserAssignedIdentityID
-	if userAssignedIdentityID != "" {
-		userAssignedID = userAssignedIdentityID
+	userAssignedID := config.CloudProviderConfig.UserAssignedIdentityID
+	if config.UserAssignedIdentityID != "" {
+		userAssignedID = config.UserAssignedIdentityID
 	}
 
 	if userAssignedID != "" {
 		logger.Info("generating MSI access token", zap.String("clientId", userAssignedID))
-		msiToken, err := adal.NewServicePrincipalTokenFromManagedIdentity(resource, &adal.ManagedIdentityOptions{
+		msiToken, err := adal.NewServicePrincipalTokenFromManagedIdentity(config.AADResource, &adal.ManagedIdentityOptions{
 			ClientID: userAssignedID,
 		})
 		if err != nil {
-			return "", makeNonRetryableGetAccessTokenFailure(fmt.Errorf("generating MSI access token: %w", err))
+			return "", fmt.Errorf("generating MSI access token: %w", err)
 		}
-		// to avoid falling too deep into exponential backoff implemented by adal, which follows the public retry guidance:
-		// https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/how-to-use-vm-token#retry-guidance
 		msiToken.MaxMSIRefreshAttempts = maxMSIRefreshAttempts
-		return c.extractAccessTokenFunc(msiToken, true)
+		return c.extractAccessTokenFunc(ctx, msiToken, true)
 	}
 
-	if cloudProviderConfig.ClientID == clientIDForMSI {
-		return "", makeNonRetryableGetAccessTokenFailure(fmt.Errorf("client ID within cloud provider config indicates usage of a managed identity, though no user-assigned identity ID was provided"))
+	if config.CloudProviderConfig.ClientID == clientIDForMSI {
+		return "", fmt.Errorf("client ID within cloud provider config indicates usage of a managed identity, though no user-assigned identity ID was provided")
 	}
 
-	servicePrincipalToken, err := getServicePrincipalToken(ctx, resource, cloudProviderConfig)
+	servicePrincipalToken, err := getServicePrincipalToken(ctx, config.AADResource, config.CloudProviderConfig)
 	if err != nil {
 		return "", err
 	}
 
-	return c.extractAccessTokenFunc(servicePrincipalToken, false)
+	return c.extractAccessTokenFunc(ctx, servicePrincipalToken, false)
 }
 
 func getServicePrincipalToken(ctx context.Context, resource string, cloudProviderConfig *cloud.ProviderConfig) (*adal.ServicePrincipalToken, error) {
@@ -85,18 +83,18 @@ func getServicePrincipalToken(ctx context.Context, resource string, cloudProvide
 
 	env, err := azure.EnvironmentFromName(cloudProviderConfig.CloudName)
 	if err != nil {
-		return nil, makeNonRetryableGetAccessTokenFailure(fmt.Errorf("getting azure environment config for cloud %q: %w", cloudProviderConfig.CloudName, err))
+		return nil, fmt.Errorf("getting azure environment config for cloud %q: %w", cloudProviderConfig.CloudName, err)
 	}
 	oauthConfig, err := adal.NewOAuthConfig(env.ActiveDirectoryEndpoint, cloudProviderConfig.TenantID)
 	if err != nil {
-		return nil, makeNonRetryableGetAccessTokenFailure(fmt.Errorf("creating oauth config with azure environment: %w", err))
+		return nil, fmt.Errorf("creating oauth config with azure environment: %w", err)
 	}
 
 	if !strings.HasPrefix(secret, certificateSecretPrefix) {
 		logger.Info("generating service principal access token with client secret", zap.String("clientId", cloudProviderConfig.ClientID))
 		token, err := adal.NewServicePrincipalToken(*oauthConfig, cloudProviderConfig.ClientID, secret, resource)
 		if err != nil {
-			return nil, makeNonRetryableGetAccessTokenFailure(fmt.Errorf("generating service principal access token with client secret: %w", err))
+			return nil, fmt.Errorf("generating service principal access token with client secret: %w", err)
 		}
 		return token, nil
 	}
@@ -105,17 +103,17 @@ func getServicePrincipalToken(ctx context.Context, resource string, cloudProvide
 
 	certData, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(secret, certificateSecretPrefix))
 	if err != nil {
-		return nil, makeNonRetryableGetAccessTokenFailure(fmt.Errorf("b64-decoding certificate data in client secret: %w", err))
+		return nil, fmt.Errorf("b64-decoding certificate data in client secret: %w", err)
 	}
 	certificate, privateKey, err := adal.DecodePfxCertificateData(certData, "")
 	if err != nil {
-		return nil, makeNonRetryableGetAccessTokenFailure(fmt.Errorf("decoding pfx certificate data in client secret: %w", err))
+		return nil, fmt.Errorf("decoding pfx certificate data in client secret: %w", err)
 	}
 
 	logger.Info("generating service principal access token with certificate", zap.String("clientId", cloudProviderConfig.ClientID))
 	token, err := adal.NewServicePrincipalTokenFromCertificate(*oauthConfig, cloudProviderConfig.ClientID, certificate, privateKey, resource)
 	if err != nil {
-		return nil, makeNonRetryableGetAccessTokenFailure(fmt.Errorf("generating service principal access token with certificate: %w", err))
+		return nil, fmt.Errorf("generating service principal access token with certificate: %w", err)
 	}
 
 	return token, nil
@@ -127,44 +125,3 @@ func maybeB64Decode(str string) string {
 	}
 	return str
 }
-
-func makeNonRetryableGetAccessTokenFailure(err error) error {
-	return &bootstrapError{
-		errorType: ErrorTypeGetAccessTokenFailure,
-		retryable: false,
-		inner:     err,
-	}
-}
-
-func tokenRefreshErrorToGetAccessTokenFailure(err error, isMSI bool) error {
-	bootstrapErr := &bootstrapError{
-		errorType: ErrorTypeGetAccessTokenFailure,
-		retryable: true, // optimistically consider the error retryable from the start
-		inner:     fmt.Errorf("obtaining fresh access token: %w", err),
-	}
-
-	rerr, ok := err.(adal.TokenRefreshError)
-	if !ok {
-		return bootstrapErr
-	}
-
-	resp := rerr.Response()
-	if resp == nil {
-		return bootstrapErr
-	}
-
-	if !isMSI {
-		bootstrapErr.retryable = resp.StatusCode >= http.StatusInternalServerError
-		return bootstrapErr
-	}
-
-	if resp.StatusCode != http.StatusBadRequest {
-		bootstrapErr.retryable = imds.IsRetryableHTTPStatusCode(resp.StatusCode)
-		return bootstrapErr
-	}
-
-	// 400s aren't normally retryable, though identity assignment can sometimes take a bit of time to propagate to IMDS,
-	// so we treat "Identity not found" errors as retryable
-	bootstrapErr.retryable = strings.Contains(strings.ToLower(err.Error()), "identity not found")
-	return bootstrapErr
-}