steamcache2/steamcache/ratelimit.go

// steamcache/ratelimit.go
// Per-client and global concurrency rate limiting, trusted proxy / client IP
// extraction logic (for safe X-Forwarded-For handling under security hardening),
// and background cleanup of idle client limiters. The clientRateLimiter wrapper owns
// the map + cleanup ticker/stop chan (SteamCache delegates; exact shutdown/wg patterns
// preserved).
package steamcache

import (
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/semaphore"
)

type clientLimiter struct {
	semaphore *semaphore.Weighted
	lastSeen  time.Time
}

// isTrustedProxy reports whether ipStr matches any CIDR or IP in trustedProxies list.
// Used for safe client IP extraction from X-Forwarded-For (rightmost untrusted proxy wins).
func isTrustedProxy(ipStr string, trustedProxies []string) bool {
	ip := net.ParseIP(strings.TrimSpace(ipStr))
	if ip == nil {
		return false
	}
	for _, c := range trustedProxies {
		c = strings.TrimSpace(c)
		if c == "" {
			continue
		}
		if !strings.Contains(c, "/") {
			if p := net.ParseIP(c); p != nil && p.Equal(ip) {
				return true
			}
			continue
		}
		if _, n, err := net.ParseCIDR(c); err == nil && n.Contains(ip) {
			return true
		}
	}
	return false
}

// getClientIP extracts the client IP address from the request.
// If trustedProxies is empty (the safe default), always use RemoteAddr only (prevents spoofing).
// When list non-empty, use rightmost-untrusted from XFF+Remote chain (proper proxy extraction, not naive first XFF).
// X-Real-IP is ignored for simplicity/safety (XFF is the standard multi-hop header).
// Security: prevents clients spoofing XFF to bypass per-client rate limits.
func getClientIP(r *http.Request, trustedProxies []string) string {
	// Normalize remote
	remoteIP := r.RemoteAddr
	if host, _, err := net.SplitHostPort(remoteIP); err == nil {
		remoteIP = host
	}

	if len(trustedProxies) == 0 {
		// Conservative safe default: never trust forwarded headers (spoof prevention)
		return remoteIP
	}

	// Build trust chain: XFF parts (left=original client) + direct remote (right=closest)
	chain := []string{}
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		for _, p := range strings.Split(xff, ",") {
			if t := strings.TrimSpace(p); t != "" {
				chain = append(chain, t)
			}
		}
	}
	chain = append(chain, remoteIP)

	// Walk from right (closest to server) to left; return first (rightmost) non-trusted = real client
	for i := len(chain) - 1; i >= 0; i-- {
		cand := chain[i]
		if !isTrustedProxy(cand, trustedProxies) {
			return cand
		}
	}
	return remoteIP
}

// clientRateLimiter owns the per-client limiters map, its mutex, the
// background cleanup stop channel, and the max-per-client config.
// Encapsulates lifecycle of the rate limiter map + its cleanup goroutine
// coordination (Phase 2). Unexported; same-package only.
type clientRateLimiter struct {
	mu           sync.RWMutex
	limiters     map[string]*clientLimiter
	cleanupStop  chan struct{}
	maxPerClient int64
}

// newClientRateLimiter constructs with its own stop chan (used by Run/Shutdown
// via thin delegates on SteamCache to preserve exact goroutine + wg patterns).
func newClientRateLimiter(maxPerClient int64) *clientRateLimiter {
	return &clientRateLimiter{
		limiters:     make(map[string]*clientLimiter),
		cleanupStop:  make(chan struct{}),
		maxPerClient: maxPerClient,
	}
}

func (crl *clientRateLimiter) getOrCreate(clientIP string) *clientLimiter {
	crl.mu.Lock()
	defer crl.mu.Unlock()

	limiter, exists := crl.limiters[clientIP]
	if !exists || time.Since(limiter.lastSeen) > 5*time.Minute {
		// Create new limiter or refresh existing one
		limiter = &clientLimiter{
			semaphore: semaphore.NewWeighted(crl.maxPerClient),
			lastSeen:  time.Now(),
		}
		crl.limiters[clientIP] = limiter
	} else {
		limiter.lastSeen = time.Now()
	}

	return limiter
}

// cleanupOld removes old client limiters to prevent memory leaks.
// Respects cleanupStop (closed by Shutdown) to allow graceful shutdown
// without hanging the wg.Wait in Run (historical shutdown hygiene).
func (crl *clientRateLimiter) cleanupOld() {
	ticker := time.NewTicker(10 * time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-crl.cleanupStop:
			return
		case <-ticker.C:
			crl.mu.Lock()
			now := time.Now()
			for ip, limiter := range crl.limiters {
				if now.Sub(limiter.lastSeen) > 30*time.Minute {
					delete(crl.limiters, ip)
				}
			}
			crl.mu.Unlock()
		}
	}
}

// getOrCreateClientLimiter delegates to the owned clientRateLimiter (preserves
// call sites in handler.go and shutdown coordination exactly).
func (sc *SteamCache) getOrCreateClientLimiter(clientIP string) *clientLimiter {
	return sc.clientRateLimiter.getOrCreate(clientIP)
}

// cleanupOldClientLimiters delegates to the internal rate limiter's cleanup
// loop. The goroutine launch + wg + stop-close patterns in Run/Shutdown are
// unchanged (critical for avoiding past goroutine leak / hang bugs).
func (sc *SteamCache) cleanupOldClientLimiters() {
	if sc.clientRateLimiter != nil {
		sc.clientRateLimiter.cleanupOld()
	}
}

// stop signals the background cleanup goroutine (started in Run) to exit by
// closing its stop channel. The exact select+default+close idiom is kept
// inside the wrapper (preserves all historical shutdown hygiene and wg.Wait
// behavior with zero change). Called from SteamCache.Shutdown.
func (crl *clientRateLimiter) stop() {
	if crl == nil || crl.cleanupStop == nil {
		return
	}
	select {
	case <-crl.cleanupStop:
	default:
		close(crl.cleanupStop)
	}
}