61 lines
1.4 KiB
Go
61 lines
1.4 KiB
Go
package middleware
|
|
|
|
import (
|
|
"RenewCMS/api/pkg/apperr"
|
|
"net/http"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type rateLimitWindow struct {
|
|
start time.Time
|
|
count int
|
|
}
|
|
|
|
func NewRateLimiter(maxAttempts int, window time.Duration) func(http.Handler) http.Handler {
|
|
var mu sync.Mutex
|
|
windows := make(map[string]*rateLimitWindow)
|
|
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
retryAfter, allowed := recordAttempt(&mu, windows, r.RemoteAddr, maxAttempts, window)
|
|
if !allowed {
|
|
w.Header().Set("Retry-After", strconv.Itoa(retryAfter))
|
|
http.Error(w, apperr.TooManyRequestsErrorMessage, http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|
|
|
|
func recordAttempt(mu *sync.Mutex, windows map[string]*rateLimitWindow, key string, maxAttempts int, window time.Duration) (int, bool) {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
|
|
now := time.Now()
|
|
current, exists := windows[key]
|
|
if !exists || now.Sub(current.start) >= window {
|
|
current = &rateLimitWindow{start: now}
|
|
windows[key] = current
|
|
}
|
|
|
|
current.count++
|
|
if current.count > maxAttempts {
|
|
return retryAfterSeconds(window - now.Sub(current.start)), false
|
|
}
|
|
return 0, true
|
|
}
|
|
|
|
func retryAfterSeconds(remaining time.Duration) int {
|
|
seconds := int(remaining / time.Second)
|
|
if remaining%time.Second > 0 {
|
|
seconds++
|
|
}
|
|
if seconds < 0 {
|
|
return 0
|
|
}
|
|
return seconds
|
|
}
|