Use WithRequestLimit with 0 to skip rate limit

context.go

@@ -9,24 +9,26 @@ const (
 	requestLimitKey
 )
 
+const _NoLimit = -1
+
 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 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)
 }
 
-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
 }

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

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=

limiter.go

@@ -72,12 +72,35 @@ 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
+	setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix()))
+
+	limit, ok := getRequestLimit(ctx)
+	if !ok {
+		limit = l.requestLimit
+	}
+	// 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
 	}
+
+	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))
-	setHeader(w, l.headers.Reset, fmt.Sprintf("%d", currentWindow.Add(l.windowLength).Unix()))
+
+	if increment > 1 {
+		setHeader(w, l.headers.Increment, fmt.Sprintf("%d", increment))
+	}
 
 	l.mu.Lock()
 	_, rateFloat, err := l.calculateRate(key, limit)
@@ -88,11 +111,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))
 

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