feat: make id_token mutator cache configurable

David-Wobrock · David-Wobrock · commit 461e6a027193 · 2025-02-07T10:46:50.000+01:00
diff --git a/.schema/config.schema.json b/.schema/config.schema.json
@@ -1140,6 +1140,25 @@
           "pattern": "^[0-9]+(ns|us|ms|s|m|h)$",
           "default": "15m",
           "examples": ["1h", "1m", "30s"]
+        },
+        "cache": {
+          "additionalProperties": false,
+          "type": "object",
+          "properties": {
+            "enabled": {
+              "title": "Enabled",
+              "type": "boolean",
+              "default": true,
+              "examples": [false, true],
+              "description": "En-/disables this component."
+            },
+            "max_cost": {
+              "type": "integer",
+              "default": 33554432,
+              "title": "Maximum Cached Cost",
+              "description": "The cost of one cached JSON Web Token is the length of its string form."
+            }
+          }
         }
       },
       "additionalProperties": false
diff --git a/.schemas/mutators.id_token.schema.json b/.schemas/mutators.id_token.schema.json
@@ -32,6 +32,25 @@
       "pattern": "^[0-9]+(ns|us|ms|s|m|h)$",
       "default": "1m",
       "examples": ["1h", "1m", "30s"]
+    },
+    "cache": {
+      "additionalProperties": false,
+      "type": "object",
+      "properties": {
+        "enabled": {
+          "title": "Enabled",
+          "type": "boolean",
+          "default": true,
+          "examples": [false, true],
+          "description": "En-/disables this component."
+        },
+        "max_cost": {
+          "type": "integer",
+          "default": 33554432,
+          "title": "Maximum Cached Cost",
+          "description": "Max cost to cache."
+        }
+      }
     }
   },
   "additionalProperties": false
