rl544
93 lines
1.7 KiB
Go
93 lines
1.7 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// visitor defines the rate limiter for a single IP
|
|
type visitor struct {
|
|
limiter *rate.Limiter
|
|
lastSeen time.Time
|
|
}
|
|
|
|
// IPRateLimiter holds a map of rate limiters per IP
|
|
type IPRateLimiter struct {
|
|
mu sync.Mutex
|
|
visitors map[string]*visitor
|
|
rate rate.Limit
|
|
burst int
|
|
}
|
|
|
|
func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
|
|
i := &IPRateLimiter{
|
|
visitors: make(map[string]*visitor),
|
|
rate: r,
|
|
burst: b,
|
|
}
|
|
|
|
// Background routine to clean up old visitors
|
|
go func() {
|
|
for {
|
|
time.Sleep(time.Minute)
|
|
i.cleanupVisitors()
|
|
}
|
|
}()
|
|
|
|
return i
|
|
}
|
|
|
|
func (i *IPRateLimiter) getLimiter(ip string) *rate.Limiter {
|
|
i.mu.Lock()
|
|
defer i.mu.Unlock()
|
|
|
|
v, exists := i.visitors[ip]
|
|
if !exists {
|
|
limiter := rate.NewLimiter(i.rate, i.burst)
|
|
i.visitors[ip] = &visitor{
|
|
limiter: limiter,
|
|
lastSeen: time.Now(),
|
|
}
|
|
return limiter
|
|
}
|
|
|
|
v.lastSeen = time.Now()
|
|
return v.limiter
|
|
}
|
|
|
|
func (i *IPRateLimiter) cleanupVisitors() {
|
|
i.mu.Lock()
|
|
defer i.mu.Unlock()
|
|
for ip, v := range i.visitors {
|
|
if time.Since(v.lastSeen) > 3*time.Minute {
|
|
delete(i.visitors, ip)
|
|
}
|
|
}
|
|
}
|
|
|
|
// RateLimit is a middleware that limits requests per IP
|
|
func RateLimit(requestsPerSecond float64, burst int) Middleware {
|
|
limiter := NewIPRateLimiter(rate.Limit(requestsPerSecond), burst)
|
|
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
// Extract IP
|
|
ip, _, err := net.SplitHostPort(r.RemoteAddr)
|
|
if err != nil {
|
|
ip = r.RemoteAddr
|
|
}
|
|
|
|
if !limiter.getLimiter(ip).Allow() {
|
|
http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|