init

internal/security/ratelimit.go

@@ -0,0 +1,80 @@
+package security
+
+import (
+	"net/http"
+	"sync"
+	"time"
+)
+
+type bucket struct {
+	tokens   float64
+	last     time.Time
+	blocked  time.Time
+}
+
+type RateLimiter struct {
+	mu        sync.Mutex
+	capacity  float64
+	refillPer float64 // tokens/sec
+	ttl       time.Duration
+	buckets   map[string]*bucket
+}
+
+func NewRateLimiter(capacity int, refillPerSec float64, ttl time.Duration) *RateLimiter {
+	return &RateLimiter{
+		capacity:  float64(capacity),
+		refillPer: refillPerSec,
+		ttl:       ttl,
+		buckets:   map[string]*bucket{},
+	}
+}
+
+func (rl *RateLimiter) Allow(key string) bool {
+	now := time.Now()
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+	b := rl.buckets[key]
+	if b == nil {
+		b = &bucket{tokens: rl.capacity, last: now}
+		rl.buckets[key] = b
+	}
+	// cleanup occasionally
+	if len(rl.buckets) > 10000 {
+		for k, v := range rl.buckets {
+			if now.Sub(v.last) > rl.ttl {
+				delete(rl.buckets, k)
+			}
+		}
+	}
+
+	elapsed := now.Sub(b.last).Seconds()
+	if elapsed > 0 {
+		b.tokens += elapsed * rl.refillPer
+		if b.tokens > rl.capacity {
+			b.tokens = rl.capacity
+		}
+		b.last = now
+	}
+	if b.tokens >= 1 {
+		b.tokens -= 1
+		return true
+	}
+	return false
+}
+
+func (rl *RateLimiter) Middleware(keyFn func(r *http.Request) string) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			key := keyFn(r)
+			if key == "" {
+				key = "anon"
+			}
+			if !rl.Allow(key) {
+				w.Header().Set("Retry-After", "2")
+				http.Error(w, "rate limited", http.StatusTooManyRequests)
+				return
+			}
+			next.ServeHTTP(w, r)
+		})
+	}
+}