From 3c6783b9cf8fedf22f8880b046a9692c1b160cff Mon Sep 17 00:00:00 2001 From: Nikos Date: Thu, 28 Mar 2024 18:19:14 +0200 Subject: [PATCH] fix: token endpoint rate limiting --- handler/rfc8628/strategy_hmacsha.go | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/handler/rfc8628/strategy_hmacsha.go b/handler/rfc8628/strategy_hmacsha.go index 4bc1ae22..6f043890 100644 --- a/handler/rfc8628/strategy_hmacsha.go +++ b/handler/rfc8628/strategy_hmacsha.go @@ -20,6 +20,8 @@ import ( enigma "github.com/ory/fosite/token/hmac" ) +const POLLING_RATE_LIMITING_LEEWAY = 200 + // DeviceFlowSession is a fosite.Session container specific for the device flow. type DeviceFlowSession interface { // GetBrowserFlowCompleted returns the flag indicating whether user has completed the browser flow or not. @@ -182,10 +184,10 @@ func (h *DefaultDeviceStrategy) ShouldRateLimit(context context.Context, code st timer.NotUntil = h.getExpirationTime(context, 1) exp, err := h.serializeExpiration(timer) if err != nil { - return false, err + return false, errorsx.WithStack(fosite.ErrServerError.WithHintf("Failed to serialize expiration struct %s", err)) } // Set the expiration time as value, and use the lifespan of the device code as TTL. - h.RateLimiterCache.Set(keyBytes, exp, int(h.Config.GetDeviceAndUserCodeLifespan(context))) + h.RateLimiterCache.Set(keyBytes, exp, int(h.Config.GetDeviceAndUserCodeLifespan(context).Seconds())) return false, nil } @@ -195,13 +197,13 @@ func (h *DefaultDeviceStrategy) ShouldRateLimit(context context.Context, code st } // The code is valid and enough time has passed since the last call. - if expiration.NotUntil.Before(time.Now()) { + if time.Since(expiration.NotUntil).Milliseconds() > -POLLING_RATE_LIMITING_LEEWAY { expiration.NotUntil = h.getExpirationTime(context, expiration.Counter) exp, err := h.serializeExpiration(expiration) if err != nil { - return false, err + return false, errorsx.WithStack(fosite.ErrServerError.WithHintf("Failed to serialize expiration struct %s", err)) } - h.RateLimiterCache.Set(keyBytes, exp, int(h.Config.GetDeviceAndUserCodeLifespan(context))) + h.RateLimiterCache.Set(keyBytes, exp, int(h.Config.GetDeviceAndUserCodeLifespan(context).Seconds())) return false, nil } @@ -210,9 +212,9 @@ func (h *DefaultDeviceStrategy) ShouldRateLimit(context context.Context, code st expiration.Counter += 1 exp, err := h.serializeExpiration(expiration) if err != nil { - return false, err + return false, errorsx.WithStack(fosite.ErrServerError.WithHintf("Failed to serialize expiration struct %s", err)) } - h.RateLimiterCache.Set(keyBytes, exp, int(h.Config.GetDeviceAndUserCodeLifespan(context))) + h.RateLimiterCache.Set(keyBytes, exp, int(h.Config.GetDeviceAndUserCodeLifespan(context).Seconds())) return true, nil }