CosmicClash/server/api/rate_limit.go

package api

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net"
	"net/http"
	"net/netip"
	"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
}

// ClientIPResolver accepts X-Forwarded-For only from explicitly trusted
// immediate peers. It walks the chain from the application backwards so an
// untrusted client cannot select its own rate-limit identity by prepending a
// forged address.
type ClientIPResolver struct {
	trustedProxies []netip.Prefix
}

func NewClientIPResolver(cidrs string) (*ClientIPResolver, error) {
	resolver := &ClientIPResolver{}
	for _, raw := range strings.Split(cidrs, ",") {
		raw = strings.TrimSpace(raw)
		if raw == "" {
			continue
		}
		prefix, err := netip.ParsePrefix(raw)
		if err != nil {
			return nil, fmt.Errorf("invalid trusted proxy CIDR %q", raw)
		}
		resolver.trustedProxies = append(resolver.trustedProxies, prefix.Masked())
	}
	return resolver, nil
}

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, resolver *ClientIPResolver) string {
	keys := requestRateKeys(r, resolver)
	if len(keys) == 0 {
		return ""
	}
	return keys[0]
}

func requestRateKeys(r *http.Request, resolver *ClientIPResolver) []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 := requestClientIP(r, resolver)
	if host == "" {
		return keys
	}
	return append(keys, "ip:"+host)
}

func requestClientIP(r *http.Request, resolver *ClientIPResolver) string {
	host := strings.TrimSpace(r.RemoteAddr)
	if parsedHost, _, err := net.SplitHostPort(host); err == nil {
		host = parsedHost
	}
	remote, err := netip.ParseAddr(strings.Trim(host, "[]"))
	if err != nil {
		return host
	}
	remote = remote.Unmap()
	if resolver == nil || !resolver.trusts(remote) {
		return remote.String()
	}
	forwarded := strings.Join(r.Header.Values("X-Forwarded-For"), ",")
	if forwarded == "" || len(forwarded) > 2048 {
		return remote.String()
	}
	parts := strings.Split(forwarded, ",")
	if len(parts) > 16 {
		return remote.String()
	}
	chain := make([]netip.Addr, 0, len(parts))
	for _, part := range parts {
		address, parseErr := netip.ParseAddr(strings.TrimSpace(part))
		if parseErr != nil {
			return remote.String()
		}
		chain = append(chain, address.Unmap())
	}
	for index := len(chain) - 1; index >= 0; index-- {
		if !resolver.trusts(chain[index]) {
			return chain[index].String()
		}
	}
	if len(chain) > 0 {
		return chain[0].String()
	}
	return remote.String()
}

func (r *ClientIPResolver) trusts(address netip.Addr) bool {
	if r == nil || !address.IsValid() {
		return false
	}
	for _, prefix := range r.trustedProxies {
		if prefix.Contains(address) {
			return true
		}
	}
	return false
}