FILE / Unmei/Backend

v1/ratelimiter.go

Исходный файл и его история в репозитории.
FILE ae81fbcce941523578569fccd828c513e38e6201
Files
Backend/v1/ratelimiter.go
T
2021-05-02 16:04:45 +03:00

64 lines
1018 B
Go

package v1
import (
"net"
"net/http"
"sync"
"time"
"golang.org/x/time/rate"
)
type visitor struct {
limiter *rate.Limiter
lastSeen time.Time
}
var visitors = make(map[string]*visitor)
var mu sync.Mutex
func CleanupVisitors() {
for {
time.Sleep(time.Minute)
mu.Lock()
for ip, v := range visitors {
if time.Since(v.lastSeen) > time.Minute*3 {
delete(visitors, ip)
}
}
mu.Unlock()
}
}
func getVisitor(ip string) *rate.Limiter {
mu.Lock()
defer mu.Unlock()
v, exists := visitors[ip]
if !exists {
limiter := rate.NewLimiter(1, 20)
visitors[ip] = &visitor{ limiter, time.Now() }
return limiter
}
v.lastSeen = time.Now()
return v.limiter
}
func LimitMiddleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ip, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
return
}
limiter := getVisitor(ip)
if !limiter.Allow() {
TooManyRequests.Ferror(w)
return
}
next.ServeHTTP(w, r)
})
}