From 7a46fd0b3674f30f24e6a924751c87fae7129839 Mon Sep 17 00:00:00 2001 From: David Wobrock Date: Fri, 9 Aug 2024 16:51:09 +0200 Subject: [PATCH] feat: make id_token mutator cache configurable --- .schema/config.schema.json | 19 ++++++++ .schemas/config.schema.json | 19 ++++++++ .schemas/mutators.id_token.schema.json | 19 ++++++++ pipeline/mutate/mutator_id_token.go | 62 ++++++++++++++++-------- pipeline/mutate/mutator_id_token_test.go | 2 +- spec/config.schema.json | 19 ++++++++ 6 files changed, 120 insertions(+), 20 deletions(-) diff --git a/.schema/config.schema.json b/.schema/config.schema.json index 4dcd8a8032..84f1ace417 100644 --- a/.schema/config.schema.json +++ b/.schema/config.schema.json @@ -1118,6 +1118,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/.schemas/config.schema.json b/.schemas/config.schema.json index 216788c8cd..5f6a112e91 100644 --- 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/.schemas/mutators.id_token.schema.json b/.schemas/mutators.id_token.schema.json index c160f0bc08..ef07faa17a 100644 --- 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 index 476e9a4c63..4fd6241b39 100644 --- a/pipeline/mutate/mutator_id_token.go +++ b/pipeline/mutate/mutator_id_token.go @@ -38,15 +38,20 @@ 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 { @@ -54,12 +59,7 @@ func (c *CredentialsIDTokenConfig) ClaimsTemplateID() string { } 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 index 2666961307..43b419a78c 100644 --- 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 index ce6ef286d5..177b0899db 100644 --- 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