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

…rieval in kubelogin

This commit introduces the ability to specify a custom timeout for Azure CLI token retrieval within the kubelogin package. The `AzureCLIToken` struct now includes a `timeout` field, allowing users to set a specific timeout duration.

main.go

@@ -3,18 +3,27 @@ package main
 import (
 	"flag"
 	"os"
+	"time"
 
 	"github.com/Azure/kubelogin/pkg/cmd"
 	"github.com/spf13/pflag"
 	klog "k8s.io/klog/v2"
 )
 
+var timeoutDuration time.Duration
+
 func main() {
 	klog.InitFlags(nil)
+
+	flag.DurationVar(&timeoutDuration, "timeout", 10*time.Second, "Timeout duration for Azure CLI token requests")
+
 	pflag.CommandLine.AddGoFlag(flag.CommandLine.Lookup("v"))
 	pflag.CommandLine.AddGoFlag(flag.CommandLine.Lookup("logtostderr"))
+	pflag.CommandLine.AddGoFlag(flag.CommandLine.Lookup("timeout"))
 	_ = pflag.CommandLine.Set("logtostderr", "true")
-	root := cmd.NewRootCmd(v.String())
+	pflag.Parse()
+
+	root := cmd.NewRootCmd(v.String(), timeoutDuration)
 	if err := root.Execute(); err != nil {
 		os.Exit(1)
 	}

pkg/cmd/convert.go

@@ -1,14 +1,16 @@
 package cmd
 
 import (
+	"time"
+
 	"github.com/Azure/kubelogin/pkg/converter"
 	"github.com/spf13/cobra"
 	"k8s.io/client-go/tools/clientcmd"
 )
 
 // NewConvertCmd provides a cobra command for convert sub command
-func NewConvertCmd() *cobra.Command {
-	o := converter.New()
+func NewConvertCmd(timeout time.Duration) *cobra.Command {
+	o := converter.New(timeout)
 
 	cmd := &cobra.Command{
 		Use:          "convert-kubeconfig",

pkg/cmd/root.go

@@ -1,11 +1,13 @@
 package cmd
 
 import (
+	"time"
+
 	"github.com/spf13/cobra"
 )
 
 // NewRootCmd provides a cobra root command
-func NewRootCmd(version string) *cobra.Command {
+func NewRootCmd(version string, timeout time.Duration) *cobra.Command {
 
 	cmd := &cobra.Command{
 		Use:          "kubelogin",
@@ -17,8 +19,8 @@ func NewRootCmd(version string) *cobra.Command {
 		},
 	}
 
-	cmd.AddCommand(NewConvertCmd())
-	cmd.AddCommand(NewTokenCmd())
+	cmd.AddCommand(NewConvertCmd(timeout))
+	cmd.AddCommand(NewTokenCmd(timeout))
 	cmd.AddCommand(NewRemoveTokenCacheCmd())
 
 	return cmd

pkg/cmd/token.go

@@ -1,13 +1,16 @@
 package cmd
 
 import (
+	"time"
+
 	"github.com/Azure/kubelogin/pkg/token"
 	"github.com/spf13/cobra"
 )
 
-// NewTokenCmd provides a cobra command for convert sub command
-func NewTokenCmd() *cobra.Command {
-	o := token.NewOptions()
+// NewTokenCmd provides a cobra command for the get-token sub command
+func NewTokenCmd(timeout time.Duration) *cobra.Command {
+	// Updated to pass the timeout directly to NewOptions
+	o := token.NewOptions(timeout)
 
 	cmd := &cobra.Command{
 		Use:          "get-token",

pkg/converter/options.go

@@ -2,6 +2,7 @@ package converter
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/Azure/kubelogin/pkg/token"
 	"github.com/spf13/pflag"
@@ -19,15 +20,17 @@ type Options struct {
 
 func stringptr(str string) *string { return &str }
 
-func New() Options {
+func New(timeout time.Duration) Options {
 	configFlags := &genericclioptions.ConfigFlags{
 		KubeConfig: stringptr(""),
 	}
-	return Options{configFlags: configFlags}
+	return Options{
+		configFlags:  configFlags,
+		TokenOptions: token.NewOptions(timeout),
+	}
 }
 
 func (o *Options) AddFlags(fs *pflag.FlagSet) {
-	o.TokenOptions = token.NewOptions()
 	if cf, ok := o.configFlags.(*genericclioptions.ConfigFlags); ok {
 		cf.AddFlags(fs)
 	}

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

@@ -2,12 +2,14 @@ package token
 
 import (
 	"testing"
+	"time"
 
 	"github.com/Azure/kubelogin/pkg/testutils"
 )
 
 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)

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
@@ -85,12 +87,13 @@ func GetSupportedLogins() string {
 	return strings.Join(supportedLogin, ", ")
 }
 
-func NewOptions() Options {
-	return Options{
-		LoginMethod:   DeviceCodeLogin,
-		Environment:   defaultEnvironmentName,
-		TokenCacheDir: DefaultTokenCacheDir,
-	}
+func NewOptions(timeout time.Duration) Options {
+    return Options{
+        LoginMethod:   DeviceCodeLogin,
+        Environment:   defaultEnvironmentName,
+        TokenCacheDir: DefaultTokenCacheDir,
+        Timeout:       timeout,
+    }
 }
 
 func (o *Options) AddFlags(fs *pflag.FlagSet) {

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, defaultTimeout)
 	case WorkloadIdentityLogin:
 		return newWorkloadIdentityToken(o.ClientID, o.FederatedTokenFile, o.AuthorityHost, o.ServerID, o.TenantID)
 	}