glchat/internal/httpx/ratelimit.go

package httpx

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

// RateLimiter — ограничитель частоты запросов по ключу (IP или пользователь).
// Реализация — token bucket без внешних зависимостей: 1 vCPU инстансу
// достаточно простого счётчика в памяти (AGENT.md 8.6).
type RateLimiter struct {
	mu      sync.Mutex
	buckets map[string]*rateBucket
	// rate — токенов в секунду, burst — максимальный запас.
	rate  float64
	burst float64
	// now подменяется в тестах, чтобы не ждать реального времени.
	now func() time.Time
}

type rateBucket struct {
	tokens float64
	seen   time.Time
}

// NewRateLimiter создаёт лимитер: perMinute запросов в минуту с запасом burst.
func NewRateLimiter(perMinute, burst int) *RateLimiter {
	if perMinute <= 0 {
		perMinute = 60
	}
	if burst <= 0 {
		burst = perMinute
	}
	return NewRateLimiterWindow(perMinute, time.Minute, burst)
}

// NewRateLimiterWindow создаёт лимитер на произвольное окно: нужно для
// суточных квот (например, 10 приглашений в день, AGENT.md 8.6).
func NewRateLimiterWindow(limit int, window time.Duration, burst int) *RateLimiter {
	if limit <= 0 {
		limit = 1
	}
	if window <= 0 {
		window = time.Minute
	}
	if burst <= 0 {
		burst = limit
	}
	return &RateLimiter{
		buckets: map[string]*rateBucket{},
		rate:    float64(limit) / window.Seconds(),
		burst:   float64(burst),
		now:     time.Now,
	}
}

// SetClock подменяет источник времени (тесты).
func (l *RateLimiter) SetClock(now func() time.Time) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.now = now
}

// Allow проверяет запрос и возвращает паузу до следующей попытки.
func (l *RateLimiter) Allow(key string) (bool, time.Duration) {
	if key == "" {
		key = "unknown"
	}
	now := l.now()

	l.mu.Lock()
	defer l.mu.Unlock()

	bucket, ok := l.buckets[key]
	if !ok {
		// Попутная уборка: словарь не должен расти бесконечно (AGENT.md 9.1).
		if len(l.buckets) > 4096 {
			l.sweepLocked(now)
		}
		bucket = &rateBucket{tokens: l.burst, seen: now}
		l.buckets[key] = bucket
	}
	elapsed := now.Sub(bucket.seen).Seconds()
	if elapsed > 0 {
		bucket.tokens = minFloat(l.burst, bucket.tokens+elapsed*l.rate)
	}
	bucket.seen = now

	if bucket.tokens < 1 {
		missing := 1 - bucket.tokens
		retryAfter := time.Duration(missing / l.rate * float64(time.Second))
		if retryAfter < time.Second {
			retryAfter = time.Second
		}
		return false, retryAfter
	}
	bucket.tokens--
	return true, 0
}

// sweepLocked удаляет записи, неактивные больше десяти минут.
func (l *RateLimiter) sweepLocked(now time.Time) {
	cutoff := now.Add(-10 * time.Minute)
	for key, bucket := range l.buckets {
		if bucket.seen.Before(cutoff) {
			delete(l.buckets, key)
		}
	}
}

// Middleware ограничивает частоту запросов по ключу keyFn.
func (l *RateLimiter) Middleware(keyFn func(*http.Request) string) Middleware {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			allowed, retryAfter := l.Allow(keyFn(r))
			if !allowed {
				WriteRateLimited(w, retryAfter)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// MiddlewareExcept ограничивает частоту запросов везде, кроме перечисленных
// путей: у живого дашборда инстанса свой лимит, иначе опрос раз в секунду
// съедал бы половину общего бюджета API (AGENT.md 8.6).
func (l *RateLimiter) MiddlewareExcept(keyFn func(*http.Request) string, except ...string) Middleware {
	skip := make(map[string]struct{}, len(except))
	for _, path := range except {
		skip[path] = struct{}{}
	}
	limited := l.Middleware(keyFn)
	return func(next http.Handler) http.Handler {
		handler := limited(next)
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if _, ok := skip[r.URL.Path]; ok {
				next.ServeHTTP(w, r)
				return
			}
			handler.ServeHTTP(w, r)
		})
	}
}

// WriteRateLimited отдаёт 429 в конверте API с подсказкой по паузе (§8.5, §8.6).
func WriteRateLimited(w http.ResponseWriter, retryAfter time.Duration) {
	milliseconds := retryAfter.Milliseconds()
	if milliseconds <= 0 {
		milliseconds = 1000
	}
	w.Header().Set("Retry-After", strconv.FormatInt((milliseconds+999)/1000, 10))
	WriteJSON(w, http.StatusTooManyRequests, map[string]any{
		"error": map[string]any{
			"code":           "rate_limited",
			"message":        "too many requests",
			"retry_after_ms": milliseconds,
		},
	})
}

func minFloat(a, b float64) float64 {
	if a < b {
		return a
	}
	return b
}