build-base

internal/middleware/security.go

@@ -0,0 +1,165 @@
+// Package middleware provides reusable HTTP middleware for security headers
+// and in-memory rate limiting (token bucket per IP).
+package middleware
+
+import (
+	"net"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+)
+
+// ---- Security Headers ----
+
+// SecureHeaders wraps an http.Handler and adds security-relevant response headers.
+// The CSP is intentionally strict: scripts only from same-origin CDN; no inline scripts.
+func SecureHeaders(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		h := w.Header()
+		h.Set("Strict-Transport-Security", "max-age=63072000; includeSubDomains; preload")
+		h.Set("X-Content-Type-Options", "nosniff")
+		h.Set("X-Frame-Options", "SAMEORIGIN")
+		h.Set("Referrer-Policy", "strict-origin-when-cross-origin")
+		h.Set("Permissions-Policy", "camera=(), microphone=(), geolocation=(), payment=()")
+		h.Set("Content-Security-Policy",
+			"default-src 'self'; "+
+				"script-src 'self' https://cdn.tailwindcss.com; "+
+				"style-src 'self' 'unsafe-inline' https://cdn.tailwindcss.com; "+
+				"img-src 'self' data:; "+
+				"font-src 'self'; "+
+				"frame-src 'self'; "+
+				"object-src 'none'; "+
+				"base-uri 'self'; "+
+				"form-action 'self'")
+		next.ServeHTTP(w, r)
+	})
+}
+
+// ---- Token-bucket rate limiter ----
+
+// bucket is one client's token state.
+type bucket struct {
+	tokens   float64
+	lastSeen time.Time
+}
+
+// RateLimiter is an in-memory token-bucket limiter keyed by IP.
+// Safe for concurrent use.
+type RateLimiter struct {
+	mu       sync.Mutex
+	buckets  map[string]*bucket
+	rate     float64 // tokens refilled per second
+	capacity float64 // max token count
+	status   int     // HTTP status on exceeded (default 429)
+}
+
+// NewRateLimiter creates a limiter.
+//   - ratePerMin: tokens refilled per minute (e.g. 20 = 20 req/min steady-state)
+//   - burst: max burst size (e.g. 5 = 5 simultaneous requests)
+func NewRateLimiter(ratePerMin int, burst int) *RateLimiter {
+	rl := &RateLimiter{
+		buckets:  make(map[string]*bucket),
+		rate:     float64(ratePerMin) / 60.0,
+		capacity: float64(burst),
+		status:   http.StatusTooManyRequests,
+	}
+	// Periodic cleanup goroutine (runs for process lifetime).
+	go rl.cleanup()
+	return rl
+}
+
+// Allow returns true if the request is within limit, consuming one token.
+func (rl *RateLimiter) Allow(ip string) bool {
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+
+	now := time.Now()
+	b, ok := rl.buckets[ip]
+	if !ok {
+		b = &bucket{tokens: rl.capacity, lastSeen: now}
+		rl.buckets[ip] = b
+	}
+
+	// Refill tokens based on elapsed time.
+	elapsed := now.Sub(b.lastSeen).Seconds()
+	b.tokens += elapsed * rl.rate
+	if b.tokens > rl.capacity {
+		b.tokens = rl.capacity
+	}
+	b.lastSeen = now
+
+	if b.tokens < 1 {
+		return false
+	}
+	b.tokens--
+	return true
+}
+
+// Middleware returns an http.Handler that rate-limits by remote IP.
+func (rl *RateLimiter) Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ip := remoteIP(r)
+		if !rl.Allow(ip) {
+			w.Header().Set("Retry-After", "60")
+			http.Error(w, http.StatusText(rl.status), rl.status)
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}
+
+// cleanup removes stale buckets every 5 minutes (no activity for >10 min).
+func (rl *RateLimiter) cleanup() {
+	ticker := time.NewTicker(5 * time.Minute)
+	defer ticker.Stop()
+	for range ticker.C {
+		cutoff := time.Now().Add(-10 * time.Minute)
+		rl.mu.Lock()
+		for ip, b := range rl.buckets {
+			if b.lastSeen.Before(cutoff) {
+				delete(rl.buckets, ip)
+			}
+		}
+		rl.mu.Unlock()
+	}
+}
+
+// remoteIP extracts the real client IP, honoring X-Real-IP and X-Forwarded-For.
+// Falls back to RemoteAddr.
+func remoteIP(r *http.Request) string {
+	if xri := r.Header.Get("X-Real-IP"); xri != "" {
+		if ip := net.ParseIP(xri); ip != nil {
+			return ip.String()
+		}
+	}
+	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
+		// First address in the chain is the original client.
+		for _, part := range splitComma(xff) {
+			trimmed := strings.TrimSpace(part)
+			if ip := net.ParseIP(trimmed); ip != nil {
+				return ip.String()
+			}
+		}
+	}
+	host, _, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		return r.RemoteAddr
+	}
+	return host
+}
+
+// splitComma splits on comma without allocating a regex.
+func splitComma(s string) []string {
+	var out []string
+	start := 0
+	for i := 0; i < len(s); i++ {
+		if s[i] == ',' {
+			out = append(out, s[start:i])
+			start = i + 1
+		}
+	}
+	out = append(out, s[start:])
+	return out
+}
+