diff --git a/pipeline/mutate/mutator_id_token.go b/pipeline/mutate/mutator_id_token.go
@@ -38,28 +38,28 @@ type MutatorIDToken struct {
 	templates     *template.Template
 	templatesLock sync.Mutex
 
-	tokenCache        *ristretto.Cache[string, *idTokenCacheContainer]
-	tokenCacheEnabled bool
+	tokenCache *ristretto.Cache[string, *idTokenCacheContainer]
 }
 
 type CredentialsIDTokenConfig struct {
-	Claims    string `json:"claims"`
-	IssuerURL string `json:"issuer_url"`
-	JWKSURL   string `json:"jwks_url"`
-	TTL       string `json:"ttl"`
+	Claims    string             `json:"claims"`
+	IssuerURL string             `json:"issuer_url"`
+	JWKSURL   string             `json:"jwks_url"`
+	TTL       string             `json:"ttl"`
+	Cache     IdTokenCacheConfig `json:"cache"`
+}
+
+type IdTokenCacheConfig struct {
+	Enabled bool `json:"enabled"`
+	MaxCost int  `json:"max_cost"`
 }
 
 func (c *CredentialsIDTokenConfig) ClaimsTemplateID() string {
 	return fmt.Sprintf("%x", md5.Sum([]byte(c.Claims)))
 }
 
 func NewMutatorIDToken(c configuration.Provider, r MutatorIDTokenRegistry) *MutatorIDToken {
-	cache, _ := ristretto.NewCache(&ristretto.Config[string, *idTokenCacheContainer]{
-		NumCounters: 10000,
-		MaxCost:     1 << 25,
-		BufferItems: 64,
-	})
-	return &MutatorIDToken{r: r, c: c, templates: x.NewTemplate("id_token"), tokenCache: cache, tokenCacheEnabled: true}
+	return &MutatorIDToken{r: r, c: c, templates: x.NewTemplate("id_token")}
 }
 
 func (a *MutatorIDToken) GetID() string {
@@ -70,10 +70,6 @@ func (a *MutatorIDToken) WithCache(t *template.Template) {
 	a.templates = t
 }
 
-func (a *MutatorIDToken) SetCaching(token bool) {
-	a.tokenCacheEnabled = token
-}
-
 type idTokenCacheContainer struct {
 	ExpiresAt time.Time
 	Token     string
@@ -86,7 +82,7 @@ func (a *MutatorIDToken) cacheKey(config *CredentialsIDTokenConfig, ttl time.Dur
 }
 
 func (a *MutatorIDToken) tokenFromCache(config *CredentialsIDTokenConfig, session *authn.AuthenticationSession, claims []byte, ttl time.Duration) (string, bool) {
-	if !a.tokenCacheEnabled {
+	if !config.Cache.Enabled {
 		return "", false
 	}
 
@@ -106,7 +102,7 @@ func (a *MutatorIDToken) tokenFromCache(config *CredentialsIDTokenConfig, sessio
 }
 
 func (a *MutatorIDToken) tokenToCache(config *CredentialsIDTokenConfig, session *authn.AuthenticationSession, claims []byte, ttl time.Duration, expiresAt time.Time, token string) {
-	if !a.tokenCacheEnabled {
+	if !config.Cache.Enabled {
 		return
 	}
 
@@ -197,7 +193,11 @@ func (a *MutatorIDToken) Validate(config json.RawMessage) error {
 }
 
 func (a *MutatorIDToken) Config(config json.RawMessage) (*CredentialsIDTokenConfig, error) {
-	var c CredentialsIDTokenConfig
+	c := CredentialsIDTokenConfig{
+		Cache: IdTokenCacheConfig{
+			Enabled: true, // default to true
+		},
+	}
 	if err := a.c.MutatorConfig(a.GetID(), config, &c); err != nil {
 		return nil, NewErrMutatorMisconfigured(a, err)
 	}
@@ -206,5 +206,29 @@ func (a *MutatorIDToken) Config(config json.RawMessage) (*CredentialsIDTokenConf
 		c.TTL = "15m"
 	}
 
+	cost := int64(c.Cache.MaxCost)
+	if cost == 0 {
+		cost = 1 << 25
+	}
+
+	if a.tokenCache == nil || a.tokenCache.MaxCost() != cost {
+		cache, err := ristretto.NewCache(&ristretto.Config[string, *idTokenCacheContainer]{
+			// Guessed approximation of max number of items.
+			NumCounters: cost * 4,
+			// Allocate a max
+			MaxCost: cost,
+			// This is a best-practice value.
+			BufferItems: 64,
+			Cost: func(container *idTokenCacheContainer) int64 {
+				return int64(len(container.Token))
+			},
+		})
+
+		if err != nil {
+			return nil, err
+		}
+		a.tokenCache = cache
+	}
+
 	return &c, nil
 }
diff --git a/pipeline/mutate/mutator_id_token_test.go b/pipeline/mutate/mutator_id_token_test.go
@@ -300,6 +300,18 @@ func TestMutatorIDToken(t *testing.T) {
 				config, _ := sjson.SetBytes(config, "jwks_url", "file://../../test/stub/jwks-hs.json")
 				assert.NotEqual(t, prev, mutate(t, *session, config))
 			})
+
+			t.Run("subcase=different tokens because cache disabled", func(t *testing.T) {
+				config, _ := sjson.SetBytes(config, "cache", map[string]bool{"enabled": false})
+				prev := mutate(t, *session, config)
+				assert.NotEqual(t, prev, mutate(t, *session, config))
+			})
+
+			t.Run("subcase=different tokens because exceeded cost", func(t *testing.T) {
+				config, _ := sjson.SetBytes(config, "cache", map[string]int{"max_cost": 1})
+				prev := mutate(t, *session, config)
+				assert.NotEqual(t, prev, mutate(t, *session, config))
+			})
 		})
 
 		t.Run("case=ensure template cache", func(t *testing.T) {
@@ -386,8 +398,8 @@ func BenchmarkMutatorIDToken(b *testing.B) {
 	} {
 		b.Run("alg="+alg, func(b *testing.B) {
 			for _, enableCache := range []bool{true, false} {
-				a.(*MutatorIDToken).SetCaching(enableCache)
 				b.Run(fmt.Sprintf("cache=%v", enableCache), func(b *testing.B) {
+					conf.SetForTest(b, "mutators.id_token.config.cache.enabled", enableCache)
 					var tc idTokenTestCase
 					var config []byte
 
diff --git a/spec/config.schema.json b/spec/config.schema.json
@@ -1140,6 +1140,25 @@
           "pattern": "^[0-9]+(ns|us|ms|s|m|h)$",
           "default": "15m",
           "examples": ["1h", "1m", "30s"]
+        },
+        "cache": {
+          "additionalProperties": false,
+          "type": "object",
+          "properties": {
+            "enabled": {
+              "title": "Enabled",
+              "type": "boolean",
+              "default": true,
+              "examples": [false, true],
+              "description": "En-/disables this component."
+            },
+            "max_cost": {
+              "type": "integer",
+              "default": 33554432,
+              "title": "Maximum Cached Cost",
+              "description": "The cost of one cached JSON Web Token is the length of its string form."
+            }
+          }
         }
       },
       "additionalProperties": false
