91 lines
2.1 KiB
Go
91 lines
2.1 KiB
Go
package middleware
|
|
|
|
import (
|
|
"RenewCMS/api/pkg/apperr"
|
|
"net/http"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type successSlots struct {
|
|
mu sync.Mutex
|
|
window time.Duration
|
|
now func() time.Time
|
|
taken map[string]time.Time
|
|
}
|
|
|
|
func NewUserSuccessLimiter(window time.Duration) func(http.Handler) http.Handler {
|
|
return NewUserSuccessLimiterAt(window, time.Now)
|
|
}
|
|
|
|
func NewUserSuccessLimiterAt(window time.Duration, now func() time.Time) func(http.Handler) http.Handler {
|
|
slots := &successSlots{window: window, now: now, taken: make(map[string]time.Time)}
|
|
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
key := userRateLimitKey(r)
|
|
claimed, retryAfter, granted := slots.claim(key)
|
|
if !granted {
|
|
w.Header().Set("Retry-After", strconv.Itoa(retryAfter))
|
|
http.Error(w, apperr.TooManyRequestsErrorMessage, http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
|
|
recorder := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
|
|
succeeded := false
|
|
defer func() {
|
|
if !succeeded {
|
|
slots.release(key, claimed)
|
|
}
|
|
}()
|
|
next.ServeHTTP(recorder, r)
|
|
succeeded = recorder.status < http.StatusBadRequest
|
|
})
|
|
}
|
|
}
|
|
|
|
func (s *successSlots) claim(key string) (time.Time, int, bool) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
now := s.now()
|
|
s.sweepExpired(now)
|
|
if claimedAt, held := s.taken[key]; held {
|
|
return time.Time{}, retryAfterSeconds(s.window - now.Sub(claimedAt)), false
|
|
}
|
|
s.taken[key] = now
|
|
return now, 0, true
|
|
}
|
|
|
|
func (s *successSlots) release(key string, claimedAt time.Time) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
if current, held := s.taken[key]; held && current.Equal(claimedAt) {
|
|
delete(s.taken, key)
|
|
}
|
|
}
|
|
|
|
func (s *successSlots) sweepExpired(now time.Time) {
|
|
for key, claimedAt := range s.taken {
|
|
if now.Sub(claimedAt) >= s.window {
|
|
delete(s.taken, key)
|
|
}
|
|
}
|
|
}
|
|
|
|
type statusRecorder struct {
|
|
http.ResponseWriter
|
|
status int
|
|
}
|
|
|
|
func (s *statusRecorder) WriteHeader(status int) {
|
|
s.status = status
|
|
s.ResponseWriter.WriteHeader(status)
|
|
}
|
|
|
|
func (s *statusRecorder) Unwrap() http.ResponseWriter {
|
|
return s.ResponseWriter
|
|
}
|