plumber/internal/web/throttle.go

package web

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

// throttle is a simple sliding-window rate limiter for auth endpoints.
type throttle struct {
	mu     sync.Mutex
	hits   map[string][]time.Time
	limit  int
	window time.Duration
}

func newThrottle(limit int, window time.Duration) *throttle {
	return &throttle{
		hits:   map[string][]time.Time{},
		limit:  limit,
		window: window,
	}
}

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)
	xs := t.hits[key]
	n := 0
	for _, ts := range xs {
		if ts.After(cutoff) {
			xs[n] = ts
			n++
		}
	}
	xs = xs[:n]
	if len(xs) >= t.limit {
		t.hits[key] = xs
		return false
	}
	t.hits[key] = append(xs, now)
	return true
}

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 (s *Server) allowLoginAttempt(w http.ResponseWriter, r *http.Request, usernameKey string) bool {
	if !s.loginIP.allow(s.clientIP(r)) || !s.loginUser.allow(usernameKey) {
		authTooMany(w)
		return false
	}
	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
}