feat: make id_token mutator cache configurable

ory · Aug 9, 2024 · 6ad3106 · 6ad3106
1 parent 6b5672b
commit 6ad3106
Show file tree

Hide file tree

Showing 4 changed files with 82 additions and 20 deletions.
diff --git a/.schemas/config.schema.json b/.schemas/config.schema.json
@@ -823,6 +823,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": "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
-	tokenCacheEnabled bool
+	tokenCache *ristretto.Cache
 }
 
 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{
-		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
 	TTL       time.Duration
@@ -87,7 +83,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
 	}
 
@@ -108,7 +104,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
 	}
 
@@ -195,7 +191,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)
 	}
@@ -204,5 +204,29 @@ func (a *MutatorIDToken) Config(config json.RawMessage) (*CredentialsIDTokenConf
 		c.TTL = "15m"
 	}
 
+	if a.tokenCache == nil {
+		cost := int64(c.Cache.MaxCost)
+		if cost == 0 {
+			cost = 1 << 25
+		}
+
+		cache, err := ristretto.NewCache(&ristretto.Config{
+			NumCounters: cost * 10,
+			// Allocate a max
+			MaxCost: cost,
+			// This is a best-practice value.
+			BufferItems: 64,
+			Cost: func(value interface{}) int64 {
+				return 1
+			},
+			IgnoreInternalCost: true,
+		})
+
+		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
@@ -386,8 +386,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", true)
 					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": "Max cost to cache."
+            }
+          }
         }
       },
       "additionalProperties": false