pub

pkg/middleware/ratelimit.go

@@ -0,0 +1,92 @@
+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)
+		})
+	}
+}