oikos/internal/httpapi/ratelimit.go

package httpapi

import (
	"context"
	"log/slog"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

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

// rateLimiter is a per-client (IP) token-bucket limiter registry. Each unique
// client gets its own *rate.Limiter; idle entries are swept periodically so a
// flood of distinct IPs can't grow the map unbounded. A rate of zero (rps==0)
// disables limiting entirely — the returned middleware is a no-op.
//
// Client identity is the source IP. Behind Caddy the real client is in
// X-Forwarded-For: Caddy appends the immediate client as the LAST hop, while
// earlier hops are client-supplied and spoofable. clientIP therefore takes the
// rightmost XFF entry (the proxy's contribution) rather than the first.
type rateLimiter struct {
	mu       sync.Mutex
	limiters map[string]*entry
	rps      rate.Limit
	burst    int
}

type entry struct {
	limiter  *rate.Limiter
	lastSeen time.Time
}

// newRateLimiter builds the registry and starts the idle-entry sweeper tied to
// ctx, so the ticker is stopped when the server shuts down.
func newRateLimiter(ctx context.Context, rps, burst int) *rateLimiter {
	rl := &rateLimiter{
		limiters: make(map[string]*entry),
		rps:      rate.Limit(rps),
		burst:    burst,
	}
	if rps > 0 {
		go rl.sweep(ctx)
	}
	return rl
}

// sweep drops entries untouched since the last sweep so the registry doesn't
// grow without bound under a rotating-IP attack or long-lived process. Exits
// (and stops its ticker) when ctx is cancelled.
func (rl *rateLimiter) sweep(ctx context.Context) {
	ticker := time.NewTicker(5 * time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			rl.mu.Lock()
			for ip, e := range rl.limiters {
				if time.Since(e.lastSeen) > 10*time.Minute {
					delete(rl.limiters, ip)
				}
			}
			rl.mu.Unlock()
		}
	}
}

func (rl *rateLimiter) get(ip string) *rate.Limiter {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	if e, ok := rl.limiters[ip]; ok {
		e.lastSeen = time.Now()
		return e.limiter
	}
	l := rate.NewLimiter(rl.rps, rl.burst)
	rl.limiters[ip] = &entry{limiter: l, lastSeen: time.Now()}
	return l
}

// middleware returns a chi-style middleware that enforces the per-IP limit.
// Call with rps==0 to get a pass-through no-op.
func (rl *rateLimiter) middleware(next http.Handler) http.Handler {
	if rl.rps <= 0 {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Exempt infra liveness probes so Caddy/compose healthchecks can't be
		// throttled into marking the service unhealthy.
		if r.URL.Path == "/healthz" {
			next.ServeHTTP(w, r)
			return
		}
		if !rl.get(clientIP(r)).Allow() {
			w.Header().Set("Retry-After", "1")
			writeProblem(w, r, http.StatusTooManyRequests, "rate limit exceeded", "")
			return
		}
		next.ServeHTTP(w, r)
	})
}

// clientIP extracts the originating client address. It takes the rightmost
// X-Forwarded-For hop — the one the reverse proxy (Caddy) appends for the
// immediate client — because earlier hops are attacker-controlled and could
// be spoofed to dodge the limit or exhaust another client's bucket. Falls
// back to r.RemoteAddr when no XFF header is present.
//
// Known limitation: this is only trustworthy when the request actually
// traverses Caddy. A client connecting directly to the published :8090 (not
// behind the proxy) can set a single-hop XFF and have it trusted. That only
// evades rate limiting (auth is still required), and rate limiting is off by
// default, so the blast radius is narrow. Fully closing it requires either
// Caddy trusted_proxies (so it overwrites XFF / sets a non-spoofable
// X-Real-Ip) or keying the limiter on the auth token instead of IP.
func clientIP(r *http.Request) string {
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		if idx := strings.LastIndex(xff, ","); idx >= 0 {
			xff = xff[idx+1:]
		}
		if ip := strings.TrimSpace(xff); ip != "" {
			return ip
		}
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

// newRateLimiterFromConfig builds the limiter from API config, logging the
// chosen policy once at startup. rps<=0 means "disabled" (returns a no-op
// middleware) so dev/single-user setups aren't throttled by default.
func newRateLimiterFromConfig(ctx context.Context, rps, burst int) *rateLimiter {
	if rps <= 0 {
		slog.Info("api rate limiting disabled (OIKOS_API_RATE_LIMIT unset)")
		return &rateLimiter{rps: 0}
	}
	if burst <= 0 {
		burst = rps * 2
	}
	slog.Info("api rate limiting enabled", "rps", rps, "burst", burst)
	return newRateLimiter(ctx, rps, burst)
}