scratchbox/internal/http/middleware.go

package httpapi

import (
	"encoding/json"
	"log"
	"net"
	"net/http"
	"net/netip"
	"strings"
	"sync"
	"time"

	"golang.org/x/time/rate"
)

type visitor struct {
	limiter  *rate.Limiter
	lastSeen time.Time
}

type RateLimiter struct {
	mu       sync.Mutex
	visitors map[string]*visitor
	rate     rate.Limit
	burst    int
	ttl      time.Duration
}

func NewRateLimiter(requestsPerMinute int, burst int) *RateLimiter {
	return &RateLimiter{
		visitors: map[string]*visitor{},
		rate:     rate.Limit(float64(requestsPerMinute) / 60.0),
		burst:    burst,
		ttl:      10 * time.Minute,
	}
}

func (r *RateLimiter) Allow(ip string) bool {
	now := time.Now()

	r.mu.Lock()
	defer r.mu.Unlock()

	for key, v := range r.visitors {
		if now.Sub(v.lastSeen) > r.ttl {
			delete(r.visitors, key)
		}
	}

	v, ok := r.visitors[ip]
	if !ok {
		v = &visitor{
			limiter: rate.NewLimiter(r.rate, r.burst),
		}
		r.visitors[ip] = v
	}

	v.lastSeen = now
	return v.limiter.Allow()
}

func AllowlistMiddleware(allowed []netip.Prefix, trustProxyHeaders bool, trustedProxies []netip.Prefix, logger *log.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		if len(allowed) == 0 {
			return next
		}

		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			clientIP, ok := extractClientIP(r, trustProxyHeaders, trustedProxies)
			if !ok {
				http.Error(w, "unable to determine client ip", http.StatusForbidden)
				return
			}

			for _, prefix := range allowed {
				if prefix.Contains(clientIP) {
					next.ServeHTTP(w, r)
					return
				}
			}

			logger.Printf("allowlist denied request method=%s path=%s ip=%s", r.Method, r.URL.Path, clientIP.String())
			http.Error(w, "ip not allowed", http.StatusForbidden)
		})
	}
}

func RateLimitMiddleware(limiter *RateLimiter, trustProxyHeaders bool, trustedProxies []netip.Prefix) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			clientIP, ok := extractClientIP(r, trustProxyHeaders, trustedProxies)
			if !ok {
				http.Error(w, "unable to determine client ip", http.StatusTooManyRequests)
				return
			}

			if limiter.Allow(clientIP.String()) {
				next.ServeHTTP(w, r)
				return
			}

			w.Header().Set("Retry-After", "60")
			writeJSON(w, http.StatusTooManyRequests, map[string]string{
				"error": "too many requests from this IP, retry later",
			})
		})
	}
}

func Chain(handler http.Handler, middlewares ...func(http.Handler) http.Handler) http.Handler {
	wrapped := handler
	for i := len(middlewares) - 1; i >= 0; i-- {
		wrapped = middlewares[i](wrapped)
	}
	return wrapped
}

func extractClientIP(r *http.Request, trustProxyHeaders bool, trustedProxies []netip.Prefix) (netip.Addr, bool) {
	remoteIP, ok := remoteIPFromRequest(r)
	if !ok {
		return netip.Addr{}, false
	}

	if trustProxyHeaders && isTrustedProxy(remoteIP, trustedProxies) {
		if xff := strings.TrimSpace(r.Header.Get("X-Forwarded-For")); xff != "" {
			parts := strings.Split(xff, ",")
			if len(parts) > 0 {
				if ip, err := netip.ParseAddr(strings.TrimSpace(parts[0])); err == nil {
					return ip, true
				}
			}
		}
		if xrip := strings.TrimSpace(r.Header.Get("X-Real-IP")); xrip != "" {
			if ip, err := netip.ParseAddr(xrip); err == nil {
				return ip, true
			}
		}
	}
	return remoteIP, true
}

func remoteIPFromRequest(r *http.Request) (netip.Addr, bool) {
	host, _, err := net.SplitHostPort(strings.TrimSpace(r.RemoteAddr))
	if err != nil {
		if ip, parseErr := netip.ParseAddr(strings.TrimSpace(r.RemoteAddr)); parseErr == nil {
			return ip, true
		}
		return netip.Addr{}, false
	}
	ip, err := netip.ParseAddr(host)
	if err != nil {
		return netip.Addr{}, false
	}
	return ip, true
}

func isTrustedProxy(remoteIP netip.Addr, trusted []netip.Prefix) bool {
	for _, prefix := range trusted {
		if prefix.Contains(remoteIP) {
			return true
		}
	}
	return false
}

func writeJSON(w http.ResponseWriter, status int, payload any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(payload)
}