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 committed Aug 9, 2024
1 parent 6b5672b commit 6146085
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 20 deletions.
62 changes: 43 additions & 19 deletions pipeline/mutate/mutator_id_token.go
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
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 {
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 @@ -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
}

Expand Down Expand Up @@ -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)
}
Expand All @@ -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
}
2 changes: 1 addition & 1 deletion pipeline/mutate/mutator_id_token_test.go
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

0 comments on commit 6146085

Please sign in to comment.