unibus/pkg/membership/ratelimit.go

package membership

import (
	"fmt"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"golang.org/x/time/rate"
)

// ipRateLimiter is a per-source-IP token-bucket rate limiter for the control
// plane. It exists to blunt pre-auth flooding: an unauthenticated peer that
// hammers the HTTP API (signature verification is not free, and io is bounded
// but still real) is throttled before it can amplify load. Like the nonceCache,
// this is transport glue specific to unibus, not a registry primitive — the
// report 0003 made the same call for the nonce cache (it would only drag a NATS
// dependency into the multi-domain registry go.mod for one helper).
//
// Each distinct IP gets its own golang.org/x/time/rate.Limiter (a standard
// token bucket already in the module graph, so no new dependency). Idle buckets
// are reaped so the map cannot grow without bound under a churn of source IPs.
type ipRateLimiter struct {
	mu      sync.Mutex
	buckets map[string]*ipBucket
	r       rate.Limit
	burst   int
	ttl     time.Duration
}

type ipBucket struct {
	lim  *rate.Limiter
	seen time.Time
}

// newIPRateLimiter builds a limiter granting r tokens/second with the given
// burst per IP. ttl bounds how long an idle bucket is retained before being
// reaped. r<=0 disables limiting (Allow always true) so dev/loopback stacks are
// unaffected.
func newIPRateLimiter(r rate.Limit, burst int, ttl time.Duration) *ipRateLimiter {
	return &ipRateLimiter{
		buckets: make(map[string]*ipBucket),
		r:       r,
		burst:   burst,
		ttl:     ttl,
	}
}

// allow reports whether a request from ip may proceed now, consuming one token
// on success. A disabled limiter (r<=0) always allows. Reaping of stale buckets
// is amortized: it runs only when the map has grown past a small threshold, so
// the common path is a single map lookup under the mutex.
func (l *ipRateLimiter) allow(ip string, now time.Time) bool {
	if l == nil || l.r <= 0 {
		return true
	}
	l.mu.Lock()
	defer l.mu.Unlock()

	if len(l.buckets) > 1024 {
		l.reapLocked(now)
	}
	b, ok := l.buckets[ip]
	if !ok {
		b = &ipBucket{lim: rate.NewLimiter(l.r, l.burst)}
		l.buckets[ip] = b
	}
	b.seen = now
	return b.lim.AllowN(now, 1)
}

// reapLocked drops buckets idle for longer than ttl. The caller holds l.mu.
func (l *ipRateLimiter) reapLocked(now time.Time) {
	for ip, b := range l.buckets {
		if now.Sub(b.seen) > l.ttl {
			delete(l.buckets, ip)
		}
	}
}

// clientIP extracts the rate-limit key for a request: the source IP, with the
// port stripped. By default it trusts the transport's RemoteAddr ONLY (no
// X-Forwarded-For parsing): honoring an attacker-supplied header would let a
// single IP fan its quota across forged identities. When the operator runs the
// control plane behind a reverse proxy they control (the same-origin Caddy
// deployment), SetTrustedProxies names that proxy's address(es); only then, and
// only when the immediate peer is one of them, is the forwarded client IP
// believed. This keeps the per-IP rate limit meaningful behind the proxy, where
// every request would otherwise share the proxy's single IP. If parsing fails the
// whole RemoteAddr is used as the key (still a stable per-connection bucket).
func (s *Server) clientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	if !s.trustedProxies.has(host) {
		return host
	}
	if fwd := forwardedClientIP(r, s.trustedProxies); fwd != "" {
		return fwd
	}
	return host
}

// forwardedClientIP returns the real client IP a trusted proxy reported, or "" if
// none is present. X-Forwarded-For is read RIGHT-TO-LEFT: the rightmost entry is
// the one our immediate (trusted) proxy appended and therefore cannot be spoofed
// by the client, which can only prepend entries to the left. Trusted-proxy hops
// are skipped so a chain of proxies we own resolves to the first address none of
// them owns — the actual external client. X-Real-IP is a single-value fallback for
// proxies that set it instead. A non-trusted immediate peer never reaches here, so
// a direct attacker's forged header is ignored entirely.
func forwardedClientIP(r *http.Request, trusted trustedProxyMatcher) string {
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		parts := strings.Split(xff, ",")
		for i := len(parts) - 1; i >= 0; i-- {
			ip := strings.TrimSpace(parts[i])
			if ip == "" || trusted.has(ip) {
				continue
			}
			if net.ParseIP(ip) != nil {
				return ip
			}
		}
	}
	if xrip := strings.TrimSpace(r.Header.Get("X-Real-IP")); xrip != "" {
		if net.ParseIP(xrip) != nil {
			return xrip
		}
	}
	return ""
}

// trustedProxyMatcher is the set of reverse-proxy addresses whose forwarding
// headers may be honored. The zero value (nil) matches nothing, so the default
// behavior is RemoteAddr-only.
type trustedProxyMatcher []*net.IPNet

// SetTrustedProxies configures the proxies whose X-Forwarded-For / X-Real-IP this
// server trusts for the per-IP rate limit. Each entry is an IP (treated as a /32
// or /128) or a CIDR. It returns an error on the first unparseable entry and
// leaves the previous configuration unchanged. Passing no entries clears the set.
func (s *Server) SetTrustedProxies(entries []string) error {
	m, err := parseTrustedProxies(entries)
	if err != nil {
		return err
	}
	s.trustedProxies = m
	return nil
}

// parseTrustedProxies turns a list of IPs/CIDRs into a matcher. A bare IP becomes
// a host route (/32 for IPv4, /128 for IPv6); blanks are skipped.
func parseTrustedProxies(entries []string) (trustedProxyMatcher, error) {
	var m trustedProxyMatcher
	for _, e := range entries {
		e = strings.TrimSpace(e)
		if e == "" {
			continue
		}
		if _, ipnet, err := net.ParseCIDR(e); err == nil {
			m = append(m, ipnet)
			continue
		}
		ip := net.ParseIP(e)
		if ip == nil {
			return nil, fmt.Errorf("trusted proxy %q is not an IP or CIDR", e)
		}
		bits := 32
		if ip.To4() == nil {
			bits = 128
		}
		m = append(m, &net.IPNet{IP: ip, Mask: net.CIDRMask(bits, bits)})
	}
	return m, nil
}

// has reports whether host (an IP string with no port) falls inside any trusted
// range. A nil matcher and an unparseable host both report false.
func (m trustedProxyMatcher) has(host string) bool {
	if len(m) == 0 {
		return false
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return false
	}
	for _, n := range m {
		if n.Contains(ip) {
			return true
		}
	}
	return false
}