feat(timeout): Implement customizable timeout for Azure CLI token ret…

pkg/token/azurecli.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"strconv"
+	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
@@ -15,18 +16,26 @@ import (
 type AzureCLIToken struct {
 	resourceID string
 	tenantID   string
+	timeout    time.Duration
 }
 
+const defaultTimeout = 30 * time.Second
+
 // newAzureCLIToken returns a TokenProvider that will fetch a token for the user currently logged into the Azure CLI.
 // Required arguments include an oAuthConfiguration object and the resourceID (which is used as the scope)
-func newAzureCLIToken(resourceID string, tenantID string) (TokenProvider, error) {
+func newAzureCLIToken(resourceID string, tenantID string, timeout time.Duration) (TokenProvider, error) {
 	if resourceID == "" {
 		return nil, errors.New("resourceID cannot be empty")
 	}
 
+	if timeout <= 0 {
+		timeout = defaultTimeout
+	}
+
 	return &AzureCLIToken{
 		resourceID: resourceID,
 		tenantID:   tenantID,
+		timeout:    timeout,
 	}, nil
 }
 
@@ -42,11 +51,15 @@ func (p *AzureCLIToken) Token() (adal.Token, error) {
 		return emptyToken, fmt.Errorf("unable to create credential. Received: %v", err)
 	}
 
-	// Use the token provider to get a new token
-	cliAccessToken, err := cred.GetToken(context.Background(), policy.TokenRequestOptions{Scopes: []string{p.resourceID}})
+	ctx, cancel := context.WithTimeout(context.Background(), p.timeout)
+	defer cancel()
+
+	// Use the token provider to get a new token with the new context
+	cliAccessToken, err := cred.GetToken(ctx, policy.TokenRequestOptions{Scopes: []string{p.resourceID}})
 	if err != nil {
 		return emptyToken, fmt.Errorf("expected an empty error but received: %v", err)
 	}
+
 	if cliAccessToken.Token == "" {
 		return emptyToken, errors.New("did not receive a token")
 	}

pkg/token/azurecli_test.go

@@ -7,7 +7,8 @@ import (
 )
 
 func TestNewAzureCLITokenEmpty(t *testing.T) {
-	_, err := newAzureCLIToken("", "")
+	// Using default timeout for testing
+	_, err := newAzureCLIToken("", "", defaultTimeout)
 
 	if !testutils.ErrorContains(err, "resourceID cannot be empty") {
 		t.Errorf("unexpected error: %v", err)
@@ -22,3 +23,29 @@ func TestNewAzureCLIToken(t *testing.T) {
 		t.Errorf("unexpected error: %v", err)
 	}
 }
+
+func TestOptionsValidateTimeout(t *testing.T) {
+	// Assuming "devicecode" is a valid login method based on your error message
+	validLoginMethod := "devicecode"
+
+	// Test with Timeout set to 0
+	options := Options{LoginMethod: validLoginMethod, Timeout: 0}
+	err := options.Validate()
+	if !testutils.ErrorContains(err, "timeout must be greater than 0") {
+		t.Errorf("expected timeout error, got: %v", err)
+	}
+
+	// Test with Timeout set to a negative value
+	options = Options{LoginMethod: validLoginMethod, Timeout: -1}
+	err = options.Validate()
+	if !testutils.ErrorContains(err, "timeout must be greater than 0") {
+		t.Errorf("expected timeout error, got: %v", err)
+	}
+
+	// Test with a valid Timeout value
+	options = Options{LoginMethod: validLoginMethod, Timeout: 10}
+	err = options.Validate()
+	if err != nil {
+		t.Errorf("unexpected error for valid timeout: %v", err)
+	}
+}

pkg/token/options.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 
 	"github.com/spf13/pflag"
 	"k8s.io/client-go/util/homedir"
@@ -22,6 +23,7 @@ type Options struct {
 	TenantID               string
 	Environment            string
 	IsLegacy               bool
+	Timeout                time.Duration
 	TokenCacheDir          string
 	tokenCacheFile         string
 	IdentityResourceID     string
@@ -121,6 +123,8 @@ func (o *Options) AddFlags(fs *pflag.FlagSet) {
 	fs.BoolVar(&o.UseAzureRMTerraformEnv, "use-azurerm-env-vars", o.UseAzureRMTerraformEnv,
 		"Use environment variable names of Terraform Azure Provider (ARM_CLIENT_ID, ARM_CLIENT_SECRET, ARM_CLIENT_CERTIFICATE_PATH, ARM_CLIENT_CERTIFICATE_PASSWORD, ARM_TENANT_ID)")
 	fs.BoolVar(&o.IsPoPTokenEnabled, "pop-enabled", o.IsPoPTokenEnabled, "set to true to use a PoP token for authentication or false to use a regular bearer token")
+	fs.DurationVar(&o.Timeout, "timeout", 30*time.Second,
+		fmt.Sprintf("Timeout duration for Azure CLI token requests. It may be specified in %s environment variable", "AZURE_CLI_TIMEOUT"))
 	fs.StringVar(&o.PoPTokenClaims, "pop-claims", o.PoPTokenClaims, "contains a comma-separated list of claims to attach to the pop token in the format `key=val,key2=val2`. At minimum, specify the ARM ID of the cluster as `u=ARM_ID`")
 }
 
@@ -145,6 +149,10 @@ func (o *Options) Validate() error {
 		return fmt.Errorf("pop-enabled flag is required to use the PoP token feature. Please provide both pop-enabled and pop-claims flags")
 	}
 
+	if o.Timeout <= 0 {
+		return fmt.Errorf("timeout must be greater than 0")
+	}
+
 	return nil
 }
 
@@ -224,18 +232,24 @@ func (o *Options) UpdateFromEnv() {
 			o.AuthorityHost = v
 		}
 	}
+	if v, ok := os.LookupEnv("AZURE_CLI_TIMEOUT"); ok {
+		if timeout, err := time.ParseDuration(v); err == nil {
+			o.Timeout = timeout
+		}
+	}
 }
 
 func (o *Options) ToString() string {
 	azureConfigDir := os.Getenv("AZURE_CONFIG_DIR")
-	return fmt.Sprintf("Login Method: %s, Environment: %s, TenantID: %s, ServerID: %s, ClientID: %s, IsLegacy: %t, msiResourceID: %s, tokenCacheDir: %s, tokenCacheFile: %s, AZURE_CONFIG_DIR: %s",
+	return fmt.Sprintf("Login Method: %s, Environment: %s, TenantID: %s, ServerID: %s, ClientID: %s, IsLegacy: %t, msiResourceID: %s, Timeout: %v, tokenCacheDir: %s, tokenCacheFile: %s, AZURE_CONFIG_DIR: %s",
 		o.LoginMethod,
 		o.Environment,
 		o.TenantID,
 		o.ServerID,
 		o.ClientID,
 		o.IsLegacy,
 		o.IdentityResourceID,
+		o.Timeout,
 		o.TokenCacheDir,
 		o.tokenCacheFile,
 		azureConfigDir,

pkg/token/options_test.go

@@ -5,6 +5,7 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/Azure/kubelogin/pkg/testutils"
 	"github.com/google/go-cmp/cmp"
@@ -104,6 +105,7 @@ func TestOptionsWithEnvVars(t *testing.T) {
 				TenantID:           tenantID,
 				LoginMethod:        DeviceCodeLogin,
 				tokenCacheFile:     "---.json",
+				Timeout:            30 * time.Second,
 			},
 		},
 		{
@@ -126,6 +128,7 @@ func TestOptionsWithEnvVars(t *testing.T) {
 				TenantID:               tenantID,
 				LoginMethod:            DeviceCodeLogin,
 				tokenCacheFile:         "---.json",
+				Timeout:                30 * time.Second,
 			},
 		},
 		{
@@ -154,6 +157,7 @@ func TestOptionsWithEnvVars(t *testing.T) {
 				AuthorityHost:      authorityHost,
 				FederatedTokenFile: tokenFile,
 				tokenCacheFile:     "---.json",
+				Timeout:            30 * time.Second,
 			},
 		},
 	}

pkg/token/provider.go

@@ -43,7 +43,7 @@ func newTokenProvider(o *Options) (TokenProvider, error) {
 	case MSILogin:
 		return newManagedIdentityToken(o.ClientID, o.IdentityResourceID, o.ServerID)
 	case AzureCLILogin:
-		return newAzureCLIToken(o.ServerID, o.TenantID)
+		return newAzureCLIToken(o.ServerID, o.TenantID, o.Timeout)
 	case WorkloadIdentityLogin:
 		return newWorkloadIdentityToken(o.ClientID, o.FederatedTokenFile, o.AuthorityHost, o.ServerID, o.TenantID)
 	}