75 lines
1.9 KiB
Go
75 lines
1.9 KiB
Go
package middleware
|
|
|
|
import (
|
|
"RenewCMS/api/pkg/apperr"
|
|
"RenewCMS/api/pkg/claims"
|
|
"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 {
|
|
return newKeyedRateLimiter(maxAttempts, window, func(r *http.Request) string { return r.RemoteAddr })
|
|
}
|
|
|
|
func NewUserRateLimiter(maxAttempts int, window time.Duration) func(http.Handler) http.Handler {
|
|
return newKeyedRateLimiter(maxAttempts, window, userRateLimitKey)
|
|
}
|
|
|
|
func userRateLimitKey(r *http.Request) string {
|
|
userId, _ := claims.UserId(r)
|
|
return strconv.FormatUint(uint64(userId), 10)
|
|
}
|
|
|
|
func newKeyedRateLimiter(maxAttempts int, window time.Duration, keyOf func(*http.Request) string) 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, keyOf(r), 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
|
|
}
|