Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 124120c

Browse files
David-WobrockDavid-Wobrock
committedJan 2, 2025··
feat: make id_token mutator cache configurable
1 parent 517f326 commit 124120c

File tree

5 files changed

+113
-20
lines changed

5 files changed

+113
-20
lines changed
 

‎.schema/config.schema.json

+19
Original file line numberDiff line numberDiff line change
@@ -1140,6 +1140,25 @@
11401140
"pattern": "^[0-9]+(ns|us|ms|s|m|h)$",
11411141
"default": "15m",
11421142
"examples": ["1h", "1m", "30s"]
1143+
},
1144+
"cache": {
1145+
"additionalProperties": false,
1146+
"type": "object",
1147+
"properties": {
1148+
"enabled": {
1149+
"title": "Enabled",
1150+
"type": "boolean",
1151+
"default": true,
1152+
"examples": [false, true],
1153+
"description": "En-/disables this component."
1154+
},
1155+
"max_cost": {
1156+
"type": "integer",
1157+
"default": 33554432,
1158+
"title": "Maximum Cached Cost",
1159+
"description": "The cost of one cached JSON Web Token is the length of its string form."
1160+
}
1161+
}
11431162
}
11441163
},
11451164
"additionalProperties": false

‎.schemas/mutators.id_token.schema.json

+19
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,25 @@
3232
"pattern": "^[0-9]+(ns|us|ms|s|m|h)$",
3333
"default": "1m",
3434
"examples": ["1h", "1m", "30s"]
35+
},
36+
"cache": {
37+
"additionalProperties": false,
38+
"type": "object",
39+
"properties": {
40+
"enabled": {
41+
"title": "Enabled",
42+
"type": "boolean",
43+
"default": true,
44+
"examples": [false, true],
45+
"description": "En-/disables this component."
46+
},
47+
"max_cost": {
48+
"type": "integer",
49+
"default": 33554432,
50+
"title": "Maximum Cached Cost",
51+
"description": "Max cost to cache."
52+
}
53+
}
3554
}
3655
},
3756
"additionalProperties": false

‎pipeline/mutate/mutator_id_token.go

+43-19
Original file line numberDiff line numberDiff line change
@@ -38,28 +38,28 @@ type MutatorIDToken struct {
3838
templates *template.Template
3939
templatesLock sync.Mutex
4040

41-
tokenCache *ristretto.Cache[string, *idTokenCacheContainer]
42-
tokenCacheEnabled bool
41+
tokenCache *ristretto.Cache[string, *idTokenCacheContainer]
4342
}
4443

4544
type CredentialsIDTokenConfig struct {
46-
Claims string `json:"claims"`
47-
IssuerURL string `json:"issuer_url"`
48-
JWKSURL string `json:"jwks_url"`
49-
TTL string `json:"ttl"`
45+
Claims string `json:"claims"`
46+
IssuerURL string `json:"issuer_url"`
47+
JWKSURL string `json:"jwks_url"`
48+
TTL string `json:"ttl"`
49+
Cache IdTokenCacheConfig `json:"cache"`
50+
}
51+
52+
type IdTokenCacheConfig struct {
53+
Enabled bool `json:"enabled"`
54+
MaxCost int `json:"max_cost"`
5055
}
5156

5257
func (c *CredentialsIDTokenConfig) ClaimsTemplateID() string {
5358
return fmt.Sprintf("%x", md5.Sum([]byte(c.Claims)))
5459
}
5560

5661
func NewMutatorIDToken(c configuration.Provider, r MutatorIDTokenRegistry) *MutatorIDToken {
57-
cache, _ := ristretto.NewCache(&ristretto.Config[string, *idTokenCacheContainer]{
58-
NumCounters: 10000,
59-
MaxCost: 1 << 25,
60-
BufferItems: 64,
61-
})
62-
return &MutatorIDToken{r: r, c: c, templates: x.NewTemplate("id_token"), tokenCache: cache, tokenCacheEnabled: true}
62+
return &MutatorIDToken{r: r, c: c, templates: x.NewTemplate("id_token")}
6363
}
6464

