agent: handle dependencies between cached leases during persistent storage restore

changelog/12765.txt

@@ -0,0 +1,3 @@
+```release-note:improvement
+agent/cache: Process persistent cache leases in dependency order during restore to ensure child leases are always correctly restored
+```

command/agent/cache/cacheboltdb/bolt.go

@@ -149,15 +149,6 @@ func (b *BoltStorage) Set(ctx context.Context, id string, plaintext []byte, inde
 	})
 }
 
-func getBucketIDs(b *bolt.Bucket) ([][]byte, error) {
-	ids := [][]byte{}
-	err := b.ForEach(func(k, v []byte) error {
-		ids = append(ids, k)
-		return nil
-	})
-	return ids, err
-}
-
 // Delete an index (token or lease) by id from bolt storage
 func (b *BoltStorage) Delete(id string) error {
 	return b.db.Update(func(tx *bolt.Tx) error {

command/agent/cache/cachememdb/index.go

@@ -60,6 +60,9 @@ type Index struct {
 	// RequestToken is the token used in the request
 	RequestToken string
 
+	// RequestTokenIndexID is the ID of the RequestToken's entry in the cache
+	RequestTokenIndexID string
+
 	// RequestHeader is the header used in the request
 	RequestHeader http.Header
 

command/agent/cache/lease_cache.go

@@ -356,6 +356,7 @@ func (c *LeaseCache) Send(ctx context.Context, req *SendRequest) (*SendResponse,
 
 		index.Lease = secret.LeaseID
 		index.LeaseToken = req.Token
+		index.RequestTokenIndexID = entry.ID
 
 		index.Type = cacheboltdb.SecretLeaseType
 
@@ -381,6 +382,7 @@ func (c *LeaseCache) Send(ctx context.Context, req *SendRequest) (*SendResponse,
 			parentCtx = entry.RenewCtxInfo.Ctx
 
 			index.TokenParent = req.Token
+			index.RequestTokenIndexID = entry.ID
 		}
 
 		renewCtxInfo = c.createCtxInfo(parentCtx)
@@ -556,7 +558,9 @@ func computeIndexID(req *SendRequest) (string, error) {
 
 	// Append req.Token into the byte slice. This is needed since auto-auth'ed
 	// requests sets the token directly into SendRequest.Token
-	b.Write([]byte(req.Token))
+	if _, err := b.Write([]byte(req.Token)); err != nil {
+		return "", fmt.Errorf("error writing token to hash input: %w", err)
+	}
 
 	return hex.EncodeToString(cryptoutil.Blake2b256Hash(string(b.Bytes()))), nil
 }
@@ -966,101 +970,145 @@ func (c *LeaseCache) Flush() error {
 	return nil
 }
 
+type indexAndChannel struct {
+	index *cachememdb.Index
+	ch    chan struct{}
+}
+
+type leaseMaps struct {
+	// Maps of index ID -> index and channel
+	authLeases   map[string]indexAndChannel
+	secretLeases map[string]indexAndChannel
+}
+
 // Restore loads the cachememdb from the persistent storage passed in. Loads
 // tokens first, since restoring a lease's renewal context and watcher requires
 // looking up the token in the cachememdb.
 func (c *LeaseCache) Restore(ctx context.Context, storage *cacheboltdb.BoltStorage) error {
-	var errors *multierror.Error
+	return c.restoreWithChannel(ctx, storage, nil)
+}
+
+func (c *LeaseCache) restoreWithChannel(ctx context.Context, storage *cacheboltdb.BoltStorage, ch chan *cachememdb.Index) error {
+	var errs *multierror.Error
 
 	// Process tokens first
 	tokens, err := storage.GetByType(ctx, cacheboltdb.TokenType)
 	if err != nil {
-		errors = multierror.Append(errors, err)
+		errs = multierror.Append(errs, err)
 	} else {
 		if err := c.restoreTokens(tokens); err != nil {
-			errors = multierror.Append(errors, err)
+			errs = multierror.Append(errs, err)
 		}
 	}
 
-	// Then process auth leases
-	authLeases, err := storage.GetByType(ctx, cacheboltdb.AuthLeaseType)
-	if err != nil {
-		errors = multierror.Append(errors, err)
-	} else {
-		if err := c.restoreLeases(authLeases); err != nil {
-			errors = multierror.Append(errors, err)
-		}
+	leaseMaps := &leaseMaps{
+		authLeases:   make(map[string]indexAndChannel),
+		secretLeases: make(map[string]indexAndChannel),
 	}
 
-	// Then process secret leases
-	secretLeases, err := storage.GetByType(ctx, cacheboltdb.SecretLeaseType)
-	if err != nil {
-		errors = multierror.Append(errors, err)
-	} else {
-		if err := c.restoreLeases(secretLeases); err != nil {
-			errors = multierror.Append(errors, err)
+	// Fetch all the auth and secret leases we'll process upfront.
+	// Doing so allows us to process dependencies between the leases if there are any.
+	// Any lease that depends on another can wait for its parent to be restored before
+	// being processed itself.
+	// This algorithm requires that the graph of dependencies is acyclic.
+	for _, leaseType := range []struct {
+		name string
+		m    map[string]indexAndChannel
+	}{
+		{cacheboltdb.AuthLeaseType, leaseMaps.authLeases},
+		{cacheboltdb.SecretLeaseType, leaseMaps.secretLeases},
+	} {
+		leases, err := storage.GetByType(ctx, leaseType.name)
+		if err != nil {
+			errs = multierror.Append(errs, err)
+			continue
+		}
+
+		for _, lease := range leases {
+			index, err := cachememdb.Deserialize(lease)
+			if err != nil {
+				errs = multierror.Append(errs, err)
+				continue
+			}
+			leaseType.m[index.ID] = indexAndChannel{
+				index: index,
+				ch:    make(chan struct{}),
+			}
 		}
 	}
 
-	return errors.ErrorOrNil()
-}
+	// Now restore the auth and secret leases.
+	wg := sync.WaitGroup{}
+	mtx := sync.Mutex{}
+	for _, leases := range []map[string]indexAndChannel{leaseMaps.authLeases, leaseMaps.secretLeases} {
+		for id, lease := range leases {
+			wg.Add(1)
+			go func(id string, lease indexAndChannel) {
+				defer wg.Done()
+				defer close(lease.ch)
 
-func (c *LeaseCache) restoreTokens(tokens [][]byte) error {
-	var errors *multierror.Error
+				c.logger.Trace("processing lease", "id", id)
+				// Check if this lease has already expired
+				expired, err := c.hasExpired(time.Now().UTC(), lease.index)
+				if err != nil {
+					c.logger.Warn("failed to check if lease is expired", "id", id, "error", err)
+				}
+				if expired {
+					return
+				}
 
-	for _, token := range tokens {
-		newIndex, err := cachememdb.Deserialize(token)
-		if err != nil {
-			errors = multierror.Append(errors, err)
-			continue
-		}
-		newIndex.RenewCtxInfo = c.createCtxInfo(nil)
-		if err := c.db.Set(newIndex); err != nil {
-			errors = multierror.Append(errors, err)
-			continue
+				if err := c.restoreLeaseRenewCtx(lease.index, leaseMaps.authLeases); err != nil {
+					mtx.Lock()
+					errs = multierror.Append(errs, err)
+					mtx.Unlock()
+					return
+				}
+				if err := c.db.Set(lease.index); err != nil {
+					mtx.Lock()
+					errs = multierror.Append(errs, err)
+					mtx.Unlock()
+					return
+				}
+				if ch != nil {
+					ch <- lease.index
+				}
+				c.logger.Trace("restored lease", "id", id, "path", lease.index.RequestPath)
+			}(id, lease)
 		}
-		c.logger.Trace("restored token", "id", newIndex.ID)
 	}
 
-	return errors.ErrorOrNil()
+	wg.Wait()
+	if ch != nil {
+		close(ch)
+	}
+
+	return errs.ErrorOrNil()
 }
 
-func (c *LeaseCache) restoreLeases(leases [][]byte) error {
+func (c *LeaseCache) restoreTokens(tokens [][]byte) error {
 	var errors *multierror.Error
 
-	for _, lease := range leases {
-		newIndex, err := cachememdb.Deserialize(lease)
+	for _, token := range tokens {
+		newIndex, err := cachememdb.Deserialize(token)
 		if err != nil {
 			errors = multierror.Append(errors, err)
 			continue
 		}
 
-		// Check if this lease has already expired
-		expired, err := c.hasExpired(time.Now().UTC(), newIndex)
-		if err != nil {
-			c.logger.Warn("failed to check if lease is expired", "id", newIndex.ID, "error", err)
-		}
-		if expired {
-			continue
-		}
-
-		if err := c.restoreLeaseRenewCtx(newIndex); err != nil {
-			errors = multierror.Append(errors, err)
-			continue
-		}
+		newIndex.RenewCtxInfo = c.createCtxInfo(nil)
 		if err := c.db.Set(newIndex); err != nil {
 			errors = multierror.Append(errors, err)
 			continue
 		}
-		c.logger.Trace("restored lease", "id", newIndex.ID, "path", newIndex.RequestPath)
+		c.logger.Trace("restored token", "id", newIndex.ID)
 	}
 
 	return errors.ErrorOrNil()
 }
 
 // restoreLeaseRenewCtx re-creates a RenewCtx for an index object and starts
 // the watcher go routine
-func (c *LeaseCache) restoreLeaseRenewCtx(index *cachememdb.Index) error {
+func (c *LeaseCache) restoreLeaseRenewCtx(index *cachememdb.Index, authLeases map[string]indexAndChannel) error {
 	if index.Response == nil {
 		return fmt.Errorf("cached response was nil for %s", index.ID)
 	}
@@ -1081,6 +1129,13 @@ func (c *LeaseCache) restoreLeaseRenewCtx(index *cachememdb.Index) error {
 	var renewCtxInfo *cachememdb.ContextInfo
 	switch {
 	case secret.LeaseID != "":
+		if parent, ok := authLeases[index.RequestTokenIndexID]; ok {
+			c.logger.Trace("waiting for parent token to restore", "id", index.RequestTokenIndexID)
+			select {
+			case <-parent.ch:
+			}
+			c.logger.Trace("parent token restored", "id", index.RequestTokenIndexID)
+		}
 		entry, err := c.db.Get(cachememdb.IndexNameToken, index.RequestToken)
 		if err != nil {
 			return err
@@ -1096,6 +1151,13 @@ func (c *LeaseCache) restoreLeaseRenewCtx(index *cachememdb.Index) error {
 	case secret.Auth != nil:
 		var parentCtx context.Context
 		if !secret.Auth.Orphan {
+			if parent, ok := authLeases[index.RequestTokenIndexID]; ok {
+				c.logger.Trace("waiting for parent token to restore", "id", index.RequestTokenIndexID)
+				select {
+				case <-parent.ch:
+				}
+				c.logger.Trace("parent token restored", "id", index.RequestTokenIndexID)
+			}
 			entry, err := c.db.Get(cachememdb.IndexNameToken, index.RequestToken)
 			if err != nil {
 				return err