samsa/internal/middleware/ratelimit.go

// samsa — a privacy-respecting metasearch engine
// Copyright (C) 2026-present metamorphosis-dev
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

package middleware

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

	"log/slog"
)

// RateLimitConfig controls per-IP rate limiting.
type RateLimitConfig struct {
	Requests        int
	Window          time.Duration
	CleanupInterval time.Duration
	// TrustedProxies is a list of CIDR ranges that are allowed to set
	// X-Forwarded-For / X-Real-IP. If empty, only r.RemoteAddr is used.
	TrustedProxies []string
}

func RateLimit(cfg RateLimitConfig, logger *slog.Logger) func(http.Handler) http.Handler {
	requests := cfg.Requests
	if requests <= 0 {
		requests = 30
	}

	window := cfg.Window
	if window <= 0 {
		window = time.Minute
	}

	cleanup := cfg.CleanupInterval
	if cleanup <= 0 {
		cleanup = 5 * time.Minute
	}

	if logger == nil {
		logger = slog.Default()
	}

	// Parse trusted proxy CIDRs.
	var trustedNets []*net.IPNet
	for _, cidr := range cfg.TrustedProxies {
		_, network, err := net.ParseCIDR(cidr)
		if err != nil {
			logger.Warn("invalid trusted proxy CIDR, skipping", "cidr", cidr, "error", err)
			continue
		}
		trustedNets = append(trustedNets, network)
	}

	limiter := &ipLimiter{
		requests: requests,
		window:   window,
		clients:  make(map[string]*bucket),
		logger:   logger,
		trusted:  trustedNets,
	}

	go limiter.cleanup(cleanup)

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ip := limiter.extractIP(r)

			if !limiter.allow(ip) {
				retryAfter := int(limiter.window.Seconds())
				w.Header().Set("Retry-After", strconv.Itoa(retryAfter))
				w.Header().Set("Content-Type", "text/plain; charset=utf-8")
				w.WriteHeader(http.StatusTooManyRequests)
				_, _ = w.Write([]byte("429 Too Many Requests\n"))
				logger.Debug("rate limited", "ip", ip)
				return
			}

			next.ServeHTTP(w, r)
		})
	}
}

type bucket struct {
	count    int
	expireAt time.Time
}

type ipLimiter struct {
	requests int
	window   time.Duration
	clients  map[string]*bucket
	mu       sync.Mutex
	logger   *slog.Logger
	trusted  []*net.IPNet
}

func (l *ipLimiter) allow(ip string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := time.Now()
	b, ok := l.clients[ip]

	if !ok || now.After(b.expireAt) {
		l.clients[ip] = &bucket{
			count:    1,
			expireAt: now.Add(l.window),
		}
		return true
	}

	b.count++
	return b.count <= l.requests
}

func (l *ipLimiter) cleanup(interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for range ticker.C {
		l.mu.Lock()
		now := time.Now()
		for ip, b := range l.clients {
			if now.After(b.expireAt) {
				delete(l.clients, ip)
			}
		}
		l.mu.Unlock()
	}
}

// extractIP extracts the client IP from the request.
// If trusted proxy CIDRs are configured, X-Forwarded-For is only used when
// the direct connection comes from a trusted proxy. Otherwise, only RemoteAddr is used.
func (l *ipLimiter) extractIP(r *http.Request) string {
	return extractIP(r, l.trusted...)
}

func extractIP(r *http.Request, trusted ...*net.IPNet) string {
	remoteIP, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		remoteIP = r.RemoteAddr
	}

	// Check if the direct connection is from a trusted proxy.
	isTrusted := false
	if len(trusted) > 0 {
		ip := net.ParseIP(remoteIP)
		if ip != nil {
			for _, network := range trusted {
				if network.Contains(ip) {
					isTrusted = true
					break
				}
			}
		}
	}

	if isTrusted {
		if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
			parts := strings.SplitN(xff, ",", 2)
			candidate := strings.TrimSpace(parts[0])
			if net.ParseIP(candidate) != nil {
				return candidate
			}
		}
		if rip := r.Header.Get("X-Real-IP"); rip != "" {
			candidate := strings.TrimSpace(rip)
			if net.ParseIP(candidate) != nil {
				return candidate
			}
		}
	}

	return remoteIP
}