6565
func (a *MutatorIDToken) GetID() string {
@@ -70,10 +70,6 @@ func (a *MutatorIDToken) WithCache(t *template.Template) {
7070
a.templates = t
7171
}
7272

73-
func (a *MutatorIDToken) SetCaching(token bool) {
74-
a.tokenCacheEnabled = token
75-
}
76-
7773
type idTokenCacheContainer struct {
7874
ExpiresAt time.Time
7975
Token string
@@ -86,7 +82,7 @@ func (a *MutatorIDToken) cacheKey(config *CredentialsIDTokenConfig, ttl time.Dur
8682
}
8783

8884
func (a *MutatorIDToken) tokenFromCache(config *CredentialsIDTokenConfig, session *authn.AuthenticationSession, claims []byte, ttl time.Duration) (string, bool) {
89-
if !a.tokenCacheEnabled {
85+
if !config.Cache.Enabled {
9086
return "", false
9187
}
9288

@@ -106,7 +102,7 @@ func (a *MutatorIDToken) tokenFromCache(config *CredentialsIDTokenConfig, sessio
106102
}
107103

108104
func (a *MutatorIDToken) tokenToCache(config *CredentialsIDTokenConfig, session *authn.AuthenticationSession, claims []byte, ttl time.Duration, expiresAt time.Time, token string) {
109-
if !a.tokenCacheEnabled {
105+
if !config.Cache.Enabled {
110106
return
111107
}
112108

@@ -197,7 +193,11 @@ func (a *MutatorIDToken) Validate(config json.RawMessage) error {
197193
}
198194

199195
func (a *MutatorIDToken) Config(config json.RawMessage) (*CredentialsIDTokenConfig, error) {
200-
var c CredentialsIDTokenConfig
196+
c := CredentialsIDTokenConfig{
197+
Cache: IdTokenCacheConfig{
198+
Enabled: true, // default to true
199+
},
200+
}
201201
if err := a.c.MutatorConfig(a.GetID(), config, &c); err != nil {
202202
return nil, NewErrMutatorMisconfigured(a, err)
203203
}
@@ -206,5 +206,29 @@ func (a *MutatorIDToken) Config(config json.RawMessage) (*CredentialsIDTokenConf
206206
c.TTL = "15m"
207207
}
208208

209+
cost := int64(c.Cache.MaxCost)
210+
if cost == 0 {
211+
cost = 1 << 25
212+
}
213+
214+
if a.tokenCache == nil || a.tokenCache.MaxCost() != cost {
215+
cache, err := ristretto.NewCache(&ristretto.Config[string, *idTokenCacheContainer]{
216+
// Guessed approximation of max number of items.
217+
NumCounters: cost * 4,
218+
// Allocate a max
219+
MaxCost: cost,
220+
// This is a best-practice value.
221+
BufferItems: 64,
222+
Cost: func(container *idTokenCacheContainer) int64 {
223+
return int64(len(container.Token))
224+
},
225+
})
226+
227+
if err != nil {
228+
return nil, err
229+
}
230+
a.tokenCache = cache
231+
}
232+
209233
return &c, nil
210234
}

‎pipeline/mutate/mutator_id_token_test.go

+13-1
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,18 @@ func TestMutatorIDToken(t *testing.T) {
300300
config, _ := sjson.SetBytes(config, "jwks_url", "file://../../test/stub/jwks-hs.json")
301301
assert.NotEqual(t, prev, mutate(t, *session, config))
302302
})
303+
304+
t.Run("subcase=different tokens because cache disabled", func(t *testing.T) {
305+
config, _ := sjson.SetBytes(config, "cache", map[string]bool{"enabled": false})
306+
prev := mutate(t, *session, config)
307+
assert.NotEqual(t, prev, mutate(t, *session, config))
308+
})
309+
310+
t.Run("subcase=different tokens because exceeded cost", func(t *testing.T) {
311+
config, _ := sjson.SetBytes(config, "cache", map[string]int{"max_cost": -1})
312+
prev := mutate(t, *session, config)
313+
assert.NotEqual(t, prev, mutate(t, *session, config))
314+
})
303315
})
304316

305317
t.Run("case=ensure template cache", func(t *testing.T) {
@@ -386,8 +398,8 @@ func BenchmarkMutatorIDToken(b *testing.B) {
386398
} {
387399
b.Run("alg="+alg, func(b *testing.B) {
388400
for _, enableCache := range []bool{true, false} {
389-
a.(*MutatorIDToken).SetCaching(enableCache)
390401
b.Run(fmt.Sprintf("cache=%v", enableCache), func(b *testing.B) {
402+
conf.SetForTest(b, "mutators.id_token.config.cache.enabled", enableCache)
391403
var tc idTokenTestCase
392404
var config []byte
393405

‎spec/config.schema.json

+19
Original file line numberDiff line numberDiff line change
@@ -1140,6 +1140,25 @@
11401140
"pattern": "^[0-9]+(ns|us|ms|s|m|h)$",
11411141
"default": "15m",
11421142
"examples": ["1h", "1m", "30s"]
1143+
},
1144+
"cache": {
1145+
"additionalProperties": false,
1146+
"type": "object",
1147+
"properties": {
1148+
"enabled": {
1149+
"title": "Enabled",
1150+
"type": "boolean",
1151+
"default": true,
1152+
"examples": [false, true],
1153+
"description": "En-/disables this component."
1154+
},
1155+
"max_cost": {
1156+
"type": "integer",
1157+
"default": 33554432,
1158+
"title": "Maximum Cached Cost",
1159+
"description": "The cost of one cached JSON Web Token is the length of its string form."
1160+
}
1161+
}
11431162
}
11441163
},
11451164
"additionalProperties": false

0 commit comments

Comments
 (0)
Please sign in to comment.