x @ v0.3.1

// Package ratelimit provides an HTTP middleware that rate-limits requests

// on a per-client basis using the token bucket algorithm.

//

// Ecample:

//

//	limiter := ratelimit.MiddlewareWithConfig(ratelimit.Config{

//	    RPS:   10,

//	    Burst: 20,

//	    TTL:   time.Minute,

//	})

//	http.Handle("/api", limiter(yourHandler))

//

package ratelimit

import (

	"net"

	"net/http"

	"strings"

	"sync"

	"time"

	"olexsmir.xyz/x/ratelimit/internal/time/rate"

)

type (

	rateLimiter struct {

		mu sync.RWMutex

		visitors map[visitorIP]*visitor

		// limit is the maximum number of requests per second

		limit rate.Limit

		// ttl is the time after which a visitor is forgotten

		ttl time.Duration

		// burst is the maximum number of requests allowed in a short burst

		burst int

	}

	visitorIP string

	visitor   struct {

		limiter  *rate.Limiter

		lastSeen time.Time

	}

)

func newLimiter(rps, burst int, ttl time.Duration) *rateLimiter {

	return &rateLimiter{

		visitors: make(map[visitorIP]*visitor),

		limit:    rate.Limit(rps),

		burst:    burst,

		ttl:      ttl,

	}

}

// getVisitor retrieves the rate limiter for the given IP, or creates one.

func (r *rateLimiter) getVisitor(ip visitorIP) *rate.Limiter {

	r.mu.RLock()

	v, exists := r.visitors[ip]

	r.mu.RUnlock()

	if !exists {

		limit := rate.NewLimiter(r.limit, r.burst)

		r.mu.Lock()

		r.visitors[ip] = &visitor{

			limiter:  limit,

			lastSeen: time.Now(),

		}

		r.mu.Unlock()

		return limit

	}

	r.mu.Lock()

	v.lastSeen = time.Now()

	r.mu.Unlock()

	return v.limiter

}

// cleanupVisitors removes visitors that haven't been seen for longer than TTL.

func (r *rateLimiter) cleanupVisitors() {

	r.mu.Lock()

	defer r.mu.Unlock()

	for ip, v := range r.visitors {

		if time.Since(v.lastSeen) > r.ttl {

			delete(r.visitors, ip)

		}

	}

}

// cleanupVisitorsLoop runs cleanupVisitors every minute.

func (r *rateLimiter) cleanupVisitorsLoop() {

	for {

		time.Sleep(time.Minute)

		r.cleanupVisitors()

	}

}

// Config configures the rate limiter middleware.

type Config struct {

	// RPS is the maximum number of requests per second per visitor.

	RPS int

	// TTL is the duration after which an idle visitor's limiter is cleaned up.

	TTL time.Duration

	// Burst is the maximum number of requests that can be made in a short burst.

	Burst int

}

// Middleware returns a rate-limiting HTTP middleware.

//

// It limits requests per client IP using the token bucket algorithm.

// When the limit is exceeded it responds with 429 Too Many Requests.

func Middleware(c Config) func(http.Handler) http.Handler {

	lmt := newLimiter(c.RPS, c.Burst, c.TTL)

	go lmt.cleanupVisitorsLoop()

	return func(next http.Handler) http.Handler {

		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

			ip := getIP(r)

			limiter := lmt.getVisitor(ip)

			if limiter == nil {

				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)

				return

			}

			if !limiter.Allow() {

				http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)

				return

			}

			next.ServeHTTP(w, r)

		})

	}

}

func getIP(r *http.Request) visitorIP {

	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {

		if i := strings.IndexByte(xff, ','); i != -1 {

			xff = xff[:i]

		}

		if ip := strings.TrimSpace(xff); ip != "" {

			return visitorIP(ip)

		}

	}

	if xri := r.Header.Get("X-Real-IP"); xri != "" {

		return visitorIP(xri)

	}

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

		return visitorIP(host)

	}

	return visitorIP(r.RemoteAddr)

}