67 lines
1.3 KiB
Go
67 lines
1.3 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
type IPRateLimiter struct {
|
|
mu sync.RWMutex
|
|
limiters map[string]*rate.Limiter
|
|
r rate.Limit
|
|
b int
|
|
}
|
|
|
|
func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
|
|
return &IPRateLimiter{
|
|
limiters: make(map[string]*rate.Limiter),
|
|
r: r,
|
|
b: b,
|
|
}
|
|
}
|
|
|
|
func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
|
|
i.mu.Lock()
|
|
defer i.mu.Unlock()
|
|
if lim, ok := i.limiters[ip]; ok {
|
|
return lim
|
|
}
|
|
lim := rate.NewLimiter(i.r, i.b)
|
|
i.limiters[ip] = lim
|
|
return lim
|
|
}
|
|
|
|
func (i *IPRateLimiter) Cleanup() {
|
|
ticker := time.NewTicker(time.Minute)
|
|
go func() {
|
|
for range ticker.C {
|
|
i.mu.Lock()
|
|
for k := range i.limiters {
|
|
delete(i.limiters, k)
|
|
}
|
|
i.mu.Unlock()
|
|
}
|
|
}()
|
|
}
|
|
|
|
func RateLimit(l *IPRateLimiter) func(http.Handler) http.Handler {
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
ip := r.RemoteAddr
|
|
if xf := r.Header.Get("X-Forwarded-For"); xf != "" {
|
|
// simples, suficiente pra começar
|
|
ip = xf
|
|
}
|
|
if !l.GetLimiter(ip).Allow() {
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
w.Write([]byte(`{"error":"muitas tentativas, tente novamente mais tarde"}`))
|
|
return
|
|
}
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|