Add flags to backfill script (#2146)

* Fix godoc comments and variable names for options

Signed-off-by: Colleen Murphy <colleenmurphy@google.com>

* Add retry options and make all configurable

Add options to the rekor client to set the minimum and maximum retry
time, and make all rekor client options configurable for the backfill
script.

Signed-off-by: Colleen Murphy <colleenmurphy@google.com>

* Add ability to set custom headers in backfill

Add flags to the backfill script to allow passing custom headers to the
Rekor service. The Rekor infrastructure may be configured to respond
in certain ways to certain headers.

Signed-off-by: Colleen Murphy <colleenmurphy@google.com>

---------

Signed-off-by: Colleen Murphy <colleenmurphy@google.com>

cmd/backfill-index/main.go

@@ -38,7 +38,9 @@ import (
 	"log"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
+	"time"
 
 	"github.com/go-openapi/runtime"
 	_ "github.com/go-sql-driver/mysql"
@@ -99,6 +101,25 @@ type mysqlClient struct {
 	client *sqlx.DB
 }
 
+type headers map[string][]string
+
+func (h *headers) String() string {
+	return fmt.Sprintf("%#v", h)
+}
+
+func (h *headers) Set(value string) error {
+	parts := strings.Split(value, "=")
+	k, v := parts[0], parts[1]
+	if *h == nil {
+		*h = make(map[string][]string)
+	}
+	if _, ok := (*h)[k]; !ok {
+		(*h)[k] = make([]string, 0)
+	}
+	(*h)[k] = append((*h)[k], v)
+	return nil
+}
+
 var (
 	redisHostname           = flag.String("redis-hostname", "", "Hostname for Redis application")
 	redisPort               = flag.String("redis-port", "", "Port to Redis application")
@@ -109,12 +130,18 @@ var (
 	startIndex              = flag.Int("start", -1, "First index to backfill")
 	endIndex                = flag.Int("end", -1, "Last index to backfill")
 	rekorAddress            = flag.String("rekor-address", "", "Address for Rekor, e.g. https://rekor.sigstore.dev")
+	rekorDisableKeepalives  = flag.Bool("rekor-disable-keepalives", true, "Disable Keep-Alive connections (defaults to true, meaning Keep-Alive is disabled)")
+	rekorRetryCount         = flag.Uint("rekor-retry-count", 3, "Maximum number of times to retry rekor requests")                          // https://github.com/sigstore/rekor/blob/5988bfa6b0761be3a810047d23b3d3191ed5af3d/pkg/client/options.go#L36
+	rekorRetryWaitMin       = flag.Duration("rekor-retry-wait-min", 1*time.Second, "Minimum time to wait between retrying rekor requests")  //nolint:revive // https://github.com/hashicorp/go-retryablehttp/blob/1542b31176d3973a6ecbc06c05a2d0df89b59afb/client.go#L49
+	rekorRetryWaitMax       = flag.Duration("rekor-retry-wait-max", 30*time.Second, "Maximum time to wait between retrying rekor requests") // https://github.com/hashicorp/go-retryablehttp/blob/1542b31176d3973a6ecbc06c05a2d0df89b59afb/client.go#L50
+	rekorHeaders            headers
 	versionFlag             = flag.Bool("version", false, "Print the current version of Backfill MySQL")
 	concurrency             = flag.Int("concurrency", 1, "Number of workers to use for backfill")
 	dryRun                  = flag.Bool("dry-run", false, "Dry run - don't actually insert into MySQL")
 )
 
 func main() {
+	flag.Var(&rekorHeaders, "rekor-header", "HTTP headers for Rekor in key=value format, repeat flag to add additional headers")
 	flag.Parse()
 
 	versionInfo := version.GetVersionInfo()
@@ -161,7 +188,12 @@ func main() {
 		log.Fatalf("creating index client: %v", err)
 	}
 
-	rekorClient, err := client.GetRekorClient(*rekorAddress, client.WithNoDisableKeepalive(true), client.WithRetryCount(10))
+	opts := []client.Option{client.WithNoDisableKeepalives(!*rekorDisableKeepalives)}
+	opts = append(opts, client.WithRetryCount(*rekorRetryCount))
+	opts = append(opts, client.WithRetryWaitMin(*rekorRetryWaitMin))
+	opts = append(opts, client.WithRetryWaitMax(*rekorRetryWaitMax))
+	opts = append(opts, client.WithHeaders(rekorHeaders))
+	rekorClient, err := client.GetRekorClient(*rekorAddress, opts...)
 	if err != nil {
 		log.Fatalf("creating rekor client: %v", err)
 	}

pkg/client/options.go

@@ -16,6 +16,7 @@ package client
 
 import (
 	"net/http"
+	"time"
 
 	"github.com/hashicorp/go-retryablehttp"
 )
@@ -26,9 +27,12 @@ type Option func(*options)
 type options struct {
 	UserAgent           string
 	RetryCount          uint
+	RetryWaitMin        time.Duration
+	RetryWaitMax        time.Duration
 	InsecureTLS         bool
 	Logger              interface{}
 	NoDisableKeepalives bool
+	Headers             map[string][]string
 }
 
 const (
@@ -63,6 +67,20 @@ func WithRetryCount(retryCount uint) Option {
 	}
 }
 
+// WithRetryWaitMin sets the minimum length of time to wait between retries.
+func WithRetryWaitMin(t time.Duration) Option {
+	return func(o *options) {
+		o.RetryWaitMin = t
+	}
+}
+
+// WithRetryWaitMax sets the minimum length of time to wait between retries.
+func WithRetryWaitMax(t time.Duration) Option {
+	return func(o *options) {
+		o.RetryWaitMax = t
+	}
+}
+
 // WithLogger sets the logger; it must implement either retryablehttp.Logger or retryablehttp.LeveledLogger; if not, this will not take effect.
 func WithLogger(logger interface{}) Option {
 	return func(o *options) {
@@ -73,39 +91,55 @@ func WithLogger(logger interface{}) Option {
 	}
 }
 
+// WithInsecureTLS disables TLS verification.
 func WithInsecureTLS(enabled bool) Option {
 	return func(o *options) {
 		o.InsecureTLS = enabled
 	}
 }
 
-func WithNoDisableKeepalive(noDisableKeepalive bool) Option {
+// WithNoDisableKeepalives unsets the default DisableKeepalives setting.
+func WithNoDisableKeepalives(noDisableKeepalives bool) Option {
 	return func(o *options) {
-		o.NoDisableKeepalives = noDisableKeepalive
+		o.NoDisableKeepalives = noDisableKeepalives
+	}
+}
+
+// WithHeaders sets default headers for every client request.
+func WithHeaders(h map[string][]string) Option {
+	return func(o *options) {
+		o.Headers = h
 	}
 }
 
 type roundTripper struct {
 	http.RoundTripper
 	UserAgent string
+	Headers   map[string][]string
 }
 
 // RoundTrip implements `http.RoundTripper`
 func (rt *roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	req.Header.Set("User-Agent", rt.UserAgent)
+	for k, v := range rt.Headers {
+		for _, h := range v {
+			req.Header.Add(k, h)
+		}
+	}
 	return rt.RoundTripper.RoundTrip(req)
 }
 
 func createRoundTripper(inner http.RoundTripper, o *options) http.RoundTripper {
 	if inner == nil {
 		inner = http.DefaultTransport
 	}
-	if o.UserAgent == "" {
+	if o.UserAgent == "" && o.Headers == nil {
 		// There's nothing to do...
 		return inner
 	}
 	return &roundTripper{
 		RoundTripper: inner,
 		UserAgent:    o.UserAgent,
+		Headers:      o.Headers,
 	}
 }

pkg/client/rekor_client.go

@@ -49,6 +49,8 @@ func GetRekorClient(rekorServerURL string, opts ...Option) (*client.Rekor, error
 		Transport: defaultTransport,
 	}
 	retryableClient.RetryMax = int(o.RetryCount)
+	retryableClient.RetryWaitMin = o.RetryWaitMin
+	retryableClient.RetryWaitMax = o.RetryWaitMax
 	retryableClient.Logger = o.Logger
 
 	httpClient := retryableClient.StandardClient()