102 lines
2.2 KiB
Go
102 lines
2.2 KiB
Go
package server
|
|
|
|
import (
|
|
"log"
|
|
"net"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type rateLimitState struct {
|
|
count int
|
|
resetAt time.Time
|
|
blockedUntil time.Time
|
|
}
|
|
|
|
type rateLimiter struct {
|
|
mu sync.Mutex
|
|
states map[string]*rateLimitState
|
|
maxRequests int
|
|
window time.Duration
|
|
blockFor time.Duration
|
|
}
|
|
|
|
func newRateLimiter(maxRequests int, window, blockFor time.Duration) *rateLimiter {
|
|
return &rateLimiter{
|
|
states: make(map[string]*rateLimitState),
|
|
maxRequests: maxRequests,
|
|
window: window,
|
|
blockFor: blockFor,
|
|
}
|
|
}
|
|
|
|
func (limiter *rateLimiter) Middleware(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if !limiter.allow(r) {
|
|
w.Header().Set("Retry-After", strconvSeconds(limiter.blockFor))
|
|
http.Error(w, "Too many requests", http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
|
|
func (limiter *rateLimiter) allow(r *http.Request) bool {
|
|
key := clientIP(r) + ":" + r.URL.Path
|
|
now := time.Now()
|
|
|
|
limiter.mu.Lock()
|
|
defer limiter.mu.Unlock()
|
|
|
|
state, ok := limiter.states[key]
|
|
if ok && now.Before(state.blockedUntil) {
|
|
return false
|
|
}
|
|
|
|
if !ok || now.After(state.resetAt) {
|
|
limiter.states[key] = &rateLimitState{
|
|
count: 1,
|
|
resetAt: now.Add(limiter.window),
|
|
}
|
|
return true
|
|
}
|
|
|
|
state.count++
|
|
if state.count > limiter.maxRequests {
|
|
state.blockedUntil = now.Add(limiter.blockFor)
|
|
state.resetAt = now.Add(limiter.window)
|
|
state.count = 0
|
|
log.Printf("Rate limit triggered for %s on %s", clientIP(r), r.URL.Path)
|
|
return false
|
|
}
|
|
|
|
return true
|
|
}
|
|
|
|
func clientIP(r *http.Request) string {
|
|
if forwardedFor := r.Header.Get("X-Forwarded-For"); forwardedFor != "" {
|
|
parts := strings.Split(forwardedFor, ",")
|
|
if ip := strings.TrimSpace(parts[0]); ip != "" {
|
|
return ip
|
|
}
|
|
}
|
|
if realIP := strings.TrimSpace(r.Header.Get("X-Real-IP")); realIP != "" {
|
|
return realIP
|
|
}
|
|
if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
|
|
return host
|
|
}
|
|
return r.RemoteAddr
|
|
}
|
|
|
|
func strconvSeconds(duration time.Duration) string {
|
|
seconds := int(duration.Seconds())
|
|
if seconds < 1 {
|
|
seconds = 1
|
|
}
|
|
return strconv.Itoa(seconds)
|
|
}
|