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 {
	return l.AllowKeys([]string{key}, now)
}

// AllowKeys atomically charges every non-empty key for a request. This lets
// the HTTP boundary enforce both the authenticated credential and source IP
// limits without charging one dimension when the other dimension rejects.
func (l *RateLimiter) AllowKeys(keys []string, now time.Time) bool {
	if l == nil || 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)
		}
	}
	unique := make([]string, 0, len(keys))
	seen := make(map[string]struct{}, len(keys))
	for _, key := range keys {
		if key == "" {
			continue
		}
		if _, exists := seen[key]; exists {
			continue
		}
		seen[key] = struct{}{}
		unique = append(unique, key)
	}
	if len(unique) == 0 {
		return false
	}
	newKeys := 0
	for _, key := range unique {
		entry, exists := l.entries[key]
		if !exists {
			newKeys++
			continue
		}
		if now.Before(entry.started.Add(l.window)) && entry.count >= l.limit {
			return false
		}
	}
	if len(l.entries)+newKeys > l.maxKeys {
		return false
	}
	for _, key := range unique {
		entry, exists := l.entries[key]
		if !exists || !now.Before(entry.started.Add(l.window)) {
			l.entries[key] = rateWindow{started: now, count: 1}
			continue
		}
		entry.count++
		l.entries[key] = entry
	}
	return true
}

func requestRateKey(r *http.Request) string {
	keys := requestRateKeys(r)
	if len(keys) == 0 {
		return ""
	}
	return keys[0]
}

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