CosmicClash/server/api/rate_limit.go

package api

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// RateLimiter is an optional fixed-window limiter for the control-plane edge.
// It is intentionally process-local: a deployment must use a shared edge
// limiter for global quotas, while this boundary still protects each replica.
type RateLimiter struct {
	mu      sync.Mutex
	limit   int
	window  time.Duration
	maxKeys int
	entries map[string]rateWindow
}

type rateWindow struct {
	started time.Time
	count   int
}

func NewRateLimiter(limit int, window time.Duration, maxKeys int) (*RateLimiter, error) {
	if limit < 1 || window <= 0 || maxKeys < 1 {
		return nil, fmt.Errorf("invalid rate limiter configuration")
	}
	return &RateLimiter{limit: limit, window: window, maxKeys: maxKeys, entries: make(map[string]rateWindow)}, nil
}

func (l *RateLimiter) Allow(key string, now time.Time) bool {
	if l == nil || key == "" || now.IsZero() {
		return false
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	for storedKey, entry := range l.entries {
		if !now.Before(entry.started.Add(l.window)) {
			delete(l.entries, storedKey)
		}
	}
	entry, exists := l.entries[key]
	if !exists {
		if len(l.entries) >= l.maxKeys {
			return false
		}
		l.entries[key] = rateWindow{started: now, count: 1}
		return true
	}
	if !now.Before(entry.started.Add(l.window)) {
		l.entries[key] = rateWindow{started: now, count: 1}
		return true
	}
	if entry.count >= l.limit {
		return false
	}
	entry.count++
	l.entries[key] = entry
	return true
}

func requestRateKey(r *http.Request) string {
	if authorization := strings.TrimSpace(r.Header.Get("Authorization")); authorization != "" {
		digest := sha256.Sum256([]byte(authorization))
		return "auth:" + hex.EncodeToString(digest[:])
	}
	host := r.RemoteAddr
	if parsedHost, _, err := net.SplitHostPort(host); err == nil {
		host = parsedHost
	}
	if host == "" {
		return ""
	}
	return "ip:" + host
}