snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)

internal/middleware/auththrottle.go

@@ -0,0 +1,144 @@
+package middleware
+
+import (
+	"net"
+	"net/http"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+)
+
+// AuthThrottle limits repeated failing requests per client IP within a time
+// window. It is meant for authentication-style endpoints (admin password auth,
+// and any other endpoint where a caller repeatedly retries credentials) to slow
+// brute-force without ever throttling a caller that succeeds.
+//
+// Only failures are counted, and the per-IP limit is checked before the wrapped
+// handler runs, which gives two properties that matter for auth endpoints:
+//
+//   - A caller presenting valid credentials is never throttled, even while a
+//     different source is actively brute-forcing: success does not accrue against
+//     the limit, and each source IP is counted independently.
+//   - Once an IP is over the limit its requests are rejected before the handler
+//     runs, so the response cannot reveal whether the submitted credentials were
+//     valid.
+//
+// It is safe for concurrent use. Wrap a handler with Middleware:
+//
+//	throttle := middleware.NewAuthThrottle(10, time.Minute, nil)
+//	mux.Handle(path, throttle.Middleware(handler))
+type AuthThrottle struct {
+	max       int
+	window    time.Duration
+	isFailure func(status int) bool
+
+	mu        sync.Mutex
+	failures  map[string]int
+	windowEnd time.Time
+}
+
+// authThrottleMapCap bounds the number of tracked IPs so a flood of requests with
+// varying source addresses cannot grow the map without limit.
+const authThrottleMapCap = 10000
+
+// NewAuthThrottle creates an AuthThrottle allowing at most max failing requests
+// per client IP within window, after which further requests from that IP are
+// rejected with 429 until the window rolls over. isFailure decides, from the
+// status the wrapped handler wrote, whether a request counts as a failure; if
+// nil, any status >= 400 counts.
+func NewAuthThrottle(max int, window time.Duration, isFailure func(status int) bool) *AuthThrottle {
+	if isFailure == nil {
+		isFailure = func(status int) bool { return status >= 400 }
+	}
+	return &AuthThrottle{
+		max:       max,
+		window:    window,
+		isFailure: isFailure,
+		failures:  make(map[string]int),
+	}
+}
+
+// Middleware wraps h with per-IP failure throttling.
+func (t *AuthThrottle) Middleware(h http.Handler) http.Handler {
+	retryAfter := strconv.Itoa(int(t.window.Seconds()))
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ip := clientIP(r)
+		if !t.allow(ip) {
+			w.Header().Set("Retry-After", retryAfter)
+			http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
+			return
+		}
+		sw := &statusResponseWriter{ResponseWriter: w, status: http.StatusOK}
+		h.ServeHTTP(sw, r)
+		if t.isFailure(sw.Status()) {
+			t.recordFailure(ip)
+		}
+	})
+}
+
+// allow reports whether another request from ip may reach the handler. It rolls
+// the window over, discarding accumulated counts once the window elapses.
+func (t *AuthThrottle) allow(ip string) bool {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	if now := time.Now(); now.After(t.windowEnd) {
+		t.failures = make(map[string]int)
+		t.windowEnd = now.Add(t.window)
+	}
+	return t.failures[ip] < t.max
+}
+
+// recordFailure counts one failed request from ip. New IPs are not tracked once
+// the map is at capacity, keeping memory bounded under an address-varying flood.
+func (t *AuthThrottle) recordFailure(ip string) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	if _, tracked := t.failures[ip]; !tracked && len(t.failures) >= authThrottleMapCap {
+		return
+	}
+	t.failures[ip]++
+}
+
+// clientIP derives the client address used as the throttle key.
+//
+// Forwarded headers are only trusted when the immediate socket peer is a private
+// or loopback address, i.e. a local reverse proxy or load balancer (Centrifugo is
+// commonly deployed behind one). For a direct public client the headers are
+// client-controlled and ignored, so a peer cannot spoof them to evade throttling
+// or lock out another IP. Header values are validated and canonicalized as IPs
+// before use, so junk cannot inflate map keys or fragment the keyspace.
+func clientIP(r *http.Request) string {
+	host, _, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		host = r.RemoteAddr
+	}
+	peer := net.ParseIP(host)
+	if peer == nil || (!peer.IsLoopback() && !peer.IsPrivate()) {
+		return host
+	}
+	// Peer is a trusted local proxy: use the forwarded client address.
+	if ip := parseIP(r.Header.Get("X-Real-IP")); ip != "" {
+		return ip
+	}
+	if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
+		if i := strings.IndexByte(fwd, ','); i >= 0 {
+			fwd = fwd[:i]
+		}
+		if ip := parseIP(strings.TrimSpace(fwd)); ip != "" {
+			return ip
+		}
+	}
+	return host
+}
+
+// parseIP returns the canonical string form of s if it is a valid IP, else "".
+func parseIP(s string) string {
+	if s == "" {
+		return ""
+	}
+	if ip := net.ParseIP(s); ip != nil {
+		return ip.String()
+	}
+	return ""
+}