81 lines
1.5 KiB
Go
81 lines
1.5 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net"
|
||
|
|
"net/http"
|
||
|
|
"strings"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type rateLimiter struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
visitors map[string]*visitor
|
||
|
|
rate int
|
||
|
|
window time.Duration
|
||
|
|
}
|
||
|
|
|
||
|
|
type visitor struct {
|
||
|
|
count int
|
||
|
|
reset time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
func newRateLimiter(rate int, window time.Duration) *rateLimiter {
|
||
|
|
rl := &rateLimiter{
|
||
|
|
visitors: make(map[string]*visitor),
|
||
|
|
rate: rate,
|
||
|
|
window: window,
|
||
|
|
}
|
||
|
|
go rl.cleanup()
|
||
|
|
return rl
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rl *rateLimiter) allow(key string) bool {
|
||
|
|
rl.mu.Lock()
|
||
|
|
defer rl.mu.Unlock()
|
||
|
|
now := time.Now()
|
||
|
|
v, ok := rl.visitors[key]
|
||
|
|
if !ok || now.After(v.reset) {
|
||
|
|
rl.visitors[key] = &visitor{count: 1, reset: now.Add(rl.window)}
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
if v.count >= rl.rate {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
v.count++
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rl *rateLimiter) cleanup() {
|
||
|
|
for {
|
||
|
|
time.Sleep(5 * time.Minute)
|
||
|
|
rl.mu.Lock()
|
||
|
|
now := time.Now()
|
||
|
|
for k, v := range rl.visitors {
|
||
|
|
if now.After(v.reset) {
|
||
|
|
delete(rl.visitors, k)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
rl.mu.Unlock()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// clientIP extracts the client address from X-Forwarded-For (set by nginx),
|
||
|
|
// falling back to X-Real-IP and finally RemoteAddr.
|
||
|
|
func clientIP(r *http.Request) string {
|
||
|
|
if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
|
||
|
|
if comma := strings.Index(xff, ","); comma > 0 {
|
||
|
|
return strings.TrimSpace(xff[:comma])
|
||
|
|
}
|
||
|
|
return strings.TrimSpace(xff)
|
||
|
|
}
|
||
|
|
if rip := r.Header.Get("X-Real-IP"); rip != "" {
|
||
|
|
return rip
|
||
|
|
}
|
||
|
|
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||
|
|
if err != nil {
|
||
|
|
return r.RemoteAddr
|
||
|
|
}
|
||
|
|
return host
|
||
|
|
}
|