spire-agent: limit JWT-SVID cache size (#5633)

* spire-agent: use a LRU cache for the JWT-SVID cache Signed-off-by: Sorin Dumitru <[email protected]>
spiffe · Nov 21, 2024 · 34c697a · 34c697a
1 parent bcf0017
commit 34c697a
Show file tree

Hide file tree

Showing 11 changed files with 181 additions and 78 deletions.
diff --git a/cmd/spire-agent/cli/run/run.go b/cmd/spire-agent/cli/run/run.go
@@ -92,6 +92,7 @@ type agentConfig struct {
 	AllowedForeignJWTClaims       []string  `hcl:"allowed_foreign_jwt_claims"`
 	AvailabilityTarget            string    `hcl:"availability_target"`
 	X509SVIDCacheMaxSize          int       `hcl:"x509_svid_cache_max_size"`
+	JWTSVIDCacheMaxSize           int       `hcl:"jwt_svid_cache_max_size"`
 
 	AuthorizedDelegates []string `hcl:"authorized_delegates"`
 
@@ -501,6 +502,11 @@ func NewAgentConfig(c *Config, logOptions []log.Option, allowUnknownConfig bool)
 	}
 	ac.X509SVIDCacheMaxSize = c.Agent.X509SVIDCacheMaxSize
 
+	if c.Agent.JWTSVIDCacheMaxSize < 0 {
+		return nil, errors.New("jwt_svid_cache_max_size should not be negative")
+	}
+	ac.JWTSVIDCacheMaxSize = c.Agent.JWTSVIDCacheMaxSize
+
 	td, err := common_cli.ParseTrustDomain(c.Agent.TrustDomain, logger)
 	if err != nil {
 		return nil, err

diff --git a/doc/spire_agent.md b/doc/spire_agent.md
@@ -70,7 +70,8 @@ This may be useful for templating configuration files, for example across differ
 | `trust_domain`                    | The trust domain that this agent belongs to (should be no more than 255 characters)                                                                                                                                                               |                                  |
 | `workload_x509_svid_key_type`     | The workload X509 SVID key type &lt;rsa-2048&vert;ec-p256&gt;                                                                                                                                                                                     | ec-p256                          |
 | `availability_target`             | The minimum amount of time desired to gracefully handle SPIRE Server or Agent downtime. This configurable influences how aggressively X509 SVIDs should be rotated. If set, must be at least 24h. See [Availability Target](#availability-target) |                                  |
-| `x509_svid_cache_max_size`        | Soft limit of max number of SVIDs that would be stored in LRU cache                                                                                                                                                                               | 1000                             |
+| `x509_svid_cache_max_size`        | Soft limit of max number of X509-SVIDs that would be stored in LRU cache                                                                                                                                                                          | 1000                             |
+| `jwt_svid_cache_max_size`         | Hard limit of max number of JWT-SVIDs that would be stored in LRU cache                                                                                                                                                                           | 1000                             |
 
 | experimental                  | Description                                                                          | Default                 |
 |:------------------------------|--------------------------------------------------------------------------------------|-------------------------|

diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go
@@ -279,7 +279,8 @@ func (a *Agent) newManager(ctx context.Context, sto storage.Storage, cat catalog
 		Storage:                  sto,
 		SyncInterval:             a.c.SyncInterval,
 		UseSyncAuthorizedEntries: a.c.UseSyncAuthorizedEntries,
-		SVIDCacheMaxSize:         a.c.X509SVIDCacheMaxSize,
+		X509SVIDCacheMaxSize:     a.c.X509SVIDCacheMaxSize,
+		JWTSVIDCacheMaxSize:      a.c.JWTSVIDCacheMaxSize,
 		SVIDStoreCache:           cache,
 		NodeAttestor:             na,
 		RotationStrategy:         rotationutil.NewRotationStrategy(a.c.AvailabilityTarget),

diff --git a/pkg/agent/api/delegatedidentity/v1/service_test.go b/pkg/agent/api/delegatedidentity/v1/service_test.go
@@ -996,7 +996,7 @@ func (m *FakeManager) SubscribeToBundleChanges() *cache.BundleStream {
 
 func newTestCache() *cache.LRUCache {
 	log, _ := test.NewNullLogger()
-	return cache.NewLRUCache(log, trustDomain1, bundle1, telemetry.Blackhole{}, cache.DefaultSVIDCacheMaxSize, clock.New())
+	return cache.NewLRUCache(log, trustDomain1, bundle1, telemetry.Blackhole{}, cache.DefaultSVIDCacheMaxSize, cache.DefaultSVIDCacheMaxSize, clock.New())
 }
 
 func generateSubscribeToX509SVIDMetrics() []fakemetrics.MetricItem {

diff --git a/pkg/agent/config.go b/pkg/agent/config.go
@@ -66,9 +66,12 @@ type Config struct {
 	// is used to sync entries from the server.
 	UseSyncAuthorizedEntries bool
 
-	// X509SVIDCacheMaxSize is a soft limit of max number of SVIDs that would be stored in cache
+	// X509SVIDCacheMaxSize is a soft limit of max number of X509-SVIDs that would be stored in cache
 	X509SVIDCacheMaxSize int
 
+	// JWTSVIDCacheMaxSize is a soft limit of max number of JWT-SVIDs that would be stored in cache
+	JWTSVIDCacheMaxSize int
+
 	// Trust domain and associated CA bundle
 	TrustDomain spiffeid.TrustDomain
 	TrustBundle []*x509.Certificate

diff --git a/pkg/agent/manager/cache/jwt_cache.go b/pkg/agent/manager/cache/jwt_cache.go
@@ -1,6 +1,7 @@
 package cache
 
 import (
+	"container/list"
 	"context"
 	"crypto/sha256"
 	"encoding/base64"
@@ -23,18 +24,36 @@ type JWTSVIDCache struct {
 	log     logrus.FieldLogger
 	metrics telemetry.Metrics
 	mu      sync.RWMutex
-	svids   map[string]*client.JWTSVID
+
+	svids   map[string]*list.Element
+	lruList *list.List
+
+	// svidCacheMaxSize is a hard limit of max number of SVIDs that would be stored in cache
+	svidCacheMaxSize int
+}
+
+type jwtSvidElement struct {
+	key  string
+	svid *client.JWTSVID
 }
 
 func (c *JWTSVIDCache) CountJWTSVIDs() int {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
 	return len(c.svids)
 }
 
-func NewJWTSVIDCache(log logrus.FieldLogger, metrics telemetry.Metrics) *JWTSVIDCache {
+func NewJWTSVIDCache(log logrus.FieldLogger, metrics telemetry.Metrics, svidCacheMaxSize int) *JWTSVIDCache {
+	if svidCacheMaxSize <= 0 {
+		svidCacheMaxSize = DefaultSVIDCacheMaxSize
+	}
 	return &JWTSVIDCache{
-		metrics: metrics,
-		log:     log,
-		svids:   make(map[string]*client.JWTSVID),
+		metrics:          metrics,
+		log:              log,
+		svids:            make(map[string]*list.Element),
+		lruList:          list.New(),
+		svidCacheMaxSize: svidCacheMaxSize,
 	}
 }
 
@@ -43,16 +62,43 @@ func (c *JWTSVIDCache) GetJWTSVID(spiffeID spiffeid.ID, audience []string) (*cli
 
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	svid, ok := c.svids[key]
-	return svid, ok
+
+	svidElement, ok := c.svids[key]
+	if !ok {
+		return nil, ok
+	}
+	c.lruList.MoveToFront(svidElement)
+
+	return svidElement.Value.(jwtSvidElement).svid, ok
 }
 
 func (c *JWTSVIDCache) SetJWTSVID(spiffeID spiffeid.ID, audience []string, svid *client.JWTSVID) {
 	key := jwtSVIDKey(spiffeID, audience)
 
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	c.svids[key] = svid
+
+	if len(c.svids) >= c.svidCacheMaxSize {
+		element := c.lruList.Back()
+		jwtSvidWithHash := element.Value.(jwtSvidElement)
+		delete(c.svids, jwtSvidWithHash.key)
+		c.lruList.Remove(element)
+	}
+
+	svidElement, ok := c.svids[key]
+	if ok {
+		svidElement.Value = jwtSvidElement{
+			key:  key,
+			svid: svid,
+		}
+		c.lruList.MoveToFront(svidElement)
+	} else {
+		svidElement = c.lruList.PushFront(jwtSvidElement{
+			key:  key,
+			svid: svid,
+		})
+		c.svids[key] = svidElement
+	}
 }
 
 func (c *JWTSVIDCache) TaintJWTSVIDs(ctx context.Context, taintedJWTAuthorities map[string]struct{}) {
@@ -64,14 +110,18 @@ func (c *JWTSVIDCache) TaintJWTSVIDs(ctx context.Context, taintedJWTAuthorities
 
 	removedKeyIDs := make(map[string]int)
 	totalCount := 0
-	for key, jwtSVID := range c.svids {
-		keyID, err := getKeyIDFromSVIDToken(jwtSVID.Token)
+	for key, element := range c.svids {
+		jwtSvidElement := element.Value.(jwtSvidElement)
+		keyID, err := getKeyIDFromSVIDToken(jwtSvidElement.svid.Token)
 		if err != nil {
 			c.log.WithError(err).Error("Could not get key ID from cached JWT-SVID")
 			continue
 		}
+
 		if _, tainted := taintedJWTAuthorities[keyID]; tainted {
 			delete(c.svids, key)
+			c.lruList.Remove(element)
+
 			removedKeyIDs[keyID]++
 			totalCount++
 		}

diff --git a/pkg/agent/manager/cache/jwt_cache_test.go b/pkg/agent/manager/cache/jwt_cache_test.go
@@ -21,13 +21,13 @@ func TestJWTSVIDCache(t *testing.T) {
 	now := time.Now()
 	tok1 := "eyJhbGciOiJFUzI1NiIsImtpZCI6ImRaRGZZaXcxdUd6TXdkTVlITDdGRVl5SzhIT0tLd0xYIiwidHlwIjoiSldUIn0.eyJhdWQiOlsidGVzdC1hdWRpZW5jZSJdLCJleHAiOjE3MjQzNjU3MzEsImlhdCI6MTcyNDI3OTQwNywic3ViIjoic3BpZmZlOi8vZXhhbXBsZS5vcmcvYWdlbnQvZGJ1c2VyIn0.dFr-oWhm5tK0bBuVXt-sGESM5l7hhoY-Gtt5DkuFoJL5Y9d4ZfmicCvUCjL4CqDB3BO_cPqmFfrO7H7pxQbGLg"
 	tok2 := "eyJhbGciOiJFUzI1NiIsImtpZCI6ImNKMXI5TVY4OTZTWXBMY0RMUjN3Q29QRHprTXpkN25tIiwidHlwIjoiSldUIn0.eyJhdWQiOlsidGVzdC1hdWRpZW5jZSJdLCJleHAiOjE3Mjg1NzEwMjUsImlhdCI6MTcyODU3MDcyNSwic3ViIjoic3BpZmZlOi8vZXhhbXBsZS5vcmcvYWdlbnQvZGJ1c2VyIn0.1YnDj7nknwIHEuNKEN0cNypXKS4SUeILXlNOsOs2XElHzfKhhDcl0sYKYtQc1Itf6cygz9C16VOQ_Yjoos2Qfg"
-	jwtSVID := &client.JWTSVID{Token: tok1, IssuedAt: now, ExpiresAt: now.Add(time.Second)}
-	jwtSVID2 := &client.JWTSVID{Token: tok2, IssuedAt: now, ExpiresAt: now.Add(time.Second)}
+	jwtSVID1 := &client.JWTSVID{Token: tok1, IssuedAt: now, ExpiresAt: now.Add(time.Minute)}
+	jwtSVID2 := &client.JWTSVID{Token: tok2, IssuedAt: now, ExpiresAt: now.Add(time.Minute)}
 
 	fakeMetrics := fakemetrics.New()
 	log, logHook := test.NewNullLogger()
 	log.Level = logrus.DebugLevel
-	cache := NewJWTSVIDCache(log, fakeMetrics)
+	cache := NewJWTSVIDCache(log, fakeMetrics, 8)
 
 	spiffeID := spiffeid.RequireFromString("spiffe://example.org/blog")
 
@@ -37,10 +37,10 @@ func TestJWTSVIDCache(t *testing.T) {
 	assert.Nil(t, actual)
 
 	// JWT is cached
-	cache.SetJWTSVID(spiffeID, []string{"bar"}, jwtSVID)
+	cache.SetJWTSVID(spiffeID, []string{"bar"}, jwtSVID1)
 	actual, ok = cache.GetJWTSVID(spiffeID, []string{"bar"})
 	assert.True(t, ok)
-	assert.Equal(t, jwtSVID, actual)
+	assert.Equal(t, jwtSVID1, actual)
 
 	// Test tainting of JWt-SVIDs
 	ctx := context.Background()
@@ -57,7 +57,7 @@ func TestJWTSVIDCache(t *testing.T) {
 			name:          "one authority tainted, one JWT-SVID",
 			taintedKeyIDs: map[string]struct{}{keyID1: {}},
 			setJWTSVIDsCached: func(cache *JWTSVIDCache) {
-				cache.SetJWTSVID(spiffeID, []string{"audience-1"}, jwtSVID)
+				cache.SetJWTSVID(spiffeID, []string{"audience-1"}, jwtSVID1)
 			},
 			expectLogs: []spiretest.LogEntry{
 				{
@@ -93,8 +93,8 @@ func TestJWTSVIDCache(t *testing.T) {
 			name:          "one authority tainted, multiple JWT-SVIDs",
 			taintedKeyIDs: map[string]struct{}{keyID1: {}},
 			setJWTSVIDsCached: func(cache *JWTSVIDCache) {
-				cache.SetJWTSVID(spiffeID, []string{"audience-1"}, jwtSVID)
-				cache.SetJWTSVID(spiffeID, []string{"audience-2"}, jwtSVID)
+				cache.SetJWTSVID(spiffeID, []string{"audience-1"}, jwtSVID1)
+				cache.SetJWTSVID(spiffeID, []string{"audience-2"}, jwtSVID1)
 			},
 			expectLogs: []spiretest.LogEntry{
 				{
@@ -130,8 +130,8 @@ func TestJWTSVIDCache(t *testing.T) {
 			name:          "multiple authorities tainted, multiple JWT-SVIDs",
 			taintedKeyIDs: map[string]struct{}{keyID1: {}, keyID2: {}},
 			setJWTSVIDsCached: func(cache *JWTSVIDCache) {
-				cache.SetJWTSVID(spiffeID, []string{"audience-1"}, jwtSVID)
-				cache.SetJWTSVID(spiffeID, []string{"audience-2"}, jwtSVID)
+				cache.SetJWTSVID(spiffeID, []string{"audience-1"}, jwtSVID1)
+				cache.SetJWTSVID(spiffeID, []string{"audience-2"}, jwtSVID1)
 				cache.SetJWTSVID(spiffeID, []string{"audience-3"}, jwtSVID2)
 			},
 			expectLogs: []spiretest.LogEntry{
@@ -176,8 +176,8 @@ func TestJWTSVIDCache(t *testing.T) {
 			name:          "none of the authorities tainted is in cache",
 			taintedKeyIDs: map[string]struct{}{"not-cached-1": {}, "not-cached-2": {}},
 			setJWTSVIDsCached: func(cache *JWTSVIDCache) {
-				cache.SetJWTSVID(spiffeID, []string{"audience-1"}, jwtSVID)
-				cache.SetJWTSVID(spiffeID, []string{"audience-2"}, jwtSVID)
+				cache.SetJWTSVID(spiffeID, []string{"audience-1"}, jwtSVID1)
+				cache.SetJWTSVID(spiffeID, []string{"audience-2"}, jwtSVID1)
 				cache.SetJWTSVID(spiffeID, []string{"audience-3"}, jwtSVID2)
 			},
 			expectMetrics: []fakemetrics.MetricItem{
@@ -203,7 +203,7 @@ func TestJWTSVIDCache(t *testing.T) {
 	} {
 		tt := tt
 		t.Run(tt.name, func(t *testing.T) {
-			cache := NewJWTSVIDCache(log, fakeMetrics)
+			cache := NewJWTSVIDCache(log, fakeMetrics, 8)
 			if tt.setJWTSVIDsCached != nil {
 				tt.setJWTSVIDsCached(cache)
 			}
@@ -221,15 +221,57 @@ func TestJWTSVIDCache(t *testing.T) {
 	}
 }
 
+func TestJWTSVIDCacheSize(t *testing.T) {
+	fakeMetrics := fakemetrics.New()
+	log, _ := test.NewNullLogger()
+	log.Level = logrus.DebugLevel
+	cache := NewJWTSVIDCache(log, fakeMetrics, 2)
+
+	now := time.Now()
+	jwtSvid1 := &client.JWTSVID{Token: "1", IssuedAt: now, ExpiresAt: now.Add(time.Minute)}
+	jwtSvid2 := &client.JWTSVID{Token: "2", IssuedAt: now, ExpiresAt: now.Add(time.Minute)}
+	jwtSvid3 := &client.JWTSVID{Token: "3", IssuedAt: now, ExpiresAt: now.Add(time.Minute)}
+
+	spiffeID := spiffeid.RequireFromString("spiffe://example.org/blog")
+	cache.SetJWTSVID(spiffeID, []string{"audience-1"}, jwtSvid1)
+	cache.SetJWTSVID(spiffeID, []string{"audience-2"}, jwtSvid2)
+	cache.SetJWTSVID(spiffeID, []string{"audience-3"}, jwtSvid3)
+
+	// The first SVID that was inserted into the cache should have been evicted.
+	_, ok := cache.GetJWTSVID(spiffeID, []string{"audience-1"})
+	assert.False(t, ok)
+
+	actual, ok := cache.GetJWTSVID(spiffeID, []string{"audience-2"})
+	assert.True(t, ok)
+	assert.Equal(t, jwtSvid2, actual)
+
+	actual, ok = cache.GetJWTSVID(spiffeID, []string{"audience-3"})
+	assert.True(t, ok)
+	assert.Equal(t, jwtSvid3, actual)
+
+	// Make the second token the most recently used token
+	_, _ = cache.GetJWTSVID(spiffeID, []string{"audience-2"})
+
+	// Insert a token
+	cache.SetJWTSVID(spiffeID, []string{"audience-1"}, jwtSvid1)
+
+	actual, ok = cache.GetJWTSVID(spiffeID, []string{"audience-2"})
+	assert.True(t, ok)
+	assert.Equal(t, jwtSvid2, actual)
+
+	_, ok = cache.GetJWTSVID(spiffeID, []string{"audience-3"})
+	assert.False(t, ok)
+}
+
 func TestJWTSVIDCacheKeyHashing(t *testing.T) {
 	spiffeID := spiffeid.RequireFromString("spiffe://example.org/blog")
 	now := time.Now()
-	expected := &client.JWTSVID{Token: "X", IssuedAt: now, ExpiresAt: now.Add(time.Second)}
+	expected := &client.JWTSVID{Token: "X", IssuedAt: now, ExpiresAt: now.Add(time.Minute)}
 
 	fakeMetrics := fakemetrics.New()
 	log, _ := test.NewNullLogger()
 	log.Level = logrus.DebugLevel
-	cache := NewJWTSVIDCache(log, fakeMetrics)
+	cache := NewJWTSVIDCache(log, fakeMetrics, 8)
 	cache.SetJWTSVID(spiffeID, []string{"ab", "cd"}, expected)
 
 	// JWT is cached