Skip to content

Commit

Permalink
feat: make id_token mutator cache configurable
Browse files Browse the repository at this point in the history
  • Loading branch information
@David-Wobrock
David-Wobrock committed Nov 11, 2024
1 parent 98f8a00 commit 5aa1069
Show file tree
Hide file tree
Showing 5 changed files with 101 additions and 20 deletions.
19 changes: 19 additions & 0 deletions .schema/config.schema.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
19 changes: 19 additions & 0 deletions .schemas/mutators.id_token.schema.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
62 changes: 43 additions & 19 deletions pipeline/mutate/mutator_id_token.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand All @@ -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
}

Expand All @@ -107,7 +103,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
}

Expand Down Expand Up @@ -194,7 +190,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)
}
Expand All @@ -203,5 +203,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[string, *idTokenCacheContainer]{
NumCounters: cost * 10,
// Allocate a max
MaxCost: cost,
// This is a best-practice value.
BufferItems: 64,
Cost: func(value interface{}) int64 {
container := value.(*idTokenCacheContainer)
return int64(len(container.Token))
},
})

if err != nil {
return nil, err
}
a.tokenCache = cache
}

return &c, nil
}
2 changes: 1 addition & 1 deletion pipeline/mutate/mutator_id_token_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
19 changes: 19 additions & 0 deletions spec/config.schema.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 5aa1069

Please sign in to comment.