Add optional caching for Vault clients, including token re-use. (#1537)

The new functionality is controlled using the newly-introduced
--experimental-enable-vault-token-cache and
--experimental-vault-token-cache-size command-line flags.

Signed-off-by: NicEggert <nicholas.eggert@target.com>

cmd/root.go

@@ -38,6 +38,7 @@ import (
 	"github.com/external-secrets/external-secrets/pkg/controllers/externalsecret"
 	"github.com/external-secrets/external-secrets/pkg/controllers/secretstore"
 	awsauth "github.com/external-secrets/external-secrets/pkg/provider/aws/auth"
+	"github.com/external-secrets/external-secrets/pkg/provider/vault"
 )
 
 var (
@@ -67,6 +68,8 @@ var (
 	certCheckInterval                     time.Duration
 	certLookaheadInterval                 time.Duration
 	enableAWSSession                      bool
+	enableVaultTokenCache                 bool
+	vaultTokenCacheSize                   int
 	tlsCiphers                            string
 	tlsMinVersion                         string
 )
@@ -176,6 +179,10 @@ var rootCmd = &cobra.Command{
 		if enableAWSSession {
 			awsauth.EnableCache = true
 		}
+		if enableVaultTokenCache {
+			vault.EnableCache = true
+			vault.VaultClientCache.Size = vaultTokenCacheSize
+		}
 		setupLog.Info("starting manager")
 		if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
 			setupLog.Error(err, "problem running manager")
@@ -207,4 +214,6 @@ func init() {
 	rootCmd.Flags().DurationVar(&storeRequeueInterval, "store-requeue-interval", time.Minute*5, "Default Time duration between reconciling (Cluster)SecretStores")
 	rootCmd.Flags().BoolVar(&enableFloodGate, "enable-flood-gate", true, "Enable flood gate. External secret will be reconciled only if the ClusterStore or Store have an healthy or unknown state.")
 	rootCmd.Flags().BoolVar(&enableAWSSession, "experimental-enable-aws-session-cache", false, "Enable experimental AWS session cache. External secret will reuse the AWS session without creating a new one on each request.")
+	rootCmd.Flags().BoolVar(&enableVaultTokenCache, "experimental-enable-vault-token-cache", false, "Enable experimental Vault token cache. External secrets will reuse the Vault token without creating a new one on each request.")
+	rootCmd.Flags().IntVar(&vaultTokenCacheSize, "experimental-vault-token-cache-size", 100, "Maximum size of Vault token cache. Only used if --experimental-enable-vault-token-cache is set.")
 }

go.mod

@@ -94,7 +94,10 @@ require (
 
 require github.com/1Password/connect-sdk-go v1.5.0
 
-require sigs.k8s.io/yaml v1.3.0
+require (
+	github.com/hashicorp/golang-lru v0.5.4
+	sigs.k8s.io/yaml v1.3.0
+)
 
 require (
 	cloud.google.com/go/compute v1.9.0 // indirect
@@ -156,7 +159,6 @@ require (
 	github.com/hashicorp/go-sockaddr v1.0.2 // indirect
 	github.com/hashicorp/go-uuid v1.0.3 // indirect
 	github.com/hashicorp/go-version v1.6.0 // indirect
-	github.com/hashicorp/golang-lru v0.5.4 // indirect
 	github.com/hashicorp/hcl v1.0.1-vault-3 // indirect
 	github.com/hashicorp/vault/sdk v0.6.0 // indirect
 	github.com/hashicorp/yamux v0.1.1 // indirect

pkg/provider/vault/cache.go

@@ -0,0 +1,106 @@
+/*
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package vault
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+
+	lru "github.com/hashicorp/golang-lru"
+
+	esv1beta1 "github.com/external-secrets/external-secrets/apis/externalsecrets/v1beta1"
+)
+
+type clientCache struct {
+	cache       *lru.Cache
+	Size        int
+	initialized bool
+	mu          sync.Mutex
+}
+
+type clientCacheKey struct {
+	Name      string
+	Namespace string
+	Kind      string
+}
+
+type clientCacheValue struct {
+	ResourceVersion string
+	Client          Client
+}
+
+func (c *clientCache) initialize() error {
+	if !c.initialized {
+		var err error
+		c.cache, err = lru.New(c.Size)
+		if err != nil {
+			return fmt.Errorf(errVaultCacheCreate, err)
+		}
+		c.initialized = true
+	}
+	return nil
+}
+
+func (c *clientCache) get(ctx context.Context, store esv1beta1.GenericStore, key clientCacheKey) (Client, bool, error) {
+	value, ok := c.cache.Get(key)
+	if ok {
+		cachedClient := value.(clientCacheValue)
+		if cachedClient.ResourceVersion == store.GetObjectMeta().ResourceVersion {
+			return cachedClient.Client, true, nil
+		}
+		// revoke token and clear old item from cache if resource has been updated
+		err := revokeTokenIfValid(ctx, cachedClient.Client)
+		if err != nil {
+			return nil, false, err
+		}
+		c.cache.Remove(key)
+	}
+	return nil, false, nil
+}
+
+func (c *clientCache) add(ctx context.Context, store esv1beta1.GenericStore, key clientCacheKey, client Client) error {
+	// don't let the LRU cache evict items
+	// remove the oldest item manually when needed so we can do some cleanup
+	for c.cache.Len() >= c.Size {
+		_, value, ok := c.cache.RemoveOldest()
+		if !ok {
+			return errors.New(errVaultCacheRemove)
+		}
+		cachedClient := value.(clientCacheValue)
+		err := revokeTokenIfValid(ctx, cachedClient.Client)
+		if err != nil {
+			return fmt.Errorf(errVaultRevokeToken, err)
+		}
+	}
+	evicted := c.cache.Add(key, clientCacheValue{ResourceVersion: store.GetObjectMeta().ResourceVersion, Client: client})
+	if evicted {
+		return errors.New(errVaultCacheEviction)
+	}
+	return nil
+}
+
+func (c *clientCache) contains(key clientCacheKey) bool {
+	return c.cache.Contains(key)
+}
+
+func (c *clientCache) lock() {
+	c.mu.Lock()
+}
+
+func (c *clientCache) unlock() {
+	c.mu.Unlock()
+}

pkg/provider/vault/fake/vault.go

@@ -88,6 +88,12 @@ type VaultListResponse struct {
 	Data     *vault.Response
 }
 
+func NewAuthTokenFn() Token {
+	return Token{nil, func(ctx context.Context) (*vault.Secret, error) {
+		return &(vault.Secret{}), nil
+	}}
+}
+
 func NewSetTokenFn(ofn ...func(v string)) MockSetTokenFn {
 	return func(v string) {
 		for _, fn := range ofn {

pkg/provider/vault/vault.go

@@ -49,14 +49,19 @@ import (
 )
 
 var (
-	_ esv1beta1.Provider      = &connector{}
-	_ esv1beta1.SecretsClient = &client{}
+	_                esv1beta1.Provider      = &connector{}
+	_                esv1beta1.SecretsClient = &client{}
+	EnableCache      bool
+	VaultClientCache clientCache
 )
 
 const (
 	serviceAccTokenPath = "/var/run/secrets/kubernetes.io/serviceaccount/token"
 
 	errVaultStore           = "received invalid Vault SecretStore resource: %w"
+	errVaultCacheCreate     = "cannot create Vault client cache: %s"
+	errVaultCacheRemove     = "error removing item from Vault client cache: %w"
+	errVaultCacheEviction   = "unexpected eviction from Vault client cache"
 	errVaultClient          = "cannot setup new vault client: %w"
 	errVaultCert            = "cannot set Vault CA certificate: %w"
 	errReadSecret           = "cannot read secret data from Vault: %w"
@@ -220,6 +225,49 @@ func newVaultClient(c *vault.Config) (Client, error) {
 	return out, nil
 }
 
+func getVaultClient(ctx context.Context, c *connector, store esv1beta1.GenericStore, cfg *vault.Config) (Client, error) {
+	isStaticToken := store.GetSpec().Provider.Vault.Auth.TokenSecretRef != nil
+	useCache := EnableCache && !isStaticToken
+
+	if useCache {
+		VaultClientCache.lock()
+		defer VaultClientCache.unlock()
+
+		err := VaultClientCache.initialize()
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	key := clientCacheKey{
+		Name:      store.GetObjectMeta().Name,
+		Namespace: store.GetObjectMeta().Namespace,
+		Kind:      store.GetTypeMeta().Kind,
+	}
+	if useCache {
+		client, ok, err := VaultClientCache.get(ctx, store, key)
+		if err != nil {
+			return nil, err
+		}
+		if ok {
+			return client, nil
+		}
+	}
+
+	client, err := c.newVaultClient(cfg)
+	if err != nil {
+		return nil, fmt.Errorf(errVaultClient, err)
+	}
+
+	if useCache && !VaultClientCache.contains(key) {
+		err = VaultClientCache.add(ctx, store, key, client)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return client, nil
+}
+
 type connector struct {
 	newVaultClient func(c *vault.Config) (Client, error)
 }
@@ -260,7 +308,7 @@ func (c *connector) newClient(ctx context.Context, store esv1beta1.GenericStore,
 		return nil, err
 	}
 
-	client, err := c.newVaultClient(cfg)
+	client, err := getVaultClient(ctx, c, store, cfg)
 	if err != nil {
 		return nil, fmt.Errorf(errVaultClient, err)
 	}
@@ -579,18 +627,12 @@ func getTypedKey(data map[string]interface{}, key string) ([]byte, error) {
 }
 
 func (v *client) Close(ctx context.Context) error {
-	// Revoke the token if we have one set and it wasn't sourced from a TokenSecretRef
-	if v.client.Token() != "" && v.store.Auth.TokenSecretRef == nil {
-		revoke, err := checkToken(ctx, v)
+	// Revoke the token if we have one set, it wasn't sourced from a TokenSecretRef,
+	// and token caching isn't enabled
+	if !EnableCache && v.client.Token() != "" && v.store.Auth.TokenSecretRef == nil {
+		err := revokeTokenIfValid(ctx, v.client)
 		if err != nil {
-			return fmt.Errorf(errVaultRevokeToken, err)
-		}
-		if revoke {
-			err = v.token.RevokeSelfWithContext(ctx, v.client.Token())
-			if err != nil {
-				return fmt.Errorf(errVaultRevokeToken, err)
-			}
-			v.client.ClearToken()
+			return err
 		}
 	}
 	return nil
@@ -631,7 +673,7 @@ func (v *client) Validate() (esv1beta1.ValidationResult, error) {
 	if v.storeKind == esv1beta1.ClusterSecretStoreKind && isReferentSpec(v.store) {
 		return esv1beta1.ValidationResultUnknown, nil
 	}
-	_, err := checkToken(context.Background(), v)
+	_, err := checkToken(context.Background(), v.token)
 	if err != nil {
 		return esv1beta1.ValidationResultError, fmt.Errorf(errInvalidCredentials, err)
 	}
@@ -816,41 +858,61 @@ func getCertFromConfigMap(v *client) ([]byte, error) {
 	return []byte(val), nil
 }
 
+/*
+setAuth gets a new token using the configured mechanism.
+If there's already a valid token, does nothing.
+*/
 func (v *client) setAuth(ctx context.Context, cfg *vault.Config) error {
-	tokenExists, err := setSecretKeyToken(ctx, v)
+	tokenExists := false
+	var err error
+	if v.client.Token() != "" {
+		tokenExists, err = checkToken(ctx, v.token)
+	}
+	if tokenExists {
+		v.log.V(1).Info("Re-using existing token")
+		return err
+	}
+
+	tokenExists, err = setSecretKeyToken(ctx, v)
 	if tokenExists {
+		v.log.V(1).Info("Set token from secret")
 		return err
 	}
 
 	tokenExists, err = setAppRoleToken(ctx, v)
 	if tokenExists {
+		v.log.V(1).Info("Retrieved new token using AppRole auth")
 		return err
 	}
 
 	tokenExists, err = setKubernetesAuthToken(ctx, v)
 	if tokenExists {
+		v.log.V(1).Info("Retrieved new token using Kubernetes auth")
 		return err
 	}
 
 	tokenExists, err = setLdapAuthToken(ctx, v)
 	if tokenExists {
+		v.log.V(1).Info("Retrieved new token using LDAP auth")
 		return err
 	}
 
 	tokenExists, err = setJwtAuthToken(ctx, v)
 	if tokenExists {
+		v.log.V(1).Info("Retrieved new token using JWT auth")
 		return err
 	}
 
 	tokenExists, err = setCertAuthToken(ctx, v, cfg)
 	if tokenExists {
+		v.log.V(1).Info("Retrieved new token using certificate auth")
 		return err
 	}
 
 	return errors.New(errAuthFormat)
 }
 
-func setAppRoleToken(ctx context.Context, v *client) (bool, error) {
+func setSecretKeyToken(ctx context.Context, v *client) (bool, error) {
 	tokenRef := v.store.Auth.TokenSecretRef
 	if tokenRef != nil {
 		token, err := v.secretKeyRef(ctx, tokenRef)
@@ -863,7 +925,7 @@ func setAppRoleToken(ctx context.Context, v *client) (bool, error) {
 	return false, nil
 }
 
-func setSecretKeyToken(ctx context.Context, v *client) (bool, error) {
+func setAppRoleToken(ctx context.Context, v *client) (bool, error) {
 	appRole := v.store.Auth.AppRole
 	if appRole != nil {
 		err := v.requestTokenWithAppRoleRef(ctx, appRole)
@@ -1007,9 +1069,9 @@ func (v *client) serviceAccountToken(ctx context.Context, serviceAccountRef esme
 }
 
 // checkToken does a lookup and checks if the provided token exists.
-func checkToken(ctx context.Context, vStore *client) (bool, error) {
+func checkToken(ctx context.Context, token Token) (bool, error) {
 	// https://www.vaultproject.io/api-docs/auth/token#lookup-a-token-self
-	resp, err := vStore.token.LookupSelfWithContext(ctx)
+	resp, err := token.LookupSelfWithContext(ctx)
 	if err != nil {
 		return false, err
 	}
@@ -1024,6 +1086,21 @@ func checkToken(ctx context.Context, vStore *client) (bool, error) {
 	return true, nil
 }
 
+func revokeTokenIfValid(ctx context.Context, client Client) error {
+	valid, err := checkToken(ctx, client.AuthToken())
+	if err != nil {
+		return fmt.Errorf(errVaultRevokeToken, err)
+	}
+	if valid {
+		err = client.AuthToken().RevokeSelfWithContext(ctx, client.Token())
+		if err != nil {
+			return fmt.Errorf(errVaultRevokeToken, err)
+		}
+		client.ClearToken()
+	}
+	return nil
+}
+
 func (v *client) requestTokenWithAppRoleRef(ctx context.Context, appRole *esv1beta1.VaultAppRole) error {
 	roleID := strings.TrimSpace(appRole.RoleID)
 

pkg/provider/vault/vault_test.go

@@ -190,9 +190,11 @@ type testCase struct {
 
 func clientWithLoginMock(c *vault.Config) (Client, error) {
 	cl := fake.VaultClient{
-		MockSetToken: fake.NewSetTokenFn(),
-		MockAuth:     fake.NewVaultAuth(),
-		MockLogical:  fake.NewVaultLogical(),
+		MockAuthToken: fake.NewAuthTokenFn(),
+		MockSetToken:  fake.NewSetTokenFn(),
+		MockToken:     fake.NewTokenFn(""),
+		MockAuth:      fake.NewVaultAuth(),
+		MockLogical:   fake.NewVaultLogical(),
 	}
 	auth := cl.Auth()
 	token := cl.AuthToken()