diff --git a/internal/middleware/ratelimiter.go b/internal/middleware/ratelimiter.go index 293dd9c..944b6d5 100644 --- a/internal/middleware/ratelimiter.go +++ b/internal/middleware/ratelimiter.go @@ -69,6 +69,7 @@ func RateLimiter(client *db.DiceDB, next http.Handler, limit, window int) http.H // Check if the request count exceeds the limit if requestCount >= limit { slog.Warn("Request limit exceeded", "count", requestCount) + addRateLimitHeaders(w, limit, limit-(requestCount+1), requestCount+1, currentWindow+int64(window)) http.Error(w, "429 - Too Many Requests", http.StatusTooManyRequests) return } @@ -77,6 +78,7 @@ func RateLimiter(client *db.DiceDB, next http.Handler, limit, window int) http.H if _, err := client.Client.Incr(ctx, key).Result(); err != nil { slog.Error("Error incrementing request count", "error", err) http.Error(w, "Internal Server Error", http.StatusInternalServerError) + addRateLimitHeaders(w, limit, limit-(requestCount+1), requestCount+1, currentWindow+int64(window)) return } @@ -87,6 +89,9 @@ func RateLimiter(client *db.DiceDB, next http.Handler, limit, window int) http.H } } + // Add rate limit headers + addRateLimitHeaders(w, limit, limit-(requestCount+1), requestCount+1, currentWindow+int64(window)) + // Log the successful request increment slog.Info("Request processed", "count", requestCount+1) @@ -94,3 +99,10 @@ func RateLimiter(client *db.DiceDB, next http.Handler, limit, window int) http.H next.ServeHTTP(w, r) }) } + +func addRateLimitHeaders(w http.ResponseWriter, limit, remaining, used int, resetTime int64) { + w.Header().Set("x-ratelimit-limit", strconv.Itoa(limit)) + w.Header().Set("x-ratelimit-remaining", strconv.Itoa(remaining)) + w.Header().Set("x-ratelimit-used", strconv.Itoa(used)) + w.Header().Set("x-rateLimit-reset", strconv.FormatInt(resetTime, 10)) +} \ No newline at end of file