Improve performance of NewJWTSVID

Author: sorindumitru

pkg/server/api/api.go

@@ -13,6 +13,7 @@ import (
 
 // AuthorizedEntryFetcher is the interface to fetch authorized entries
 type AuthorizedEntryFetcher interface {
+	IsAuthorizedForEntryID(ctx context.Context, id spiffeid.ID, entryID string) *types.Entry
 	// FetchAuthorizedEntries fetches the entries that the specified
 	// SPIFFE ID is authorized for
 	FetchAuthorizedEntries(ctx context.Context, id spiffeid.ID) ([]*types.Entry, error)

pkg/server/api/svid/v1/service.go

@@ -205,6 +205,18 @@ func (s *Service) BatchNewX509SVID(ctx context.Context, req *svidv1.BatchNewX509
 	return &svidv1.BatchNewX509SVIDResponse{Results: results}, nil
 }
 
+// fetchEntries fetches authorized entries using caller ID from context
+func (s *Service) isAuthorizedForSPIFFEID(ctx context.Context, entryID string, log logrus.FieldLogger) (*types.Entry, error) {
+	callerID, ok := rpccontext.CallerID(ctx)
+	if !ok {
+		return nil, api.MakeErr(log, codes.Internal, "caller ID missing from request context", nil)
+	}
+
+	entry := s.ef.IsAuthorizedForEntryID(ctx, callerID, entryID)
+	return entry, nil
+
+}
+
 // fetchEntries fetches authorized entries using caller ID from context
 func (s *Service) fetchEntries(ctx context.Context, log logrus.FieldLogger) (map[string]*types.Entry, error) {
 	callerID, ok := rpccontext.CallerID(ctx)
@@ -350,14 +362,12 @@ func (s *Service) NewJWTSVID(ctx context.Context, req *svidv1.NewJWTSVIDRequest)
 		return nil, api.MakeErr(log, status.Code(err), "rejecting request due to JWT signing request rate limiting", err)
 	}
 
-	// Fetch authorized entries
-	entriesMap, err := s.fetchEntries(ctx, log)
+	entry, err := s.isAuthorizedForSPIFFEID(ctx, req.EntryId, log)
 	if err != nil {
 		return nil, err
 	}
 
-	entry, ok := entriesMap[req.EntryId]
-	if !ok {
+	if entry == nil {
 		return nil, api.MakeErr(log, codes.NotFound, "entry not found or not authorized", nil)
 	}
 

pkg/server/cache/entrycache/fullcache.go

@@ -28,6 +28,7 @@ var _ Cache = (*FullEntryCache)(nil)
 // Cache contains a snapshot of all registration entries and Agent selectors from the data source
 // at a particular moment in time.
 type Cache interface {
+	IsAuthorizedForEntryID(agentID spiffeid.ID, entryID string) *types.Entry
 	GetAuthorizedEntries(agentID spiffeid.ID) []*types.Entry
 }
 
@@ -173,6 +174,34 @@ func Build(ctx context.Context, entryIter EntryIterator, agentIter AgentIterator
 	}, nil
 }
 
+func (c *FullEntryCache) IsAuthorizedForEntryID(agentID spiffeid.ID, entryID string) *types.Entry {
+	seen := allocSeenSet()
+	defer freeSeenSet(seen)
+
+	return c.isAuthorizedForEntryID(spiffeIDFromID(agentID), entryID, seen)
+}
+
+func (c *FullEntryCache) isAuthorizedForEntryID(id spiffeID, entryID string, seen map[spiffeID]struct{}) *types.Entry {
+	entries := c.crawl(id, seen)
+
+	for _, descendant := range entries {
+		if descendant.GetId() == entryID {
+			return descendant
+		}
+		if entry := c.isAuthorizedForEntryID(id, entryID, seen); entry != nil {
+			return entry
+		}
+	}
+
+	for _, alias := range c.aliases[id] {
+		if entry := c.isAuthorizedForEntryID(alias.id, entryID, seen); entry != nil {
+			return entry
+		}
+	}
+
+	return nil
+}
+
 // GetAuthorizedEntries gets all authorized registration entries for a given Agent SPIFFE ID.
 func (c *FullEntryCache) GetAuthorizedEntries(agentID spiffeid.ID) []*types.Entry {
 	seen := allocSeenSet()

pkg/server/endpoints/authorized_entryfetcher.go

@@ -56,6 +56,10 @@ func NewAuthorizedEntryFetcherWithEventsBasedCache(ctx context.Context, log logr
 	}, nil
 }
 
+func (a *AuthorizedEntryFetcherWithEventsBasedCache) IsAuthorizedForEntryID(_ context.Context, agentID spiffeid.ID, entryID string) *types.Entry {
+	return nil
+}
+
 func (a *AuthorizedEntryFetcherWithEventsBasedCache) FetchAuthorizedEntries(_ context.Context, agentID spiffeid.ID) ([]*types.Entry, error) {
 	return a.cache.GetAuthorizedEntries(agentID), nil
 }

pkg/server/endpoints/entryfetcher.go

@@ -49,6 +49,12 @@ func NewAuthorizedEntryFetcherWithFullCache(ctx context.Context, buildCache entr
 	}, nil
 }
 
+func (a *AuthorizedEntryFetcherWithFullCache) IsAuthorizedForEntryID(_ context.Context, agentID spiffeid.ID, entryID string) *types.Entry {
+	a.mu.RLock()
+	defer a.mu.RUnlock()
+	return a.cache.IsAuthorizedForEntryID(agentID, entryID)
+}
+
 func (a *AuthorizedEntryFetcherWithFullCache) FetchAuthorizedEntries(_ context.Context, agentID spiffeid.ID) ([]*types.Entry, error) {
 	a.mu.RLock()
 	defer a.mu.RUnlock()