Support auth method with snapshot agent [ENT] (#15020)

Port of hashicorp/consul-enterprise#3303

command/flags/flag_map_value.go

@@ -35,3 +35,15 @@ func (h *FlagMapValue) Set(value string) error {
 
 	return nil
 }
+
+// Merge will overlay this value if it has been set.
+func (h *FlagMapValue) Merge(onto map[string]string) {
+	if h == nil || onto == nil {
+		return
+	}
+	for k, v := range *h {
+		if _, ok := onto[k]; !ok {
+			onto[k] = v
+		}
+	}
+}

command/flags/flag_map_value_test.go

@@ -3,6 +3,8 @@ package flags
 import (
 	"fmt"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestFlagMapValueSet(t *testing.T) {
@@ -78,3 +80,75 @@ func TestFlagMapValueSet(t *testing.T) {
 		}
 	})
 }
+
+func TestFlagMapValueMerge(t *testing.T) {
+	cases := map[string]struct {
+		src FlagMapValue
+		dst map[string]string
+		exp map[string]string
+	}{
+		"empty source and destination": {},
+		"empty source": {
+			dst: map[string]string{
+				"key": "val",
+			},
+			exp: map[string]string{
+				"key": "val",
+			},
+		},
+		"empty destination": {
+			src: map[string]string{
+				"key": "val",
+			},
+			dst: make(map[string]string),
+			exp: map[string]string{
+				"key": "val",
+			},
+		},
+		"non-overlapping keys": {
+			src: map[string]string{
+				"key1": "val1",
+			},
+			dst: map[string]string{
+				"key2": "val2",
+			},
+			exp: map[string]string{
+				"key1": "val1",
+				"key2": "val2",
+			},
+		},
+		"overlapping keys": {
+			src: map[string]string{
+				"key1": "val1",
+			},
+			dst: map[string]string{
+				"key1": "val2",
+			},
+			exp: map[string]string{
+				"key1": "val2",
+			},
+		},
+		"multiple keys": {
+			src: map[string]string{
+				"key1": "val1",
+				"key2": "val2",
+			},
+			dst: map[string]string{
+				"key1": "val2",
+				"key3": "val3",
+			},
+			exp: map[string]string{
+				"key1": "val2",
+				"key2": "val2",
+				"key3": "val3",
+			},
+		},
+	}
+
+	for name, c := range cases {
+		t.Run(name, func(t *testing.T) {
+			c.src.Merge(c.dst)
+			require.Equal(t, c.exp, c.dst)
+		})
+	}
+}

lib/retry/retry.go

@@ -2,6 +2,7 @@ package retry
 
 import (
 	"context"
+	"fmt"
 	"math/rand"
 	"time"
 )
@@ -30,7 +31,7 @@ func NewJitter(percent int64) Jitter {
 }
 
 // Waiter records the number of failures and performs exponential backoff when
-// when there are consecutive failures.
+// there are consecutive failures.
 type Waiter struct {
 	// MinFailures before exponential backoff starts. Any failures before
 	// MinFailures is reached will wait MinWait time.
@@ -117,3 +118,23 @@ func (w *Waiter) Wait(ctx context.Context) error {
 func (w *Waiter) NextWait() time.Duration {
 	return w.delay()
 }
+
+// RetryLoop retries an operation until either operation completes without error
+// or Waiter's context is canceled.
+func (w *Waiter) RetryLoop(ctx context.Context, operation func() error) error {
+	var lastError error
+	for {
+		if err := w.Wait(ctx); err != nil {
+			// The error will only be non-nil if the context is canceled.
+			return fmt.Errorf("could not retry operation: %w", lastError)
+		}
+
+		if err := operation(); err == nil {
+			// Reset the failure count seen by the waiter if there was no error.
+			w.Reset()
+			return nil
+		} else {
+			lastError = err
+		}
+	}
+}

lib/retry/retry_test.go

@@ -2,6 +2,7 @@ package retry
 
 import (
 	"context"
+	"fmt"
 	"math"
 	"testing"
 	"time"
@@ -157,6 +158,38 @@ func TestWaiter_Wait(t *testing.T) {
 	})
 }
 
+func TestWaiter_RetryLoop(t *testing.T) {
+	if testing.Short() {
+		t.Skip("too slow for testing.Short")
+	}
+	// Change the default factor so that we retry faster.
+	w := &Waiter{Factor: 1 * time.Millisecond}
+
+	t.Run("exits if operation is successful after a few reties", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		t.Cleanup(cancel)
+		numRetries := 0
+		err := w.RetryLoop(ctx, func() error {
+			if numRetries < 2 {
+				numRetries++
+				return fmt.Errorf("operation not successful")
+			}
+			return nil
+		})
+		require.NoError(t, err)
+	})
+
+	t.Run("errors if operation is never successful", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+		t.Cleanup(cancel)
+		err := w.RetryLoop(ctx, func() error {
+			return fmt.Errorf("operation not successful")
+		})
+		require.NotNil(t, err)
+		require.EqualError(t, err, "could not retry operation: operation not successful")
+	})
+}
+
 func runWait(ctx context.Context, w *Waiter) (time.Duration, error) {
 	before := time.Now()
 	err := w.Wait(ctx)

website/content/commands/snapshot/agent.mdx

@@ -143,6 +143,12 @@ Usage: `consul snapshot agent [options]`
     "key_file": "",
     "license_path": "",
     "tls_server_name": "",
+    "login": {
+      "auth_method": "",
+      "bearer_token": "",
+      "bearer_token_file": "",
+      "meta": {},
+    },
     "log": {
       "level": "INFO",
       "enable_syslog": false,
@@ -238,6 +244,16 @@ if desired.
 - `-syslog-facility` - Sets the facility to use for forwarding logs to syslog.
   Defaults to "LOCAL0".
 
+- `login-auth-method` - Auth method name to use to log into Consul. If provided, the token obtained with this auth method
+   will be used instead of a static token if it is provided. Currently, only `kubernetes` auth method type is supported.
+
+- `login-bearer-token` - Bearer token to use to log into Consul. Used only if `-login-auth-method` is set.
+
+- `login-bearer-token-file` - A file container bearer token to use for logging into Consul.
+  `-login-bearer-token` is ignored if this flag is provided.
+
+- `login-meta` - Metadata to set on the token, formatted as key=value. This flag may be provided multiple times.
+
 #### Local Storage Options
 
 - `-local-path` - Location to store snapshots locally. The default behavior