EvoBGP/internal/httpapi/ratelimit.go

package httpapi

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

const maxAuthRateLimitClients = 4096

type authRateHit struct {
	start time.Time
	count int
}

// authRateLimiter keeps a bounded fixed-window counter per client IP.
type authRateLimiter struct {
	mu         sync.Mutex
	hits       map[string]authRateHit
	limit      int
	window     time.Duration
	maxClients int
	nextSweep  time.Time
}

func newAuthRateLimiter() *authRateLimiter {
	limit := 60
	if n, err := strconv.Atoi(strings.TrimSpace(os.Getenv("EVOBGP_AUTH_RATE_LIMIT"))); err == nil && n > 0 && n <= 1000 {
		limit = n
	}
	return &authRateLimiter{
		hits:       make(map[string]authRateHit),
		limit:      limit,
		window:     time.Minute,
		maxClients: maxAuthRateLimitClients,
	}
}

func (l *authRateLimiter) allow(ip string) bool {
	if l == nil {
		return true
	}
	return l.allowAt(ip, time.Now())
}

func (l *authRateLimiter) allowAt(ip string, now time.Time) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	if !l.nextSweep.IsZero() && !now.Before(l.nextSweep) {
		l.nextSweep = time.Time{}
		for key, hit := range l.hits {
			expiry := hit.start.Add(l.window)
			if !now.Before(expiry) {
				delete(l.hits, key)
			} else if l.nextSweep.IsZero() || expiry.Before(l.nextSweep) {
				l.nextSweep = expiry
			}
		}
	}
	hit, exists := l.hits[ip]
	if !exists && len(l.hits) >= l.maxClients {
		// Never evict an active counter: a new identity must wait for expiry.
		return false
	}
	if exists && hit.count >= l.limit {
		return false
	}
	if !exists {
		hit.start = now
		expiry := now.Add(l.window)
		if l.nextSweep.IsZero() || expiry.Before(l.nextSweep) {
			l.nextSweep = expiry
		}
	}
	hit.count++
	l.hits[ip] = hit
	return true
}

func (s *Server) withAuthRateLimit(next http.HandlerFunc) http.HandlerFunc {
	if s.authLimiter == nil {
		s.authLimiter = newAuthRateLimiter()
	}
	return func(w http.ResponseWriter, r *http.Request) {
		ip := clientIP(r)
		if !s.authLimiter.allow(ip) {
			writeProblem(w, http.StatusTooManyRequests, "Too Many Requests", "auth rate limit exceeded")
			return
		}
		next(w, r)
	}
}