plumber/internal/web/throttle.go

package web

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

const defaultThrottleMaxKeys = 10_000

// throttle is a sliding-window rate limiter with expired-key eviction and a cap.
type throttle struct {
	mu      sync.Mutex
	hits    map[string][]time.Time
	limit   int
	window  time.Duration
	maxKeys int
}

func newThrottle(limit int, window time.Duration, maxKeys int) *throttle {
	if maxKeys <= 0 {
		maxKeys = defaultThrottleMaxKeys
	}
	return &throttle{
		hits:    map[string][]time.Time{},
		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()
	t.evictExpiredLocked(now)

	xs := pruneTimes(t.hits[key], now.Add(-t.window))
	if len(xs) >= t.limit {
		if len(xs) == 0 {
			delete(t.hits, key)
		} else {
			t.hits[key] = xs
		}
		return false
	}
	if _, exists := t.hits[key]; !exists && len(t.hits) >= t.maxKeys {
		t.evictExpiredLocked(now)
		if len(t.hits) >= t.maxKeys {
			return false
		}
	}
	t.hits[key] = append(xs, now)
	return true
}

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

func (t *throttle) evictExpiredLocked(now time.Time) {
	cutoff := now.Add(-t.window)
	for k, xs := range t.hits {
		xs = pruneTimes(xs, cutoff)
		if len(xs) == 0 {
			delete(t.hits, k)
		} else {
			t.hits[k] = xs
		}
	}
}

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 {
		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 {
	if s.cfg.TrustProxy {
		if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
			return strings.TrimSpace(strings.Split(xff, ",")[0])
		}
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

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
}