add retry command line flag on rekor-cli (#1097)

* enable retries with backoff in RekorClient Signed-off-by: Bob Callaway <bcallaway@google.com> * add ability to explicitly disable logger Signed-off-by: Bob Callaway <bcallaway@google.com> * pass diff logger types, ensure User-Agent is still set Signed-off-by: Bob Callaway <bcallaway@google.com> * add retry flag to rekor-cli Signed-off-by: Bob Callaway <bcallaway@google.com> Signed-off-by: Bob Callaway <bcallaway@google.com>
sigstore · Oct 6, 2022 · 45f17ce · 45f17ce
1 parent c0382c2
commit 45f17ce
Show file tree

Hide file tree

Showing 9 changed files with 72 additions and 6 deletions.
diff --git a/cmd/rekor-cli/app/get.go b/cmd/rekor-cli/app/get.go
@@ -81,7 +81,7 @@ var getCmd = &cobra.Command{
 	},
 	Run: format.WrapCmd(func(args []string) (interface{}, error) {
 		ctx := context.Background()
-		rekorClient, err := client.GetRekorClient(viper.GetString("rekor_server"), client.WithUserAgent(UserAgent()))
+		rekorClient, err := client.GetRekorClient(viper.GetString("rekor_server"), client.WithUserAgent(UserAgent()), client.WithRetryCount(viper.GetUint("retry")), client.WithLogger(log.CliLogger))
 		if err != nil {
 			return nil, err
 		}

diff --git a/cmd/rekor-cli/app/log_info.go b/cmd/rekor-cli/app/log_info.go
@@ -70,7 +70,7 @@ var logInfoCmd = &cobra.Command{
 	Run: format.WrapCmd(func(args []string) (interface{}, error) {
 		serverURL := viper.GetString("rekor_server")
 		ctx := context.Background()
-		rekorClient, err := client.GetRekorClient(serverURL, client.WithUserAgent(UserAgent()))
+		rekorClient, err := client.GetRekorClient(serverURL, client.WithUserAgent(UserAgent()), client.WithRetryCount(viper.GetUint("retry")), client.WithLogger(log.CliLogger))
 		if err != nil {
 			return nil, err
 		}

diff --git a/cmd/rekor-cli/app/log_proof.go b/cmd/rekor-cli/app/log_proof.go
@@ -26,6 +26,7 @@ import (
 	"github.com/sigstore/rekor/cmd/rekor-cli/app/format"
 	"github.com/sigstore/rekor/pkg/client"
 	"github.com/sigstore/rekor/pkg/generated/client/tlog"
+	"github.com/sigstore/rekor/pkg/log"
 )
 
 type logProofOutput struct {
@@ -72,7 +73,7 @@ var logProofCmd = &cobra.Command{
 		return nil
 	},
 	Run: format.WrapCmd(func(args []string) (interface{}, error) {
-		rekorClient, err := client.GetRekorClient(viper.GetString("rekor_server"), client.WithUserAgent(UserAgent()))
+		rekorClient, err := client.GetRekorClient(viper.GetString("rekor_server"), client.WithUserAgent(UserAgent()), client.WithRetryCount(viper.GetUint("retry")), client.WithLogger(log.CliLogger))
 		if err != nil {
 			return nil, err
 		}

diff --git a/cmd/rekor-cli/app/pflags.go b/cmd/rekor-cli/app/pflags.go
@@ -50,6 +50,7 @@ const (
 	formatFlag         FlagType = "format"
 	timeoutFlag        FlagType = "timeout"
 	base64Flag         FlagType = "base64"
+	uintFlag           FlagType = "uint"
 )
 
 type newPFlagValueFunc func() pflag.Value
@@ -121,6 +122,10 @@ func initializePFlagMap() {
 			// This validates the string is in base64 format
 			return valueFactory(base64Flag, validateBase64, "")
 		},
+		uintFlag: func() pflag.Value {
+			// This validates the string is in base64 format
+			return valueFactory(uintFlag, validateUint, "")
+		},
 	}
 }
 
@@ -320,3 +325,19 @@ func useValidator(flagType FlagType, s interface{}) error {
 
 	return nil
 }
+
+// validateUint ensures that the supplied string is a valid unsigned integer >= 0
+func validateUint(v string) error {
+	i, err := strconv.Atoi(v)
+	if err != nil {
+		return err
+	}
+	if i < 0 {
+		return fmt.Errorf("invalid unsigned int: %v", v)
+	}
+	u := struct {
+		Uint uint `validate:"gte=0"`
+	}{uint(i)}
+
+	return useValidator(uintFlag, u)
+}
diff --git a/cmd/rekor-cli/app/pflags_test.go b/cmd/rekor-cli/app/pflags_test.go
@@ -505,6 +505,48 @@ func TestValidateRekorServerURL(t *testing.T) {
 	}
 }
 
+// TestValidateRetryCount tests the validation of the retry count flag
+func TestValidateRetryCount(t *testing.T) {
+	type test struct {
+		caseDesc      string
+		retryCount    string
+		expectSuccess bool
+	}
+
+	tests := []test{
+		{
+			caseDesc:      "value not specified",
+			expectSuccess: false,
+		},
+		{
+			caseDesc:      "valid retry_count value: 0",
+			retryCount:    "0",
+			expectSuccess: true,
+		},
+		{
+			caseDesc:      "valid retry_count value: 1",
+			retryCount:    "1",
+			expectSuccess: true,
+		},
+		{
+			caseDesc:      "invalid retry_count value: asdf",
+			retryCount:    "asdf",
+			expectSuccess: false,
+		},
+		{
+			caseDesc:      "invalid retry_count value: -1",
+			retryCount:    "-1",
+			expectSuccess: false,
+		},
+	}
+
+	for _, tc := range tests {
+		if err := rootCmd.PersistentFlags().Set("retry", tc.retryCount); (err == nil) != tc.expectSuccess {
+			t.Errorf("unexpected result in '%v': %v", tc.caseDesc, err)
+		}
+	}
+}
+
 func TestSearchPFlags(t *testing.T) {
 	type test struct {
 		caseDesc              string

diff --git a/cmd/rekor-cli/app/root.go b/cmd/rekor-cli/app/root.go
@@ -24,6 +24,7 @@ import (
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 
+	"github.com/sigstore/rekor/pkg/client"
 	"github.com/sigstore/rekor/pkg/log"
 
 	// these imports are to call the packages' init methods
@@ -64,6 +65,7 @@ func init() {
 	rootCmd.PersistentFlags().Var(NewFlagValue(urlFlag, "https://rekor.sigstore.dev"), "rekor_server", "Server address:port")
 	rootCmd.PersistentFlags().Var(NewFlagValue(formatFlag, "default"), "format", "Command output format")
 	rootCmd.PersistentFlags().Var(NewFlagValue(timeoutFlag, "30s"), "timeout", "HTTP timeout")
+	rootCmd.PersistentFlags().Var(NewFlagValue(uintFlag, fmt.Sprintf("%d", client.DefaultRetryCount)), "retry", "Number of times to retry HTTP requests")
 
 	// these are bound here and not in PreRun so that all child commands can use them
 	if err := viper.BindPFlags(rootCmd.PersistentFlags()); err != nil {

diff --git a/cmd/rekor-cli/app/search.go b/cmd/rekor-cli/app/search.go
@@ -98,7 +98,7 @@ var searchCmd = &cobra.Command{
 	},
 	Run: format.WrapCmd(func(args []string) (interface{}, error) {
 		log := log.CliLogger
-		rekorClient, err := client.GetRekorClient(viper.GetString("rekor_server"), client.WithUserAgent(UserAgent()))
+		rekorClient, err := client.GetRekorClient(viper.GetString("rekor_server"), client.WithUserAgent(UserAgent()), client.WithRetryCount(viper.GetUint("retry")), client.WithLogger(log))
 		if err != nil {
 			return nil, err
 		}

diff --git a/cmd/rekor-cli/app/upload.go b/cmd/rekor-cli/app/upload.go
@@ -69,7 +69,7 @@ var uploadCmd = &cobra.Command{
 	Long: `This command takes the public key, signature and URL of the release artifact and uploads it to the rekor server.`,
 	Run: format.WrapCmd(func(args []string) (interface{}, error) {
 		ctx := context.Background()
-		rekorClient, err := client.GetRekorClient(viper.GetString("rekor_server"), client.WithUserAgent(UserAgent()))
+		rekorClient, err := client.GetRekorClient(viper.GetString("rekor_server"), client.WithUserAgent(UserAgent()), client.WithRetryCount(viper.GetUint("retry")), client.WithLogger(log.CliLogger))
 		if err != nil {
 			return nil, err
 		}

diff --git a/cmd/rekor-cli/app/verify.go b/cmd/rekor-cli/app/verify.go
@@ -92,7 +92,7 @@ var verifyCmd = &cobra.Command{
 	},
 	Run: format.WrapCmd(func(args []string) (interface{}, error) {
 		ctx := context.Background()
-		rekorClient, err := client.GetRekorClient(viper.GetString("rekor_server"), client.WithUserAgent(UserAgent()))
+		rekorClient, err := client.GetRekorClient(viper.GetString("rekor_server"), client.WithUserAgent(UserAgent()), client.WithRetryCount(viper.GetUint("retry")), client.WithLogger(log.CliLogger))
 		if err != nil {
 			return nil, err
 		}