kafka/internal/middleware/ratelimit.go

// kafka — 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 using a sliding window counter.
type RateLimitConfig struct {
	// Requests is the max number of requests allowed per window.
	Requests int
	// Window is the time window duration (e.g. "1m").
	Window time.Duration
	// CleanupInterval is how often stale entries are purged (default: 5m).
	CleanupInterval time.Duration
}

// RateLimit returns a middleware that limits requests per IP address.
// Uses an in-memory sliding window counter. When the limit is exceeded,
// responds with HTTP 429 and a Retry-After header.
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()
	}

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

	go limiter.cleanup(cleanup)

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ip := 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
}

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()
	}
}

func extractIP(r *http.Request) string {
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		parts := strings.SplitN(xff, ",", 2)
		return strings.TrimSpace(parts[0])
	}
	if rip := r.Header.Get("X-Real-IP"); rip != "" {
		return strings.TrimSpace(rip)
	}

	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}