Prototype commit

server/ratelimit.go

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