Core: call callbacks on invalidated Client

Previously, invalidated Clients were removed from the Client cache without calling any of the registered ClientCallbacks. Now any callbacks registered with ClientCallbackOnCacheRemoval be invoked when a client is removed from cache.

Author: benashz

controllers/vaultdynamicsecret_controller.go

@@ -252,7 +252,8 @@ func (r *VaultDynamicSecretReconciler) Reconcile(ctx context.Context, req ctrl.R
 			} else if !vault.IsLeaseNotFoundError(err) {
 				r.Recorder.Eventf(o, corev1.EventTypeWarning, consts.ReasonSecretLeaseRenewalError,
 					"Could not renew lease, lease_id=%s, err=%s", leaseID, err)
-			} else if !vault.IsForbiddenError(err) {
+			} else if vault.IsForbiddenError(err) {
+				logger.V(consts.LogLevelWarning).Info("Tainting client", "err", err)
 				vClient.Taint()
 			}
 			syncReason = consts.ReasonSecretLeaseRenewalError
@@ -275,6 +276,10 @@ func (r *VaultDynamicSecretReconciler) Reconcile(ctx context.Context, req ctrl.R
 	secretLease, staticCredsUpdated, err := r.syncSecret(ctx, vClient, o, transOption)
 	if err != nil {
 		r.SyncRegistry.Add(req.NamespacedName)
+		if vault.IsForbiddenError(err) {
+			logger.V(consts.LogLevelWarning).Info("Tainting client", "err", err)
+			vClient.Taint()
+		}
 		entry, _ := r.BackOffRegistry.Get(req.NamespacedName)
 		horizon := entry.NextBackOff()
 		r.Recorder.Eventf(o, corev1.EventTypeWarning, consts.ReasonSecretSyncError,
@@ -504,7 +509,7 @@ func (r *VaultDynamicSecretReconciler) SetupWithManager(mgr ctrl.Manager, opts c
 
 	r.ClientFactory.RegisterClientCallbackHandler(
 		vault.ClientCallbackHandler{
-			On:       vault.ClientCallbackOnLifetimeWatcherDone,
+			On:       vault.ClientCallbackOnLifetimeWatcherDone | vault.ClientCallbackOnCacheRemoval,
 			Callback: r.vaultClientCallback,
 		},
 	)

internal/vault/client.go

@@ -28,7 +28,7 @@ import (
 
 type ClientOptions struct {
 	SkipRenewal   bool
-	WatcherDoneCh chan<- Client
+	WatcherDoneCh chan<- *ClientCallbackHandlerRequest
 }
 
 func defaultClientOptions() *ClientOptions {
@@ -129,11 +129,10 @@ func NewClientFromStorageEntry(ctx context.Context, client ctrlclient.Client, en
 		return nil, fmt.Errorf("restored client's cacheKey %s does not match expected %s", cacheKey, entry.CacheKey)
 	}
 
-	if err := c.Validate(); err != nil {
-		return nil, err
-	}
+	c.Taint()
+	defer c.Untaint()
 
-	if _, err := c.Read(ctx, NewReadRequest("auth/token/lookup-self", nil)); err != nil {
+	if err := c.Validate(ctx); err != nil {
 		return nil, err
 	}
 
@@ -154,7 +153,7 @@ type Client interface {
 	Restore(context.Context, *api.Secret) error
 	GetTokenSecret() *api.Secret
 	CheckExpiry(int64) (bool, error)
-	Validate() error
+	Validate(ctx context.Context) error
 	GetVaultAuthObj() *secretsv1beta1.VaultAuth
 	GetVaultConnectionObj() *secretsv1beta1.VaultConnection
 	GetCredentialProvider() provider.CredentialProviderBase
@@ -184,7 +183,7 @@ type defaultClient struct {
 	inClosing          bool
 	closed             bool
 	lastWatcherErr     error
-	watcherDoneCh      chan<- Client
+	watcherDoneCh      chan<- *ClientCallbackHandlerRequest
 	tainted            bool
 	once               sync.Once
 	mu                 sync.RWMutex
@@ -220,7 +219,7 @@ func (c *defaultClient) Taint() {
 // Validate the client, returning an error for any validation failures.
 // Typically, an invalid Client would be discarded and replaced with a new
 // instance.
-func (c *defaultClient) Validate() error {
+func (c *defaultClient) Validate(ctx context.Context) error {
 	c.mu.RLock()
 	defer c.mu.RUnlock()
 
@@ -245,6 +244,16 @@ func (c *defaultClient) Validate() error {
 		return errors.New("client token expired")
 	}
 
+	if c.client == nil {
+		return errors.New("client not set")
+	}
+
+	if c.tainted {
+		if _, err := c.Read(ctx, NewReadRequest("auth/token/lookup-self", nil)); err != nil {
+			return fmt.Errorf("tainted client is invalid: %w", err)
+		}
+	}
+
 	return nil
 }
 
@@ -492,7 +501,10 @@ func (c *defaultClient) startLifetimeWatcher(ctx context.Context) error {
 				if c.watcherDoneCh != nil {
 					if !c.inClosing {
 						logger.V(consts.LogLevelTrace).Info("Writing to watcherDone channel")
-						c.watcherDoneCh <- c
+						c.watcherDoneCh <- &ClientCallbackHandlerRequest{
+							Client: c,
+							On:     ClientCallbackOnLifetimeWatcherDone,
+						}
 					} else {
 						logger.V(consts.LogLevelTrace).Info("In closing, not writing to watcherDone channel")
 					}
@@ -759,12 +771,22 @@ func (c *defaultClient) init(ctx context.Context, client ctrlclient.Client,
 }
 
 func (c *defaultClient) observeTime(ts time.Time, operation string) {
+	if c.connObj == nil {
+		// should not happen on a properly initialized Client
+		return
+	}
+
 	clientOperationTimes.WithLabelValues(operation, ctrlclient.ObjectKeyFromObject(c.connObj).String()).Observe(
 		time.Since(ts).Seconds(),
 	)
 }
 
 func (c *defaultClient) incrementOperationCounter(operation string, err error) {
+	if c.connObj == nil {
+		// should not happen on a properly initialized Client
+		return
+	}
+
 	vaultConn := ctrlclient.ObjectKeyFromObject(c.connObj).String()
 	clientOperations.WithLabelValues(operation, vaultConn).Inc()
 	if err != nil {

internal/vault/client_factory.go

@@ -28,15 +28,34 @@ import (
 )
 
 // ClientCallbackOn is an enumeration of possible client callback events.
-type ClientCallbackOn int
+type ClientCallbackOn uint32
 
 const (
+	NamePrefixVCC = "vso-cc-"
+
 	// ClientCallbackOnLifetimeWatcherDone is a ClientCallbackOn that handles client
 	// lifetime watcher done events.
-	ClientCallbackOnLifetimeWatcherDone ClientCallbackOn = iota
-	NamePrefixVCC                                        = "vso-cc-"
+	ClientCallbackOnLifetimeWatcherDone ClientCallbackOn = 1 << iota
+	// ClientCallbackOnCacheRemoval is a ClientCallbackOn that handles client cache removal events.
+	ClientCallbackOnCacheRemoval
 )
 
+func (o ClientCallbackOn) String() string {
+	switch o {
+	case ClientCallbackOnLifetimeWatcherDone:
+		return "LifetimeWatcherDone"
+	case ClientCallbackOnCacheRemoval:
+		return "CacheRemoval"
+	default:
+		return "Unknown"
+	}
+}
+
+type ClientCallbackHandlerRequest struct {
+	On     ClientCallbackOn
+	Client Client
+}
+
 // ClientCallback is a function type that takes a context, a Client, and an error as parameters.
 // It is used in the context of a ClientCallbackHandler.
 type ClientCallback func(ctx context.Context, c Client)
@@ -99,7 +118,7 @@ type cachingClientFactory struct {
 	pruneStorageOnEvict    bool
 	ctrlClient             ctrlclient.Client
 	clientCallbacks        []ClientCallbackHandler
-	callbackHandlerCh      chan Client
+	callbackHandlerCh      chan *ClientCallbackHandlerRequest
 	mu                     sync.RWMutex
 	onceDoWatcher          sync.Once
 	callbackHandlerCancel  context.CancelFunc
@@ -182,7 +201,10 @@ func (m *cachingClientFactory) prune(ctx context.Context, client ctrlclient.Clie
 	if !skipCallbacks {
 		for _, c := range pruned {
 			// the callback handler will remove the client from the storage
-			m.callbackHandlerCh <- c
+			m.callbackHandlerCh <- &ClientCallbackHandlerRequest{
+				On:     ClientCallbackOnCacheRemoval,
+				Client: c,
+			}
 		}
 	} else {
 		// for all cache entries pruned, remove the corresponding storage entries.
@@ -404,38 +426,31 @@ func (m *cachingClientFactory) Get(ctx context.Context, client ctrlclient.Client
 	if ok {
 		// return the Client from the cache if it is still Valid
 		tainted = c.Tainted()
-		logger.V(consts.LogLevelTrace).Info("Got client from cache", "clientID", c.ID(), "tainted", tainted)
-		if tainted {
-			// if the Client is tainted, we need to validate its token.
-			if _, err := c.Read(ctx, NewReadRequest("auth/token/lookup-self", nil)); err == nil {
-				defer c.Untaint()
-				tainted = false
-				if err := c.Validate(); err == nil {
-					return namespacedClient(c)
-				}
+		logger.V(consts.LogLevelTrace).Info("Got client from cache",
+			"clientID", c.ID(), "tainted", tainted)
+		if err := c.Validate(ctx); err != nil {
+			logger.V(consts.LogLevelDebug).Error(err, "Invalid client",
+				"tainted", tainted)
+			m.callbackHandlerCh <- &ClientCallbackHandlerRequest{
+				On:     ClientCallbackOnCacheRemoval,
+				Client: c,
 			}
-		} else if err := c.Validate(); err == nil {
+		} else {
+			c.Untaint()
 			return namespacedClient(c)
 		}
-
-		logger.V(consts.LogLevelDebug).Error(err, "Invalid client",
-			"tainted", tainted)
-
-		// remove the parent Client from the cache in order to prune any of its clones.
-		m.cache.Remove(cacheKey)
 	} else {
 		logger.V(consts.LogLevelTrace).Info("Client not found in cache", "cacheKey", fmt.Sprintf("%#v", cacheKey))
-	}
-
-	if !ok && m.storageEnabled() {
-		// try and restore from Client storage cache, if properly configured to do so.
-		restored, err := m.restoreClientFromCacheKey(ctx, client, cacheKey)
-		if restored != nil {
-			return namespacedClient(restored)
-		}
+		if m.storageEnabled() {
+			// try and restore from Client storage cache, if properly configured to do so.
+			restored, err := m.restoreClientFromCacheKey(ctx, client, cacheKey)
+			if restored != nil {
+				return namespacedClient(restored)
+			}
 
-		if !IsStorageEntryNotFoundErr(err) {
-			logger.Error(err, "Failed to restore client from storage")
+			if !IsStorageEntryNotFoundErr(err) {
+				logger.Error(err, "Failed to restore client from storage")
+			}
 		}
 	}
 
@@ -642,7 +657,7 @@ func (m *cachingClientFactory) storageEncryptionClient(ctx context.Context, clie
 		// ensure that the cached Vault Client is not expired, and if it is then call storageEncryptionClient() again.
 		// This operation should be safe since we are setting m.clientCacheKeyEncrypt to empty string,
 		// so there should be no risk of causing a maximum recursion error.
-		if reason := c.Validate(); reason != nil {
+		if reason := c.Validate(ctx); reason != nil {
 			m.logger.V(consts.LogLevelWarning).Info("Restored Vault client is invalid, recreating it",
 				"cacheKey", m.clientCacheKeyEncrypt, "reason", reason)
 
@@ -677,7 +692,7 @@ func (m *cachingClientFactory) startClientCallbackHandler(ctx context.Context) {
 
 	go func() {
 		if m.callbackHandlerCh == nil {
-			m.callbackHandlerCh = make(chan Client)
+			m.callbackHandlerCh = make(chan *ClientCallbackHandlerRequest)
 		}
 		defer func() {
 			close(m.callbackHandlerCh)
@@ -689,16 +704,20 @@ func (m *cachingClientFactory) startClientCallbackHandler(ctx context.Context) {
 			case <-callbackCtx.Done():
 				logger.Info("Client callback handler done")
 				return
-			case c, stillOpen := <-m.callbackHandlerCh:
+			case req, stillOpen := <-m.callbackHandlerCh:
 				if !stillOpen {
 					logger.Info("Client callback handler channel closed")
 					return
 				}
-				if c.IsClone() {
+				if req == nil {
 					continue
 				}
 
-				cacheKey, err := c.GetCacheKey()
+				if req.Client.IsClone() {
+					continue
+				}
+
+				cacheKey, err := req.Client.GetCacheKey()
 				if err != nil {
 					logger.Error(err, "Invalid client, client callbacks not executed",
 						"cacheKey", cacheKey)
@@ -708,30 +727,62 @@ func (m *cachingClientFactory) startClientCallbackHandler(ctx context.Context) {
 				// remove the client from the cache, it will be recreated when a reconciler
 				// requests it.
 				logger.V(consts.LogLevelDebug).Info("Removing client from cache", "cacheKey", cacheKey)
-				m.cache.Remove(cacheKey)
-				if m.storageEnabled() {
-					if _, err := m.pruneStorage(ctx, m.ctrlClient, cacheKey); err != nil {
-						logger.Info("Warning: failed to prune storage", "cacheKey", cacheKey)
+				if req.On&ClientCallbackOnLifetimeWatcherDone != 0 {
+					m.cache.Remove(cacheKey)
+					if m.storageEnabled() {
+						if _, err := m.pruneStorage(ctx, m.ctrlClient, cacheKey); err != nil {
+							logger.Info("Warning: failed to prune storage", "cacheKey", cacheKey)
+						}
 					}
 				}
 
-				for idx, cbReq := range m.clientCallbacks {
-					if cbReq.On != ClientCallbackOnLifetimeWatcherDone {
-						continue
-					}
-
-					logger.Info("Calling client callback on lifetime watcher done",
-						"index", idx, "cacheKey", cacheKey, "clientID", c.ID())
-					// call in a go routine to avoid blocking the channel
-					go func(cbReq ClientCallbackHandler) {
-						cbReq.Callback(ctx, c)
-					}(cbReq)
-				}
+				m.callClientCallbacks(ctx, req.Client, req.On, false)
 			}
 		}
 	}()
 }
 
+// callClientCallbacks calls all registered client callbacks for the specified
+// event. If wait is true, it will block until all callbacks have been executed.
+// Note: wait is only for testing purposes.
+func (m *cachingClientFactory) callClientCallbacks(ctx context.Context, c Client, on ClientCallbackOn, wait bool) {
+	logger := log.FromContext(ctx).WithName("callClientCallbacks")
+
+	var cbs []ClientCallbackHandler
+	for _, cbReq := range m.clientCallbacks {
+		x := on & cbReq.On
+		if x != 0 {
+			cbs = append(cbs, cbReq)
+			continue
+		}
+	}
+
+	if len(cbs) == 0 {
+		return
+	}
+
+	var wg sync.WaitGroup
+	if wait {
+		wg.Add(len(cbs))
+	}
+
+	for idx, cbReq := range cbs {
+		logger.Info("Calling client callback",
+			"index", idx, "clientID", c.ID(), "on", on)
+		// call in a go routine to avoid blocking the channel
+		go func(cbReq ClientCallbackHandler) {
+			if wait {
+				defer wg.Done()
+			}
+			cbReq.Callback(ctx, c)
+		}(cbReq)
+	}
+
+	if wait {
+		wg.Wait()
+	}
+}
+
 // NewCachingClientFactory returns a CachingClientFactory with ClientCache initialized.
 // The ClientCache's onEvictCallback is registered with the factory's onClientEvict(),
 // to ensure any evictions are handled by the factory (this is very important).
@@ -741,7 +792,7 @@ func NewCachingClientFactory(ctx context.Context, client ctrlclient.Client, cach
 		recorder:           config.Recorder,
 		persist:            config.Persist,
 		ctrlClient:         client,
-		callbackHandlerCh:  make(chan Client),
+		callbackHandlerCh:  make(chan *ClientCallbackHandlerRequest),
 		encryptionRequired: config.StorageConfig.EnforceEncryption,
 		clientLocks:        make(map[ClientCacheKey]*sync.RWMutex, config.ClientCacheSize),
 		logger: zap.New().WithName("clientCacheFactory").WithValues(