init

back/internal/api/middleware/ratelimit.go

@@ -0,0 +1,72 @@
+package middleware
+
+import (
+	"net"
+	"net/http"
+	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+type ipLimiter struct {
+	limiter  *rate.Limiter
+	lastSeen time.Time
+}
+
+type Limiter struct {
+	mu       sync.Mutex
+	visitors map[string]*ipLimiter
+	r        rate.Limit
+	burst    int
+}
+
+func New(r rate.Limit, burst int) *Limiter {
+	l := &Limiter{
+		visitors: make(map[string]*ipLimiter),
+		r:        r,
+		burst:    burst,
+	}
+	go l.cleanupLoop()
+	return l
+}
+
+func (l *Limiter) getLimiter(ip string) *rate.Limiter {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	v, exists := l.visitors[ip]
+	if !exists {
+		v = &ipLimiter{limiter: rate.NewLimiter(l.r, l.burst)}
+		l.visitors[ip] = v
+	}
+	v.lastSeen = time.Now()
+	return v.limiter
+}
+
+func (l *Limiter) cleanupLoop() {
+	ticker := time.NewTicker(10 * time.Minute)
+	defer ticker.Stop()
+	for range ticker.C {
+		l.mu.Lock()
+		for ip, v := range l.visitors {
+			if time.Since(v.lastSeen) > 10*time.Minute {
+				delete(l.visitors, ip)
+			}
+		}
+		l.mu.Unlock()
+	}
+}
+
+func (l *Limiter) Handler(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ip, _, err := net.SplitHostPort(r.RemoteAddr)
+		if err != nil {
+			ip = r.RemoteAddr
+		}
+		if !l.getLimiter(ip).Allow() {
+			http.Error(w, `{"error":"rate limit exceeded, please slow down"}`, http.StatusTooManyRequests)
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}