97 lines
2.3 KiB
Go
97 lines
2.3 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
|
|
}
|
|
|
|
type rateLimitCounters struct {
|
|
mu sync.Mutex
|
|
window time.Duration
|
|
windows map[string]*rateLimitWindow
|
|
lastSweep time.Time
|
|
}
|
|
|
|
func NewRateLimiter(maxAttempts int, window time.Duration) func(http.Handler) http.Handler {
|
|
return newKeyedRateLimiter(maxAttempts, window, connectionHostKey)
|
|
}
|
|
|
|
func connectionHostKey(r *http.Request) string {
|
|
return hostOf(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 {
|
|
counters := &rateLimitCounters{window: window, windows: make(map[string]*rateLimitWindow)}
|
|
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
retryAfter, allowed := counters.recordAttempt(keyOf(r), maxAttempts, time.Now())
|
|
if !allowed {
|
|
w.Header().Set("Retry-After", strconv.Itoa(retryAfter))
|
|
http.Error(w, apperr.TooManyRequestsErrorMessage, http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|
|
|
|
func (c *rateLimitCounters) recordAttempt(key string, maxAttempts int, now time.Time) (int, bool) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
|
|
c.sweepExpired(now)
|
|
current, exists := c.windows[key]
|
|
if !exists || now.Sub(current.start) >= c.window {
|
|
current = &rateLimitWindow{start: now}
|
|
c.windows[key] = current
|
|
}
|
|
|
|
current.count++
|
|
if current.count > maxAttempts {
|
|
return retryAfterSeconds(c.window - now.Sub(current.start)), false
|
|
}
|
|
return 0, true
|
|
}
|
|
|
|
func (c *rateLimitCounters) sweepExpired(now time.Time) {
|
|
if now.Sub(c.lastSweep) < c.window {
|
|
return
|
|
}
|
|
c.lastSweep = now
|
|
for key, current := range c.windows {
|
|
if now.Sub(current.start) >= c.window {
|
|
delete(c.windows, key)
|
|
}
|
|
}
|
|
}
|
|
|
|
func retryAfterSeconds(remaining time.Duration) int {
|
|
seconds := int(remaining / time.Second)
|
|
if remaining%time.Second > 0 {
|
|
seconds++
|
|
}
|
|
if seconds < 0 {
|
|
return 0
|
|
}
|
|
return seconds
|
|
}
|