100 lines
1.9 KiB
Go
100 lines
1.9 KiB
Go
package mw
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// RateLimiter implements a simple in-memory rate limiter
|
|
type RateLimiter struct {
|
|
requests map[string][]time.Time
|
|
mu sync.RWMutex
|
|
limit int
|
|
window time.Duration
|
|
}
|
|
|
|
func NewRateLimiter(limit int, window time.Duration) *RateLimiter {
|
|
rl := &RateLimiter{
|
|
requests: make(map[string][]time.Time),
|
|
limit: limit,
|
|
window: window,
|
|
}
|
|
// Cleanup old entries periodically
|
|
go func() {
|
|
for {
|
|
time.Sleep(window)
|
|
rl.cleanup()
|
|
}
|
|
}()
|
|
return rl
|
|
}
|
|
|
|
func (rl *RateLimiter) cleanup() {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
now := time.Now()
|
|
for key, times := range rl.requests {
|
|
var valid []time.Time
|
|
for _, t := range times {
|
|
if now.Sub(t) < rl.window {
|
|
valid = append(valid, t)
|
|
}
|
|
}
|
|
if len(valid) == 0 {
|
|
delete(rl.requests, key)
|
|
} else {
|
|
rl.requests[key] = valid
|
|
}
|
|
}
|
|
}
|
|
|
|
func (rl *RateLimiter) Allow(key string) bool {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
now := time.Now()
|
|
windowStart := now.Add(-rl.window)
|
|
|
|
var valid []time.Time
|
|
for _, t := range rl.requests[key] {
|
|
if t.After(windowStart) {
|
|
valid = append(valid, t)
|
|
}
|
|
}
|
|
|
|
if len(valid) >= rl.limit {
|
|
rl.requests[key] = valid
|
|
return false
|
|
}
|
|
|
|
rl.requests[key] = append(valid, now)
|
|
return true
|
|
}
|
|
|
|
// RateLimit middleware - limits requests per IP
|
|
func RateLimit(limit int, window time.Duration) func(http.Handler) http.Handler {
|
|
limiter := NewRateLimiter(limit, window)
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
// Get client IP
|
|
ip := r.Header.Get("CF-Connecting-IP")
|
|
if ip == "" {
|
|
ip = r.Header.Get("X-Real-IP")
|
|
}
|
|
if ip == "" {
|
|
ip = r.Header.Get("X-Forwarded-For")
|
|
}
|
|
if ip == "" {
|
|
ip = r.RemoteAddr
|
|
}
|
|
|
|
if !limiter.Allow(ip) {
|
|
http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|