From c32694ff67dae7654a85a8a5e063035d7036500b Mon Sep 17 00:00:00 2001 From: Alex Guerrieri Date: Thu, 19 Sep 2024 15:08:02 +0200 Subject: [PATCH 1/7] Use `WithRequestLimit` with `0` to skip rate limit --- context.go | 16 ++++++---------- go.mod | 2 +- go.sum | 2 -- limiter.go | 24 ++++++++++++++++-------- 4 files changed, 23 insertions(+), 21 deletions(-) diff --git a/context.go b/context.go index 138110d..3988dd6 100644 --- a/context.go +++ b/context.go @@ -13,20 +13,16 @@ func WithIncrement(ctx context.Context, value int) context.Context { return context.WithValue(ctx, incrementKey, value) } -func getIncrement(ctx context.Context) int { - if value, ok := ctx.Value(incrementKey).(int); ok { - return value - } - return 1 +func getIncrement(ctx context.Context) (int, bool) { + value, ok := ctx.Value(incrementKey).(int) + return value, ok } func WithRequestLimit(ctx context.Context, value int) context.Context { return context.WithValue(ctx, requestLimitKey, value) } -func getRequestLimit(ctx context.Context) int { - if value, ok := ctx.Value(requestLimitKey).(int); ok { - return value - } - return 0 +func getRequestLimit(ctx context.Context) (int, bool) { + value, ok := ctx.Value(requestLimitKey).(int) + return value, ok } diff --git a/go.mod b/go.mod index 998cbf5..a63a858 100644 --- a/go.mod +++ b/go.mod @@ -4,4 +4,4 @@ go 1.17 require github.com/cespare/xxhash/v2 v2.3.0 -require golang.org/x/sync v0.7.0 // indirect +require golang.org/x/sync v0.7.0 diff --git a/go.sum b/go.sum index 09aebbf..2cc97dd 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,3 @@ -github.com/cespare/xxhash/v2 v2.1.2 h1:YRXhKfTDauu4ajMg1TPgFO5jnlC2HCbmLXMcTG5cbYE= -github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M= diff --git a/limiter.go b/limiter.go index d5ef436..fe52545 100644 --- a/limiter.go +++ b/limiter.go @@ -72,13 +72,26 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string currentWindow := time.Now().UTC().Truncate(l.windowLength) ctx := r.Context() - limit := l.requestLimit - if val := getRequestLimit(ctx); val > 0 { - limit = val + limit, ok := getRequestLimit(ctx) + if !ok { + limit = l.requestLimit } + + if limit <= 0 { + return false + } + setHeader(w, l.headers.Limit, fmt.Sprintf("%d", limit)) setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix())) + increment, ok := getIncrement(r.Context()) + if !ok { + increment = 1 + } + if increment > 1 { + setHeader(w, l.headers.Increment, fmt.Sprintf("%d", increment)) + } + l.mu.Lock() _, rateFloat, err := l.calculateRate(key, limit) if err != nil { @@ -88,11 +101,6 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string } rate := int(math.Round(rateFloat)) - increment := getIncrement(r.Context()) - if increment > 1 { - setHeader(w, l.headers.Increment, fmt.Sprintf("%d", increment)) - } - if rate+increment > limit { setHeader(w, l.headers.Remaining, fmt.Sprintf("%d", limit-rate)) From 7eaea8fb5c19385eb787171e586d116fd7034bb4 Mon Sep 17 00:00:00 2001 From: Alex Guerrieri Date: Thu, 19 Sep 2024 15:53:01 +0200 Subject: [PATCH 2/7] NoCost --- context.go | 6 ++++++ limiter.go | 20 +++++++++++++++----- limiter_test.go | 4 ++-- 3 files changed, 23 insertions(+), 7 deletions(-) diff --git a/context.go b/context.go index 3988dd6..93965e6 100644 --- a/context.go +++ b/context.go @@ -9,6 +9,8 @@ const ( requestLimitKey ) +const _NoLimit = -1 + func WithIncrement(ctx context.Context, value int) context.Context { return context.WithValue(ctx, incrementKey, value) } @@ -18,6 +20,10 @@ func getIncrement(ctx context.Context) (int, bool) { return value, ok } +func WithNoLimit(ctx context.Context) context.Context { + return context.WithValue(ctx, requestLimitKey, _NoLimit) +} + func WithRequestLimit(ctx context.Context, value int) context.Context { return context.WithValue(ctx, requestLimitKey, value) } diff --git a/limiter.go b/limiter.go index fe52545..011f586 100644 --- a/limiter.go +++ b/limiter.go @@ -72,22 +72,32 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string currentWindow := time.Now().UTC().Truncate(l.windowLength) ctx := r.Context() + setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix())) + limit, ok := getRequestLimit(ctx) if !ok { limit = l.requestLimit } - - if limit <= 0 { + // If the limit is set to 0, we are always over limit + if limit == 0 { + return true + } + // If the limit is set to -1, we are never over limit + if limit == _NoLimit { return false } - setHeader(w, l.headers.Limit, fmt.Sprintf("%d", limit)) - setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix())) - increment, ok := getIncrement(r.Context()) if !ok { increment = 1 } + // If the increment is 0, we are always on limit + if increment == 0 { + return false + } + + setHeader(w, l.headers.Limit, fmt.Sprintf("%d", limit)) + if increment > 1 { setHeader(w, l.headers.Increment, fmt.Sprintf("%d", increment)) } diff --git a/limiter_test.go b/limiter_test.go index 5ac41c1..cc4a249 100644 --- a/limiter_test.go +++ b/limiter_test.go @@ -141,8 +141,8 @@ func TestResponseHeaders(t *testing.T) { requestsLimit: 5, increments: []int{0, 0, 0, 0, 0, 0}, respCodes: []int{200, 200, 200, 200, 200, 200}, - respLimitHeader: []string{"5", "5", "5", "5", "5", "5"}, - respRemainingHeader: []string{"5", "5", "5", "5", "5", "5"}, + respLimitHeader: []string{"", "", "", "", "", ""}, + respRemainingHeader: []string{"", "", "", "", "", ""}, }, { name: "always block", From ac748ea5a6b8921aef7a442985c77657a989932f Mon Sep 17 00:00:00 2001 From: Alex Guerrieri Date: Fri, 20 Sep 2024 10:16:08 +0200 Subject: [PATCH 3/7] cr suggestion --- context.go | 4 +--- limiter.go | 8 ++++---- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/context.go b/context.go index 93965e6..3e2c2bb 100644 --- a/context.go +++ b/context.go @@ -9,8 +9,6 @@ const ( requestLimitKey ) -const _NoLimit = -1 - func WithIncrement(ctx context.Context, value int) context.Context { return context.WithValue(ctx, incrementKey, value) } @@ -21,7 +19,7 @@ func getIncrement(ctx context.Context) (int, bool) { } func WithNoLimit(ctx context.Context) context.Context { - return context.WithValue(ctx, requestLimitKey, _NoLimit) + return context.WithValue(ctx, requestLimitKey, -1) } func WithRequestLimit(ctx context.Context, value int) context.Context { diff --git a/limiter.go b/limiter.go index 011f586..c893ec2 100644 --- a/limiter.go +++ b/limiter.go @@ -78,12 +78,12 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string if !ok { limit = l.requestLimit } - // If the limit is set to 0, we are always over limit + // If the limit is set to 0, we always limit if limit == 0 { return true } - // If the limit is set to -1, we are never over limit - if limit == _NoLimit { + // If the limit is set to -1, there is no limit + if limit == -1 { return false } @@ -91,7 +91,7 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string if !ok { increment = 1 } - // If the increment is 0, we are always on limit + // If the increment is 0, we do not limit if increment == 0 { return false } From 9719e456e4a20cd3b87ba59ec0ef52ebd28eb188 Mon Sep 17 00:00:00 2001 From: Alex Guerrieri Date: Fri, 20 Sep 2024 10:20:41 +0200 Subject: [PATCH 4/7] more cr suggestions --- limiter.go | 10 ++++------ limiter_test.go | 2 +- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/limiter.go b/limiter.go index c893ec2..a6077dc 100644 --- a/limiter.go +++ b/limiter.go @@ -86,22 +86,20 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string if limit == -1 { return false } + setHeader(w, l.headers.Limit, fmt.Sprintf("%d", limit)) increment, ok := getIncrement(r.Context()) if !ok { increment = 1 } + if increment != 1 { + setHeader(w, l.headers.Increment, fmt.Sprintf("%d", increment)) + } // If the increment is 0, we do not limit if increment == 0 { return false } - setHeader(w, l.headers.Limit, fmt.Sprintf("%d", limit)) - - if increment > 1 { - setHeader(w, l.headers.Increment, fmt.Sprintf("%d", increment)) - } - l.mu.Lock() _, rateFloat, err := l.calculateRate(key, limit) if err != nil { diff --git a/limiter_test.go b/limiter_test.go index cc4a249..50d7b07 100644 --- a/limiter_test.go +++ b/limiter_test.go @@ -141,7 +141,7 @@ func TestResponseHeaders(t *testing.T) { requestsLimit: 5, increments: []int{0, 0, 0, 0, 0, 0}, respCodes: []int{200, 200, 200, 200, 200, 200}, - respLimitHeader: []string{"", "", "", "", "", ""}, + respLimitHeader: []string{"5", "5", "5", "5", "5", "5"}, respRemainingHeader: []string{"", "", "", "", "", ""}, }, { From 1677b6d6cd15e710b91c54c395519186d8fffb0a Mon Sep 17 00:00:00 2001 From: Alex Guerrieri Date: Fri, 20 Sep 2024 10:22:12 +0200 Subject: [PATCH 5/7] set headers if we are actually RLing --- limiter.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/limiter.go b/limiter.go index a6077dc..0a80ceb 100644 --- a/limiter.go +++ b/limiter.go @@ -72,8 +72,6 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string currentWindow := time.Now().UTC().Truncate(l.windowLength) ctx := r.Context() - setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix())) - limit, ok := getRequestLimit(ctx) if !ok { limit = l.requestLimit @@ -86,6 +84,8 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string if limit == -1 { return false } + + setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix())) setHeader(w, l.headers.Limit, fmt.Sprintf("%d", limit)) increment, ok := getIncrement(r.Context()) From 37be242073949f8ed3e6c127e913108c3e684fc3 Mon Sep 17 00:00:00 2001 From: Alex Guerrieri Date: Fri, 20 Sep 2024 10:27:58 +0200 Subject: [PATCH 6/7] move headers set --- limiter.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/limiter.go b/limiter.go index 0a80ceb..c1cb23c 100644 --- a/limiter.go +++ b/limiter.go @@ -76,8 +76,12 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string if !ok { limit = l.requestLimit } + setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix())) + setHeader(w, l.headers.Limit, fmt.Sprintf("%d", limit)) + // If the limit is set to 0, we always limit if limit == 0 { + setHeader(w, l.headers.Remaining, "0") return true } // If the limit is set to -1, there is no limit @@ -85,9 +89,6 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string return false } - setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix())) - setHeader(w, l.headers.Limit, fmt.Sprintf("%d", limit)) - increment, ok := getIncrement(r.Context()) if !ok { increment = 1 From 8d182db7564e436280511fab752db97c45021ec8 Mon Sep 17 00:00:00 2001 From: Alex Guerrieri Date: Fri, 20 Sep 2024 10:29:24 +0200 Subject: [PATCH 7/7] move header up --- limiter.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/limiter.go b/limiter.go index c1cb23c..eba48b7 100644 --- a/limiter.go +++ b/limiter.go @@ -76,8 +76,8 @@ func (l *RateLimiter) OnLimit(w http.ResponseWriter, r *http.Request, key string if !ok { limit = l.requestLimit } - setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix())) setHeader(w, l.headers.Limit, fmt.Sprintf("%d", limit)) + setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix())) // If the limit is set to 0, we always limit if limit == 0 {