plumber/internal/web/throttle.go

package web

import (
	"container/list"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const defaultThrottleMaxKeys = 10_000

// throttle is a sliding-window rate limiter with LRU eviction at capacity.
type throttle struct {
	mu      sync.Mutex
	hits    map[string]*throttleEntry
	lru     *list.List // front = most recently used
	limit   int
	window  time.Duration
	maxKeys int
	rejects atomic.Uint64
	lastLog time.Time
}

type throttleEntry struct {
	times []time.Time
	el    *list.Element
}

func newThrottle(limit int, window time.Duration, maxKeys int) *throttle {
	if maxKeys <= 0 {
		maxKeys = defaultThrottleMaxKeys
	}
	return &throttle{
		hits:    map[string]*throttleEntry{},
		lru:     list.New(),
		limit:   limit,
		window:  window,
		maxKeys: maxKeys,
	}
}

func (t *throttle) allow(key string) bool {
	if t == nil || key == "" {
		return true
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	now := time.Now()
	cutoff := now.Add(-t.window)

	ent, ok := t.hits[key]
	if ok {
		ent.times = pruneTimes(ent.times, cutoff)
		if len(ent.times) == 0 {
			t.removeLocked(key, ent)
			ok = false
		}
	}
	if ok {
		if len(ent.times) >= t.limit {
			t.touchLocked(key, ent)
			return false
		}
		ent.times = append(ent.times, now)
		t.touchLocked(key, ent)
		return true
	}

	// New key: make room via LRU if needed.
	for len(t.hits) >= t.maxKeys {
		oldest := t.lru.Back()
		if oldest == nil {
			break
		}
		oldKey := oldest.Value.(string)
		t.removeLocked(oldKey, t.hits[oldKey])
		n := t.rejects.Add(1)
		if time.Since(t.lastLog) > time.Minute {
			log.Printf("throttle: LRU evicted key at capacity=%d rejects=%d", t.maxKeys, n)
			t.lastLog = now
		}
	}
	ent = &throttleEntry{times: []time.Time{now}}
	ent.el = t.lru.PushFront(key)
	t.hits[key] = ent
	return true
}

func (t *throttle) touchLocked(key string, ent *throttleEntry) {
	if ent.el != nil {
		t.lru.MoveToFront(ent.el)
	}
}

func (t *throttle) removeLocked(key string, ent *throttleEntry) {
	if ent == nil {
		return
	}
	if ent.el != nil {
		t.lru.Remove(ent.el)
	}
	delete(t.hits, key)
}

func (t *throttle) lenKeys() int {
	t.mu.Lock()
	defer t.mu.Unlock()
	return len(t.hits)
}

func pruneTimes(xs []time.Time, cutoff time.Time) []time.Time {
	n := 0
	for _, ts := range xs {
		if ts.After(cutoff) {
			xs[n] = ts
			n++
		}
	}
	return xs[:n]
}

// failureTracker records auth failures for progressive delay (not a hard lockout).
type failureTracker struct {
	mu      sync.Mutex
	fails   map[string]failState
	window  time.Duration
	maxKeys int
}

type failState struct {
	count int
	last  time.Time
}

func newFailureTracker(window time.Duration, maxKeys int) *failureTracker {
	if maxKeys <= 0 {
		maxKeys = defaultThrottleMaxKeys
	}
	return &failureTracker{
		fails:   map[string]failState{},
		window:  window,
		maxKeys: maxKeys,
	}
}

func (f *failureTracker) delay(key string) time.Duration {
	if f == nil || key == "" {
		return 0
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	now := time.Now()
	f.evictExpiredLocked(now)
	st, ok := f.fails[key]
	if !ok {
		return 0
	}
	return progressiveDelay(st.count)
}

func (f *failureTracker) record(key string) {
	if f == nil || key == "" {
		return
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	now := time.Now()
	f.evictExpiredLocked(now)
	st := f.fails[key]
	if st.count == 0 && len(f.fails) >= f.maxKeys {
		// Drop an arbitrary expired-or-oldest entry.
		for k, v := range f.fails {
			if now.Sub(v.last) > f.window/2 {
				delete(f.fails, k)
				break
			}
		}
		if len(f.fails) >= f.maxKeys {
			return
		}
	}
	st.count++
	st.last = now
	f.fails[key] = st
}

func (f *failureTracker) clear(key string) {
	if f == nil || key == "" {
		return
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	delete(f.fails, key)
}

func (f *failureTracker) lenKeys() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return len(f.fails)
}

func (f *failureTracker) evictExpiredLocked(now time.Time) {
	cutoff := now.Add(-f.window)
	for k, st := range f.fails {
		if st.last.Before(cutoff) {
			delete(f.fails, k)
		}
	}
}

func progressiveDelay(failCount int) time.Duration {
	switch {
	case failCount <= 1:
		return 0
	case failCount == 2:
		return 200 * time.Millisecond
	case failCount == 3:
		return 500 * time.Millisecond
	case failCount == 4:
		return time.Second
	default:
		return 2 * time.Second
	}
}

func (s *Server) clientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	peer := net.ParseIP(host)
	if peer == nil || !ipInNets(peer, s.cfg.TrustedProxies) {
		return host
	}
	xff := r.Header.Get("X-Forwarded-For")
	if xff == "" {
		return host
	}
	parts := strings.Split(xff, ",")
	// Walk right-to-left; skip trusted hops; first untrusted is the client.
	for i := len(parts) - 1; i >= 0; i-- {
		p := net.ParseIP(strings.TrimSpace(parts[i]))
		if p == nil {
			continue
		}
		if !ipInNets(p, s.cfg.TrustedProxies) {
			return p.String()
		}
	}
	return host
}

func ipInNets(ip net.IP, nets []*net.IPNet) bool {
	for _, n := range nets {
		if n.Contains(ip) {
			return true
		}
	}
	return false
}

func parseCIDRs(raw string) []*net.IPNet {
	var out []*net.IPNet
	for _, part := range strings.Split(raw, ",") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		_, n, err := net.ParseCIDR(part)
		if err != nil {
			log.Printf("trusted proxy CIDR ignored %q: %v", part, err)
			continue
		}
		out = append(out, n)
	}
	return out
}

func authTooMany(w http.ResponseWriter) {
	http.Error(w, "too many attempts; try again later", http.StatusTooManyRequests)
}

func loginFailKey(ip, usernameKey string) string {
	return ip + "\x00" + usernameKey
}

func (s *Server) allowLoginAttempt(w http.ResponseWriter, r *http.Request, usernameKey string) bool {
	ip := s.clientIP(r)
	if !s.loginIP.allow(ip) {
		authTooMany(w)
		return false
	}
	if d := s.loginFail.delay(loginFailKey(ip, usernameKey)); d > 0 {
		time.Sleep(d)
	}
	return true
}

func (s *Server) allowRegisterAttempt(w http.ResponseWriter, r *http.Request) bool {
	if !s.registerIP.allow(s.clientIP(r)) {
		authTooMany(w)
		return false
	}
	return true
}