first commit

internal/ratelimit/http.go

@@ -0,0 +1,42 @@
+package ratelimit
+
+import (
+	"net"
+	"net/http"
+	"strings"
+)
+
+// HTTPMiddleware wraps next with per-client-IP rate limiting, returning 429
+// for requests over the limit. realIPHeader (e.g. "X-Forwarded-For"), if
+// non-empty, is trusted for the client IP instead of RemoteAddr — only set
+// this when the server is genuinely behind a reverse proxy that sets it;
+// trusting it otherwise lets any client spoof their rate-limit identity.
+func (l *Limiter) HTTPMiddleware(realIPHeader string, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ip := ClientIP(r, realIPHeader)
+		if !l.Allow(ip) {
+			w.Header().Set("Retry-After", "60")
+			http.Error(w, "rate limit exceeded, try again shortly", http.StatusTooManyRequests)
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}
+
+// ClientIP resolves the request's client IP the same way HTTPMiddleware
+// does, for callers that need it outside a rate-limit context (e.g. an IP
+// allowlist middleware). See HTTPMiddleware's doc comment for the
+// realIPHeader trust caveat.
+func ClientIP(r *http.Request, realIPHeader string) string {
+	if realIPHeader != "" {
+		if v := r.Header.Get(realIPHeader); v != "" {
+			parts := strings.Split(v, ",")
+			return strings.TrimSpace(parts[0])
+		}
+	}
+	host, _, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		return r.RemoteAddr
+	}
+	return host
+}

internal/ratelimit/ratelimit.go

@@ -0,0 +1,109 @@
+// Package ratelimit implements a per-key token-bucket rate limiter — no
+// third-party rate-limiting library. A token bucket (rather than a hard
+// fixed-window reset) is used deliberately: it smooths out bursts at
+// window boundaries that a naive "reset every 60s" counter would allow
+// (e.g. 20 requests at 0:59 plus another 20 at 1:01 both passing a
+// "20/min" limit reset at the minute boundary). Safe for concurrent use.
+package ratelimit
+
+import (
+	"sync"
+	"time"
+)
+
+type bucket struct {
+	tokens     float64
+	lastRefill time.Time
+}
+
+// Limiter enforces "at most ratePerMinute events per key, per minute" with
+// burst tolerance up to ratePerMinute tokens banked at once (i.e. a key
+// that's been idle can burst up to the full per-minute allowance instantly,
+// then is throttled to the steady-state rate — standard token-bucket
+// behavior, not a stricter "evenly spaced" enforcement).
+type Limiter struct {
+	ratePerMinute float64
+	mu            sync.Mutex
+	buckets       map[string]*bucket
+
+	stopCleanup chan struct{}
+}
+
+// New creates a limiter allowing ratePerMinute events per key. Pass 0 to
+// disable limiting entirely (Allow always returns true) — this is how a
+// zero/unset config value opts a listener out of rate limiting rather than
+// silently blocking everything.
+func New(ratePerMinute int) *Limiter {
+	l := &Limiter{
+		ratePerMinute: float64(ratePerMinute),
+		buckets:       make(map[string]*bucket),
+		stopCleanup:   make(chan struct{}),
+	}
+	if ratePerMinute > 0 {
+		go l.cleanupLoop()
+	}
+	return l
+}
+
+// Allow reports whether an event for key is permitted right now, consuming
+// one token if so. Safe to call from many goroutines concurrently.
+func (l *Limiter) Allow(key string) bool {
+	if l.ratePerMinute <= 0 {
+		return true
+	}
+
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	now := time.Now()
+	b, ok := l.buckets[key]
+	if !ok {
+		b = &bucket{tokens: l.ratePerMinute - 1, lastRefill: now}
+		l.buckets[key] = b
+		return true
+	}
+
+	elapsed := now.Sub(b.lastRefill).Seconds()
+	refill := elapsed * (l.ratePerMinute / 60.0)
+	b.tokens += refill
+	if b.tokens > l.ratePerMinute {
+		b.tokens = l.ratePerMinute
+	}
+	b.lastRefill = now
+
+	if b.tokens < 1 {
+		return false
+	}
+	b.tokens--
+	return true
+}
+
+// cleanupLoop periodically evicts buckets idle long enough to have fully
+// refilled, so a limiter tracking many distinct one-off IPs doesn't grow
+// unboundedly over a long-running server's lifetime.
+func (l *Limiter) cleanupLoop() {
+	ticker := time.NewTicker(10 * time.Minute)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-l.stopCleanup:
+			return
+		case <-ticker.C:
+			l.mu.Lock()
+			now := time.Now()
+			for key, b := range l.buckets {
+				if now.Sub(b.lastRefill) > 30*time.Minute {
+					delete(l.buckets, key)
+				}
+			}
+			l.mu.Unlock()
+		}
+	}
+}
+
+// Stop releases the background cleanup goroutine.
+func (l *Limiter) Stop() {
+	if l.ratePerMinute > 0 {
+		close(l.stopCleanup)
+	}
+}