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
}

// rateLimiterMaxVisitors caps the per-limiter visitor map (token buckets keyed by client IP).
// Prevents unbounded memory growth under many distinct clients (or DoS) while prune
// still removes stale (>ttl) entries on every Allow. 1024 is generous for the intended
// LAN/minimal-public use; scans are O(cap) worst case.
const (
	rateLimiterMaxVisitors = 1024
	rateLimiterPruneTTL    = 10 * time.Minute
)

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:      rateLimiterPruneTTL,
	}
}

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
		if len(r.visitors) > rateLimiterMaxVisitors {
			// evict arbitrary entry (prune already removed stales); keeps map bounded
			for k := range r.visitors {
				delete(r.visitors, k)
				break
			}
		}
	}

	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 AccessLogMiddleware(logger *log.Logger, trustProxyHeaders bool, trustedProxies []netip.Prefix) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		if logger == nil {
			return next
		}
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if !shouldLogAccessRequest(r.Method, r.URL.Path) {
				next.ServeHTTP(w, r)
				return
			}

			start := time.Now()
			recorder := &accessLogResponseWriter{ResponseWriter: w}
			next.ServeHTTP(recorder, r)

			status := recorder.status
			if status == 0 {
				status = http.StatusOK
			}
			clientIP := "-"
			if ip, ok := extractClientIP(r, trustProxyHeaders, trustedProxies); ok {
				clientIP = ip.String()
			}
			logger.Printf(
				"method=%s path=%s status=%d bytes=%d duration_ms=%d ip=%s ua=%q",
				r.Method,
				r.URL.RequestURI(),
				status,
				recorder.bytes,
				time.Since(start).Milliseconds(),
				clientIP,
				r.UserAgent(),
			)
		})
	}
}

func shouldLogAccessRequest(method string, path string) bool {
	if method == http.MethodPost && path == "/api/scratch" {
		return true
	}
	if method == http.MethodGet && strings.HasPrefix(path, "/api/raw/") && len(path) > len("/api/raw/") {
		return true
	}
	return false
}

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
}

type accessLogResponseWriter struct {
	http.ResponseWriter
	status int
	bytes  int
}

func (w *accessLogResponseWriter) WriteHeader(statusCode int) {
	w.status = statusCode
	w.ResponseWriter.WriteHeader(statusCode)
}

func (w *accessLogResponseWriter) Write(p []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	n, err := w.ResponseWriter.Write(p)
	w.bytes += n
	return n, err
}

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

// SecurityHeadersMiddleware sets safe default security headers on all responses.
// HSTS is conditional on the hstsEnabled toggle (default true in config for public proxy+TLS use;
// false allows pure local HTTP as noted in README LAN Deployment Notes).
// Other headers always-on per minimal security defaults.
// CSP uses 'unsafe-inline' for style-src (and data: for media) due to current Vite/Svelte
// build output; this is an acknowledged defense-in-depth tradeoff for the minimal trusted UI
// (script-src remains strict 'self'; no user-controlled content in styles).
func SecurityHeadersMiddleware(hstsEnabled bool, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		h := w.Header()
		h.Set("X-Content-Type-Options", "nosniff")
		h.Set("Referrer-Policy", "strict-origin-when-cross-origin")
		h.Set("X-Frame-Options", "DENY")
		if hstsEnabled {
			h.Set("Strict-Transport-Security", "max-age=31536000; includeSubDomains")
		}
		h.Set("Content-Security-Policy", "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; media-src 'self' data:; font-src 'self'; connect-src 'self'; object-src 'none'; base-uri 'self'; frame-ancestors 'none'; form-action 'self';")
		h.Set("Permissions-Policy", "camera=(), microphone=(), geolocation=()")
		next.ServeHTTP(w, r)
	})
}