Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: make id_token mutator cache configurable #1177

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions .schema/config.schema.json
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
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
76 changes: 51 additions & 25 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[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,13 +70,8 @@ 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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven't looked at this code in detail, but ExpiresAt+TTL seem to be semantically the same thing? Do we even need to store these in the cache?

Copy link
Contributor Author

@David-Wobrock David-Wobrock Sep 13, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, we don't!

I think that's what was buggy in the first place: we don't set an actual TTL in the cache, but instead cache the TTL along with the token.

I pushed an edit.

TTL time.Duration
Token string
}

Expand All @@ -87,7 +82,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,16 +102,20 @@ 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
}

key := a.cacheKey(config, ttl, claims, session)
a.tokenCache.Set(key, &idTokenCacheContainer{
TTL: ttl,
ExpiresAt: expiresAt,
Token: token,
}, 0)
a.tokenCache.SetWithTTL(
key,
&idTokenCacheContainer{
ExpiresAt: expiresAt,
Token: token,
},
0,
ttl,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In other caches, I believe we set the TTL to min(TTL, time.Until(expiresAt). That would make sense here too IMO.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

They should be same normally. Since calling this method does:

	now := time.Now().UTC()
	exp := now.Add(ttl)
	[...]
	a.tokenToCache(c, session, templateClaims, ttl, exp, signed)

Or what am I missing? :)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

)
}

func (a *MutatorIDToken) Mutate(r *http.Request, session *authn.AuthenticationSession, config json.RawMessage, rl pipeline.Rule) error {
Expand Down Expand Up @@ -194,7 +193,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 +206,28 @@ 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(container *idTokenCacheContainer) int64 {
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
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
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": {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We have a bunch of other caches in the config schema already. Can you make your change so that it is more similar (identical) to those other cache configurations? The max_cost parameter in particular is really opaque and its impossible to come up with a reasonable value without knowing the implementation.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with the remark on max_cost 👍
A value with similar semantics is used today for the OAuth2 introspection authenticator, so I mainly mimicked this one.
Else, there is also max_tokens, used by the OAuth2 client credentials authenticator.

I think the reason for the cost instead comes from the fact that the cached objects have variable lengths, so storing a certain number of objects will result in a different cache memory usage depending on your config.

However, one can also make the decision to let the user make this decision anyway :)
Let me know what makes most sense to you, and Ory's strategy around these questions (should the product have the same defaults for everyone, or is the user trusted to configure this accordingly).


For the enabled/disabled value, I didn't re-use the existing

    "handlerSwitch": {
      "title": "Enabled",
      "type": "boolean",
      "default": false,
      "examples": [true],
      "description": "En-/disables this component."
    },

to avoid introducing a breaking change.

Currently, the id_token cache is enabled by default, and didn't wanna change the default value to false - so I couldn't re-use this configuration.


And finally, this cache config has no TTL, because the id_token mutator already has a TTL config value for the JWT expiration date.
It make sense to me to re-use the same value => cache for 15 min if the JWT is valid 15min.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ideally, we would probably have one namespaced cache for all of oathkeeper and then use that everywhere, but this is good for now!

"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
Loading